{- Possible solution of third homework for IB016, semester spring 2019 -}

-- Used packages: base, filepath, directory, Glob

module Main ( main ) where

import Numeric ( showGFloat )
import System.Console.GetOpt
import System.Environment ( getArgs )
import System.FilePath
import System.FilePath.Glob
import System.Directory
import Control.Monad ( forM, mapM, when )
import Control.Exception ( catch, IOException )
import Data.Semigroup hiding ( Option )
import Data.Functor ( (<&>) )

data Format = Human | SI
    deriving Show

-- Configuration record
data DuConf = DuConf
    { optAll :: Any
    , optFormat :: Maybe (Last Format)
    , optDepth :: Maybe (Last Int)
    , optTotal :: Any
    , optHelp :: Any
    , optExclude :: [Pattern]
    }

instance Semigroup DuConf where
    a <> b = DuConf
        { optAll = optAll a <> optAll b
        , optFormat = optFormat a <> optFormat b
        , optDepth = optDepth a <> optDepth b
        , optTotal = optTotal a <> optTotal b
        , optHelp = optHelp a <> optHelp b
        , optExclude = optExclude a <> optExclude b
        }
    stimes = stimesMonoid

instance Monoid DuConf where
    mempty = DuConf
        { optAll = Any False
        , optFormat = Nothing
        , optDepth = Nothing
        , optTotal = Any False
        , optHelp = Any False
        , optExclude = []
        }

-- Main function
main :: IO ()
main = getArgs <&> getOpt RequireOrder options >>= buildConfigAndRun

-- First line of usage information
usageHeader :: String
usageHeader = "USAGE: du [OPTION]... [--] [PATH]..."

-- Available options (used with getOpt)
options :: [OptDescr DuConf]
options =
    [ Option "a" ["all"]
        (NoArg mempty {optAll = Any True})
        "display size of all files, not only directories"
    , Option "s" ["summarize"]
        (NoArg mempty {optDepth = Just $ Last 0})
        "display only a total for each argument"
    , Option "c" ["total"]
        (NoArg mempty {optTotal = Any True})
        "produce a grand total"
    , Option "h" ["human-readable"]
        (NoArg mempty {optFormat = Just $ Last Human})
        "display sizes in human readable format"
    , Option "" ["si"]
        (NoArg mempty {optFormat = Just $ Last SI})
        "as --human-readable, but use SI prefices"
    , Option "d" ["max-depth"]
        (ReqArg (\d -> mempty {optDepth = (Just . Last . read) d}) "N")
        "only display files and directories that are N or fewer levels deep."
    , Option "?" ["help"] (NoArg mempty {optHelp = Any True})
        "display this help and exit"
    , Option "" ["exclude"]
        (ReqArg (\pat -> mempty {optExclude = [compile pat]}) "PATTERN")
        "exclude files and directories that match PATTERN"
    ]

-- Takes result of getOpt and either prints usage info or performs actual 'du'
buildConfigAndRun :: ([DuConf], [String], [String]) -> IO ()
buildConfigAndRun (opts, paths, []) =
    if getAny $ optHelp conf
        then putStr $ usageInfo usageHeader options
        else runFor conf realpaths
    where conf      = mconcat opts
          paths'    = if null paths then ["."] else paths
          realpaths = filter (notExcluded conf) paths'
buildConfigAndRun (_, _, errors) = putStr $
    concat errors ++ usageInfo usageHeader options

notExcluded :: DuConf -> FilePath -> Bool
notExcluded conf p = not $ any (`match` p) (optExclude conf)

-- Runs the program for the specified top-level files/directories.
-- This is basically the "real" main after option parsing
runFor :: DuConf -> [FilePath] -> IO ()
runFor conf paths = forM paths (duOne conf 0) >>= printTotal conf

-- Prints a grand total, if config requires it
printTotal :: DuConf -> [Integer] -> IO ()
printTotal conf sizes =
    when (getAny $ optTotal conf)
         (putStrLn $ formatSize conf (sum sizes) ++ "\ttotal")

-- Returns size of one file or (recursively) directory, possibly printing it
-- Second argument is depth of current directory
-- If an exception is caught, it is printed and zero is returned
duOne :: DuConf -> Int -> FilePath -> IO Integer
duOne conf depth path = catch (duOneUnsafe conf depth path)
                              (\e -> print (e :: IOException) >> return 0)

-- As the previous function, only it doesn't catch exceptions (that's what the
-- first one is for)
duOneUnsafe :: DuConf -> Int -> FilePath -> IO Integer
duOneUnsafe conf depth path = do
    isdir <- doesDirectoryExist path
    issymlink <- pathIsSymbolicLink path
    sz <- if isdir && not issymlink
            then do
                children <- filter (notExcluded conf) <$> listDirectory path
                let childpaths = map (path </>) children
                sum <$> mapM (duOne conf (depth + 1)) childpaths
            else getFileSize path
    printOne conf depth path sz

-- Possibly prints size and name of one object; returns second argument in IO
printOne :: DuConf -> Int -> FilePath -> Integer -> IO Integer
printOne conf depth path size = size <$ do
    let maxdepth = (fmap getLast . optDepth) conf
    isdir <- doesDirectoryExist path
    -- Only print when depth is under maximal depth or no maximum is set, and
    -- only directories, unless -a is given or it is a command line argument
    -- (then its depth equals 0)
    when (maybe True (>= depth) maxdepth &&
            (isdir || (getAny . optAll) conf || depth == 0))
         (putStrLn (formatSize conf size ++ "\t" ++ path))

-- transforms the size in bytes into format requested in config
formatSize :: DuConf -> Integer -> String
formatSize conf = case getLast <$> optFormat conf of
    Nothing -> show . flip div 1024
    Just SI -> formatSizeLogBase 10 3
    Just Human -> formatSizeLogBase 2 10

-- Actual size formatter doing the hard work.
formatSizeLogBase :: Int -> Int -> Integer -> String
formatSizeLogBase b d sizei = showGFloat decimalPlaces dividedSize unit
    where prefices = if b == 10 then si_prefices else bin_prefices
          sizef = fromInteger sizei :: Double
          bf = fromIntegral b
          df = fromIntegral d
          magn = min (length prefices - 1) $ floor $ logBase bf sizef / df
          dividedSize = sizef / bf ^ (magn * d)
          decimalPlaces = Just $ if dividedSize >= 10 then 0 else 1
          shared_prefices = ["M", "G", "T", "P", "E", "Z", "Y"]
          bin_prefices = "" : "K" : shared_prefices
          si_prefices = "" : "k" : shared_prefices
          maybeI = if magn == 0 || b == 10 then "" else "i"
          unit = (prefices !! magn) ++ maybeI ++ "B"