-- Name: Adam Matousek
-- UID: 445763
-- Used packages: base, filepath, directory
-- If version of 'directory' is < 1.2.7, then you can use 'unix' package and
-- uncomment code on lines 132-136 (sorry for not using preprocessor).

module Main ( main ) where

import Numeric ( showGFloat )
import System.Console.GetOpt
import System.Environment ( getArgs )
import System.FilePath
import System.Posix.Files ( getFileStatus, fileSize )
import System.Directory
import Control.Monad ( forM, mapM, liftM2 )
import Control.Exception ( catch, IOException )
import Data.Maybe ( isJust, fromJust )
import Data.Monoid

-- Used size units, because why use a boring bool
data Format = Human | SI
    deriving Show

-- Configuration record
data DuConf = DuConf
    { optAll :: Any
    , optFormat :: Last Format
    , optDepth :: Last Int
    , optTotal :: Any
    , optHelp :: Any
    }
instance Monoid DuConf where
    mappend a b = DuConf
        { optAll = optAll a `mappend` optAll b
        , optFormat = optFormat a `mappend` optFormat b
        , optDepth = optDepth a `mappend` optDepth b
        , optTotal = optTotal a `mappend` optTotal b
        , optHelp = optHelp a `mappend` optHelp b
        }
    mempty = DuConf
        { optAll = Any False
        , optFormat = Last Nothing
        , optDepth = Last Nothing
        , optTotal = Any False
        , optHelp = Any False
        }

-- Main function
main :: IO ()
main = getArgs >>= return . 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 = Last $ Just 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 = Last $ Just Human})
        "display sizes in human readable format"
    , Option "" ["si"]
        (NoArg mempty {optFormat = Last $ Just SI})
        "as --human-readable, but use SI prefices"
    , Option "d" ["max-depth"]
        (ReqArg (\d -> mempty {optDepth = (Last . Just . 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"
    ]

-- 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
          realpaths = if null paths then ["."] else paths
parseOptionsAndRun (_,_, errors) =
    putStr $ concat errors ++ usageInfo usageHeader options

-- 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 =
    if getAny $ optTotal conf
        then putStrLn $ show (foldl (+) 0 sizes) ++ "\ttotal"
        else return ()

-- 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 =
    (doesDirectoryExist path) >>= (\isdir ->
        (pathIsSymbolicLink path) >>= (\issymlink ->
            if isdir && not issymlink
                then fmap (map (path </>)) (listDirectory path)
                    >>= mapM (duOne conf (depth + 1)) >>= return . foldr (+) 0
                else getFileSize path))
    >>= printOne conf depth path

-- Function with the same name and type has been a part of directory package
-- since 1.2.7, but my Haskell stack is apparently still too old.
-- Not wanting to use another additional package, I opted to use the POSIX API.
-- Using hFileSize is NOT the same thing, because of permissions.
-- This should be, however, the only platform-dependent function.
--
-- Apparently, isSymbolicLink has been depreceated in favor of
-- pathIsSymbolicLink, but again, not on my system yet.
{-
getFileSize :: FilePath -> IO Integer
getFileSize path = fmap (toInteger . fileSize) (getFileStatus path)

pathIsSymbolicLink :: FilePath -> IO Bool
pathIsSymbolicLink = isSymbolicLink
-}


-- 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 =
    let maxdepth = (getLast . optDepth) conf in
    doesFileExist path >>= (\isfile ->
        if isJust maxdepth && depth > fromJust maxdepth
            -- Nothing is printed, if maximal depth has been reached
            then return size
            -- Regular file is printed if -a is used or if it is a command line
            -- argument on its own (then depth == 0)
            else if isfile && (not . getAny . optAll) conf && depth > 0
                then return size
                else putStrLn (formatSize conf size ++ "\t" ++ path)
                    >> return size)

-- transforms the size in bytes into format requested in config
formatSize :: DuConf -> Integer -> String
formatSize conf =
    case getLast (optFormat conf) of
    Nothing -> show
    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 (Just 1) (sizef / bf ^ (magn * d)) unit
    where shared_prefices = ["M", "G", "T", "P", "E", "Z", "Y"]
          bin_prefices = "" : "K" : shared_prefices
          si_prefices = "" : "k" : shared_prefices
          prefices = if b == 10 then si_prefices else bin_prefices
          sizef = fromInteger sizei
          bf = fromIntegral b
          df = fromIntegral d
          magn = min (length prefices - 1) $ floor $ logBase bf sizef / df
          unit = ' ' : (prefices !! magn) ++ maybeI ++ "B"
          maybeI = if magn == 0 || b == 10 then "" else "i"