hcdu/app/Main.hs

{-# LANGUAGE BlockArguments #-}
module Main where

import           Prelude         hiding ( lookup )
import           System.Posix.Directory ( getWorkingDirectory, changeWorkingDirectory )
import           System.Directory       ( pathIsSymbolicLink
                                        , doesFileExist
                                        , getFileSize
                                        , doesDirectoryExist )
import           Data.Functor           ( (<&>) )
import           Data.IORef
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              ( isPrefixOf )
import           Data.List.Extra        ( sortOn
                                        , (!?))
import qualified Data.Ord
import           Data.Maybe
import           Data.IOHashMap hiding ( foldr
                                       , (!?) )

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 ->
  IORef (IOHashMap FilePath [(FilePath, Integer)]) ->
  IO [(FilePath, Integer)]
getSizeSubpaths x i = do
  memo <- readIORef i
  val <- lookup x memo
  case val of
    Just b -> pure b
    _      -> do
        sub   <- listNonSymDirectories x
        subby <- mapM
                  (\y -> (y,) . sum . map snd <$>
                    getSizeSubpaths y i)
                sub
        syms  <- listSym x
        local <- listNonSymFiles x >>=
          (\y -> mapM
            (flip catchIOError handlePermissionError . getFileSize) y <&>
            zip y)
        let temp = sortOn (Data.Ord.Down . snd) $ local ++ subby ++ zip syms [0..0]
        insert x temp memo
        writeIORef i memo
        pure temp


data ScrollDirection where
  SUp   :: ScrollDirection
  SDown :: ScrollDirection

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

app :: App AppS e ()
app =
    App { appDraw = pure . browse
        , appHandleEvent = eventHandler
        , appStartEvent = put =<< liftIO . changeDir =<< get
        , 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) (appMemo 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 :: AppS -> IO AppS
overDir s = do
  changeWorkingDirectory ".."
  a <- getWorkingDirectory
  if a `isPrefixOf` appRoot s && a /= appRoot s then
    pure s else
    changeDir $ s {
                    appCWD = a
                  , appCursor = 0
                  , appFocus = pure a
                  , appSubFiles = []
                  }

scroll :: ScrollDirection -> AppS -> AppS
scroll d s = s {
                 appCursor = newCursor
               , appFocus  = maybeNewPath
               }
  where cursor =
          appCursor s +
          case d of
            SUp   -> (-1)
            SDown -> 1
        newCursor = clamp 0 (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 s)
      KLeft       -> put =<< liftIO (overDir s)

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

      _           -> continueWithoutRedraw
eventHandler (VtyEvent (EvResize _ _)) = pure ()
eventHandler _ = undefined

initialState ::
  FilePath ->
  IORef (IOHashMap FilePath [(FilePath, Integer)]) ->
  AppS
initialState f i =
  AppS {
          appCWD = f
        , appCursor = 0
        , appFocus = pure f
        , appSubFiles = []
        , appMemo = i
        , appRoot = f
        }

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