no-space.hs

-- start snippet imports
{-# LANGUAGE LambdaCase #-}

module Main where

import Data.Char (isDigit)
import Data.List (sortOn)
import qualified Data.Map.Strict as Map
import Text.ParserCombinators.ReadP ((<++))
import qualified Text.ParserCombinators.ReadP as P
-- end snippet imports

-- start snippet fs-types
data File = File {fName :: String, fSize :: Int}

instance Show File where
  show (File name size) = name <> "(file, size=" <> show size <> ")"

data Dir = Dir
  { dName :: String
  , dSize :: Int
  , dFiles :: Map.Map String File
  , dDirs :: Map.Map String Dir
  }

instance Show Dir where
  showsPrec d (Dir name _ files dirs) =
    showString (concat $ replicate d "  ")
      . showString "- "
      . showString name
      . showString " (dir)\n"
      . showDirs (sortOn dName $ Map.elems dirs)
      . showFiles (sortOn fName $ Map.elems files)
    where
      showDirs :: [Dir] -> ShowS
      showDirs = foldr (.) id . fmap (showsPrec $ d + 1)

      showFiles :: [File] -> ShowS
      showFiles = foldr ((.) . showFile) id

      showFile :: File -> ShowS
      showFile (File name' size) =
        showString (concat $ replicate (d + 1) "  ")
          . showString "- "
          . showString name'
          . showString " (file, size="
          . shows size
          . showString ")\n"

emptyDir :: String -> Dir
emptyDir name = Dir name 0 Map.empty Map.empty
-- end snippet fs-types

-- start snippet command-types
data CdArg = CdDir String | CdUp | CdRoot deriving (Show)
data Command = Cd CdArg | Ls deriving (Show)
data Output = OutputFile File | OutputDir Dir deriving (Show)
data Line = LCommand Command | LOutput Output deriving (Show)
-- end snippet command-types

-- start snippet parsers
commandParser :: P.ReadP Command
commandParser =
  P.char '$' *> P.skipSpaces *> P.choice [Cd <$> cdParser, Ls <$ lsParser]
  where
    lsParser = P.string "ls"
    cdParser =
      P.string "cd"
        *> P.skipSpaces
        *> ((CdUp <$ P.string "..")
             <++ (CdRoot <$ P.string "/")
             <++ (CdDir <$> P.munch1 (/= ' ')))

outputParser :: P.ReadP Output
outputParser = P.choice [OutputFile <$> fileParser, OutputDir <$> dirParser]
  where
    fileParser =
      flip File
        <$> (read <$> P.munch1 isDigit)
        <*> (P.skipSpaces *> P.munch1 (/= ' '))
    dirParser = emptyDir <$> (P.string "dir " *> P.munch1 (/= ' '))

lineParser :: P.ReadP Line
lineParser = P.choice [LOutput <$> outputParser, LCommand <$> commandParser]

parseLine :: String -> Line
parseLine s = case P.readP_to_S lineParser s of
  [(l, "")] -> l
  _ -> error $ "Failed to parse line: " <> s
-- end snippet parsers

-- start snippet zippers
data FsZipper = FsZipper {zPath :: [Dir], zCurrent :: Dir} deriving (Show)

moveUp :: FsZipper -> FsZipper
moveUp = \case
  FsZipper [] _ -> error "Can't move up from root"
  FsZipper (d : ds) cur ->
    FsZipper ds $ d {dDirs = Map.insert (dName cur) cur $ dDirs d}

moveDown :: String -> FsZipper -> FsZipper
moveDown name (FsZipper ds d) = FsZipper (d : ds) $ findDir d
  where
    findDir Dir {dDirs = ds'} = case Map.lookup name ds' of
      Nothing -> error $ "Can't find directory " <> name
      Just d' -> d'

moveToRoot :: FsZipper -> FsZipper
moveToRoot zipper = case zipper of
  FsZipper [] _ -> zipper
  _ -> moveToRoot $ moveUp zipper

addFile :: File -> FsZipper -> FsZipper
addFile f (FsZipper ds d) =
  FsZipper ds d {dFiles = Map.insert (fName f) f $ dFiles d}

addDir :: Dir -> FsZipper -> FsZipper
addDir d (FsZipper ds d') =
  FsZipper ds d' {dDirs = Map.insert (dName d) d $ dDirs d'}

toZipper :: Dir -> FsZipper
toZipper = FsZipper []

fromZipper :: FsZipper -> Dir
fromZipper = zCurrent . moveToRoot
-- end snippet zippers

-- start snippet interpreter
interpretLine :: FsZipper -> Line -> FsZipper
interpretLine zipper = \case
  LCommand (Cd CdUp) -> moveUp zipper
  LCommand (Cd CdRoot) -> moveToRoot zipper
  LCommand (Cd (CdDir name)) -> moveDown name zipper
  LCommand Ls -> zipper
  LOutput (OutputFile f) -> addFile f zipper
  LOutput (OutputDir d) -> addDir d zipper

interpret :: [Line] -> Dir
interpret = fromZipper . foldl interpretLine (toZipper $ emptyDir "/")
-- end snippet interpreter

-- start snippet solution
calcAndSetDirSize :: Dir -> Dir
calcAndSetDirSize d@Dir {dFiles = fs, dDirs = ds} =
  let ds' = fmap calcAndSetDirSize ds
   in d {dSize = sum $ fmap fSize fs <> fmap dSize ds', dDirs = ds'}

findDirsSmallerThan :: Int -> Dir -> [Dir]
findDirsSmallerThan size d@(Dir {dSize = dSize', dDirs = ds}) =
  [d | dSize' <= size] <> concatMap (findDirsSmallerThan size) ds
-- end snippet solution

-- start snippet part1
part1 :: Dir -> Int
part1 = sum . map dSize . findDirsSmallerThan 100000
-- end snippet part1

-- start snippet part2
findDirsLargerThan :: Int -> Dir -> [Dir]
findDirsLargerThan size d@(Dir {dSize = dSize', dDirs = ds}) =
  [d | dSize' >= size] <> concatMap (findDirsLargerThan size) ds

part2 :: Dir -> Int
part2 fs =
  let freeSpace = totalSpace - dSize fs
      spaceRequired = updateSpace - freeSpace
      dirs = findDirsLargerThan spaceRequired fs
   in minimum $ map dSize dirs
  where
    totalSpace = 70000000
    updateSpace = 30000000
-- end snippet part2

-- start snippet main
main :: IO ()
main = do
  input <- lines <$> getContents
  let fs = calcAndSetDirSize $ interpret $ map parseLine input
  print fs
  print $ part1 fs
  print $ part2 fs
-- end snippet main