hcdu/app/Main.hs

{-# LANGUAGE BlockArguments #-}
module Main where

import           System.Posix.Directory ( getWorkingDirectory, changeWorkingDirectory )
import           System.Directory       ( pathIsSymbolicLink
                                        , doesFileExist
                                        , getFileSize
                                        , doesDirectoryExist )
import           Data.Functor           ( (<&>) )
import           Control.Monad          ( filterM )
import           Control.Monad.IO.Class
import           System.Directory.Extra ( listContents )
import           System.IO.Error
import           Brick
import           Brick.Widgets.Edit
import           Graphics.Vty
import           Data.List.Extra
import qualified Data.Ord
import           Data.Maybe


listNonSymFiles :: FilePath -> IO [FilePath]
listNonSymFiles dir = filterM ((not <$>) . pathIsSymbolicLink) =<< filterM doesFileExist =<< listContents dir

listNonSymDirectories :: FilePath -> IO [FilePath]
listNonSymDirectories dir = filterM ((not <$>) . pathIsSymbolicLink) =<< filterM doesDirectoryExist =<< listContents dir

listSym :: FilePath -> IO [FilePath]
listSym dir = filterM pathIsSymbolicLink =<< listContents dir

handlePermissionError :: IOError -> IO Integer
handlePermissionError e = if isPermissionError e then pure 0 else ioError e

getSizeSubpaths :: FilePath -> IO [(FilePath, Integer)]
getSizeSubpaths x = do
  sub   <- listNonSymDirectories x
  subby <- mapM
             (\y -> (y,) . sum . map snd <$>
               getSizeSubpaths y)
           sub
  syms  <- listSym x
  local <- listNonSymFiles x >>=
    (\y -> mapM
      (flip catchIOError handlePermissionError . getFileSize) y <&>
       zip y)
  pure $ sortOn (Data.Ord.Down . snd) $ local ++ subby ++ zip syms [0..0]


data ScrollDirection where
  SUp   :: ScrollDirection
  SDown :: ScrollDirection

data AppS = AppS
  {
    appCWD :: FilePath
  , appCursor :: Int
  , appFocus :: Maybe FilePath
  , appSubFiles :: [(FilePath, Integer)]
  }

app :: App AppS e ()
app =
    App { appDraw = pure . browse
        , appHandleEvent = eventHandler
        , appStartEvent = pure ()
        , appAttrMap = attributeMap
        , appChooseCursor = showFirstCursor
        }

attributeMap :: AppS -> AttrMap
attributeMap = const $ attrMap defAttr
                  [ (attrName "selected", withStyle defAttr reverseVideo)
                  , (editAttr, fg brightBlack)
                  ]

select :: Widget () -> Widget ()
select = withAttr (attrName "selected")

unitSize :: Integer -> String
unitSize s
  | s3 >= 1 / 2 =
      show (f s3) ++ " GB"
  | s2 >= 1 / 2 =
      show (f s2) ++ " MB"
  | s1 >= 1 / 2 =
      show (f s1) ++ " KB"
  | otherwise   =
      show  s     ++ "  B"
  where k2 = 1024 :: Double
        s3 = fromInteger s / (k2^(3 :: Int))
        s2 = fromInteger s / (k2^(2 :: Int))
        s1 = fromInteger s / (k2^(1 :: Int))
        f :: Double -> Double
        f q = fromInteger (truncate $ q * 100) / 100

browse :: AppS -> Widget ()
browse s =
  str "Path" <+> padLeft Max (str "Size") <=>
  viewport () Vertical (foldr (widgetCons s) emptyWidget (appSubFiles s))

widgetCons :: AppS -> (FilePath, Integer) -> Widget () -> Widget ()
widgetCons s w@(f,_) ws =
  (<=> ws) if Just f == appFocus s then
           select . visible $ pathWidget w else
           pathWidget w

pathWidget :: (FilePath, Integer) -> Widget ()
pathWidget (f, s) = str (show f) <+> padLeft Max (str (unitSize s))

sizeDir :: AppS -> IO AppS
sizeDir s = do
  subFiles <- getSizeSubpaths $ appCWD s
  pure $
    s {
        appCursor = 0
      , appFocus = map fst subFiles !? 0
      , appSubFiles = subFiles
      }

changeDir :: AppS -> IO AppS
changeDir so
  | isJust $ appFocus so = do
    allowed <- doesDirectoryExist path
    if allowed then do
      changeWorkingDirectory path
      let s =
            so {
                appCWD       = path
              , appSubFiles  = []
              } in
                sizeDir s else
      pure so
  | otherwise             = pure so
  where path = fromJust $ appFocus so

overDir :: IO AppS
overDir = do
  changeWorkingDirectory ".."
  a <- getWorkingDirectory
  changeDir $ initialState a

scroll :: ScrollDirection -> AppS -> AppS
scroll d s = s {
                 appCursor = newCursor
               , appFocus  = maybeNewPath
               }
  where cursor =
          appCursor s +
          case d of
            SUp   -> (-1)
            SDown -> 1
        newCursor = max 0 (min (subtract 1 . length $ appSubFiles s) cursor)
        maybeNewPath = fst <$> appSubFiles s !? newCursor

eventHandler :: BrickEvent () e -> EventM () AppS ()
eventHandler (VtyEvent (EvKey k _)) = do
  s <- get
  case k of
      (KChar 'q') -> halt
      KEsc        -> halt

      (KChar 'j') -> put $ scroll SDown s
      KDown       -> put $ scroll SDown s
      (KChar 'k') -> put $ scroll SUp s
      KUp         -> put $ scroll SUp s

      (KChar 'h') -> put =<< liftIO overDir
      KLeft       -> put =<< liftIO overDir

      (KChar ' ') -> put =<< liftIO (changeDir s)
      KEnter      -> put =<< liftIO (changeDir s)
      (KChar 'l') -> put =<< liftIO (changeDir s)
      KRight      -> put =<< liftIO (changeDir s)

      _           -> continueWithoutRedraw

eventHandler _ = undefined

initialState :: FilePath -> AppS
initialState f =
  AppS {
          appCWD = f
        , appCursor = 0
        , appFocus = pure f
        , appSubFiles = []
        }

main :: IO ()
main = do
  a <- getWorkingDirectory
  b <- changeDir $ initialState a
  _ <- defaultMain app b
  return ()