#!/usr/bin/env runhaskell
import Control.Monad (forM, forM_)
import Control.Exception (bracket, handle)

import Data.Maybe (fromMaybe)
import Data.List (sortBy)
import Data.Function (on)

import Text.Printf (printf)

import Numeric (showFFloat)

import System.Directory (doesDirectoryExist, getDirectoryContents)
import System.FilePath ((</>), takeFileName)
import System.IO (IOMode(..), hClose, hFileSize, openFile)
import System.Environment (getArgs)

data DirTree = Directory String [DirTree] | File String
    deriving Show

name (Directory n _) = n
name (File n) = n

data SizeTree = Node String Integer [SizeTree]
    deriving Show

size (Node _ s _) = s

getDirectoryTree :: FilePath -> IO DirTree
getDirectoryTree topdir = do
  names <- getDirectoryContents topdir
  let properNames = filter (`notElem` [".", ".."]) names
  children <- forM properNames $ \name -> do
    let path = topdir </> name
    isDirectory <- doesDirectoryExist path
    if isDirectory
      then getDirectoryTree path
      else return (File path)
  let children' = sortBy (compare `on` name) children
  return (Directory topdir children')

getFileSize :: FilePath -> IO (Maybe Integer)
getFileSize path = handle (const $ return Nothing) $
  bracket (openFile path ReadMode) hClose $ \h -> do
    size <- hFileSize h
    return (Just size)

getSizes :: DirTree -> IO SizeTree
getSizes (Directory name children) = do
    childSizes <- mapM getSizes children
    let totalSize = sum . map size $ childSizes
    return $ Node name totalSize childSizes
getSizes (File name) = do
    size <- getFileSize name
    return $ Node name (fromMaybe 0 size) []

printSizes :: SizeTree -> IO ()
printSizes tree = go' 0 tree
  where go' indent (Node name size children) = do
        putStrLn $ printf "%s[%s] %s" (replicate indent ' ') (prettySize size) (takeFileName name)
        mapM_ (go' (indent + 2)) children

prettySize :: (Integral a) => a -> String
prettySize s = printf "%.2f %s" (size'::Double) unit
    where (size', unit) = last $ (s', head units) : zip (takeWhile ((>= 1) . abs) (iterate (/1024) (s'/1024))) (tail units)
          s' = fromIntegral s
          units = ["B", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"]

main :: IO ()
main = do
    args <- getArgs
    forM_ args $ \arg -> (getDirectoryTree arg) >>= getSizes >>= printSizes