import Control.Monad.Trans.State as S
import Data.List
import Data.Map as M

type Cwd = [String]
data Entry = File Integer | Folder (M.Map String Entry) deriving Show
data S = S {pwd :: Cwd, root :: Entry}
type Shell = S.State S

addEntry :: [String] -> String -> Entry -> Entry -> Entry
addEntry [] name entry (Folder m) = Folder $ M.insert name entry m
addEntry (x:xs) name entry (Folder m) = Folder $ M.insert x new m where
    new = addEntry xs name entry (m M.! x)

handleInput :: [String] -> Shell ()
handleInput ["$", "cd", ".."] = do
    s <- S.get
    let new = tail (pwd s)
    put (s {pwd = new})

handleInput ["$", "cd", "/"] = do
    s <- S.get; put $ s {pwd = []}

handleInput ["$", "cd", x] = do
    s <- S.get
    let new = x : pwd s
    put (s {pwd = new})

handleInput ["$", "ls"] = pure ()
handleInput ["dir", name] = do
    s <- S.get
    let new = addEntry (reverse $ pwd s) name (Folder M.empty) (root s)
    put $ s {root = new}

handleInput [size, name] = do
    s <- S.get
    let new = addEntry (reverse $ pwd s) name (File (read size)) (root s)
    put $ s {root = new}

sumSize :: Entry -> Integer
sumSize (File s) = s
sumSize (Folder m) = sum $ fmap sumSize m

allDirs :: Entry -> [Entry]
allDirs (File _) = []
allDirs f@(Folder m) = f : concatMap allDirs (M.elems m)


main :: IO ()
main = do
    x <- lines <$> readFile "inputs/day07"
    let init = S [] (Folder M.empty)
    let actions = mapM_ (handleInput . words) x
    let tree = root $ execState actions init

    putStr "part 1: "
    let dirSizes = fmap sumSize (allDirs tree)
    print $ sum $ Prelude.filter (<= 100000) dirSizes

    putStr "part 2: "
    let unused = 70000000 - sumSize tree
    let required = 30000000 - unused
    print $ minimum $ Prelude.filter (>= required) dirSizes