import Data.List
import Data.List1
import Data.Vect
import Data.String
import System.File.ReadWrite

%default total

data DirEntry = File String Nat | Directory String

Show DirEntry where
  show (File name len) = name ++ " (" ++ show len ++ ")"
  show (Directory name) = name

parseDirEntry : String -> Maybe DirEntry
parseDirEntry entry =
  let worded = words entry
  in do fst <- head' worded
        lst <- last' worded
        if fst == "dir"
          then Just (Directory lst)
          else do len <- parsePositive fst
                  Just (File lst len)

data Command : Type where
  Cd : String -> Command
  Ls : List DirEntry -> Command

Show Command where
  show (Cd str) = "cd " ++ str
  show (Ls dirs) = "ls: " ++ (unwords . map show $ dirs)

-- Parse one command from a list of lines
parseCmd : List String -> Maybe (Command, List String)
parseCmd [] = Nothing
parseCmd (cmd :: lines) =
  if isPrefixOf "$ cd" cmd
  then case last' . words $ cmd of
         Just dir => Just (Cd dir, lines)
         Nothing => Nothing
  else if isPrefixOf "$ ls" cmd
       then let ours = takeWhile (not . isPrefixOf "$") lines
                theirs = dropWhile (not . isPrefixOf "$") lines
            in do children <- traverse parseDirEntry ours
                  Just (Ls children, theirs)
       else Nothing

-- Parse all the commands in a list
parseCmds : List String -> Maybe (List Command)
parseCmds lines = map reverse $ helper [] lines
  where helper : List Command -> List String -> Maybe (List Command)
        helper [] [] = Nothing
        helper cmds [] = Just cmds
        helper cmds lines =
          do (command, rest) <- parseCmd lines
             helper (command :: cmds) (assert_smaller lines rest)

data Map : Nat -> Type -> Type -> Type where
  Nil : Map 0 k v
  (::) : (k, v) -> Map n k v -> Map (S n) k v
%name Map map, map2, map3

get : Eq k => k -> Map n k v -> Maybe v
get x [] = Nothing
get x ((key, value) :: map) =
  if key == x
    then Just value
    else get x map

values : Map n k v -> Vect n v
values [] = []
values ((x, y) :: map) = y :: values map

mapToList : Map n k v -> List (k, v)
mapToList [] = []
mapToList (x :: map) = x :: mapToList map

insert : k -> v -> Map n k v -> Map (S n) k v
insert key value map = (key, value) :: map

replace : Eq k => k -> v -> Map n k v -> Maybe (Map n k v)
replace key value [] = Nothing
replace key value ((x, y) :: map) =
  if key == x
    then Just $ (key, value) :: map
    else do rest <- replace key value map
            Just $ (x, y) :: rest

insertPair : k -> v -> (n: Nat ** Map n k v) -> (m: Nat ** Map m k v)
insertPair key value ((fst ** snd)) = (S fst ** insert key value snd)

record Node where
  constructor MkDir
  name : String
  files : (n: Nat ** Map n String (Nat, String))
  children : (m : Nat ** Map m String Node)

size : Node -> Nat
size self@(MkDir nodename (fileCount ** files) (childCount ** children)) =
  let fileTotal = sum . map fst . values $ files
      dirTotal = sum . map (\x => size (assert_smaller self x)) . values $ children
  in fileTotal + dirTotal

showDepth : Nat -> Node -> List String
showDepth depth self@(MkDir name (fileCount ** files) (childCount ** children)) =
  let fileLines = toList . map (\(size, name) => leader depth ++ name ++ " (" ++ show size ++ ")" ) . values $ files
      header = leader depth ++ name ++ ": (" ++ show (size self) ++ ")"
      childrenLines = foldMap (\x => showDepth (depth + 1) (assert_smaller self x)) . values $ children
  in header :: (fileLines ++ childrenLines)
  where leader : Nat -> String
        leader k = replicate k ' '

Show Node where
  show = unlines . showDepth 0

emptyNode : Node
emptyNode = MkDir "Root" (0 ** Nil) (0 ** Nil)

directories : Node -> List Node
directories self@(MkDir nodename files (childCount ** children)) =
  self :: foldMap (\x => directories (assert_smaller self x)) (values children)

insertFile : List String -> String -> Nat -> Node -> Maybe Node
insertFile [] name size (MkDir nodename files children) = Just $ MkDir nodename (insertPair name (size, name) files) children
insertFile (x :: xs) name size (MkDir nodename files (childCount ** children)) =
  case get x children of
    Nothing => Nothing
    (Just y) => do result <- insertFile xs name size y
                   replaced <- replace x result children
                   Just $ MkDir nodename files (childCount ** replaced)

insertDir : List String -> String -> Node -> Maybe Node
insertDir [] name (MkDir nodename files children) = Just $ MkDir nodename files (insertPair name (MkDir name (0 ** Nil) (0 ** Nil)) children)
insertDir (x :: xs) name (MkDir nodename files (childCount ** children)) =
  case get x children of
    Nothing => Nothing
    (Just y) => do result <- insertDir xs name y
                   replaced <- replace x result children
                   Just $ MkDir nodename files (childCount ** replaced)

insertEntry : List String -> Node -> DirEntry -> Maybe Node
insertEntry path node (File name size) = insertFile path name size node
insertEntry path node (Directory name) = insertDir path name node

data PathStack = Stk (List String)

current : PathStack -> List String
current (Stk strs) = reverse strs

pop : PathStack -> PathStack
pop (Stk []) = Stk []
pop (Stk (x :: xs)) = Stk xs

push : String -> PathStack -> PathStack
push "/" stack = stack
push ".." stack = pop stack
push path (Stk strs) = Stk (path :: strs)

applyCommands : List Command -> Node -> Maybe Node
applyCommands xs x = helper (Stk []) xs x
  where helper : PathStack -> List Command -> Node -> Maybe Node
        helper stack [] node = Just node
        helper stack ((Cd name) :: ys) node =
          helper (push name stack) ys node
        helper stack ((Ls entries) :: ys) node =
          do result <- foldlM (insertEntry (current stack)) node entries
             helper stack ys result

part1 : String ->  IO ()
part1 input =
  case parseCmds (lines input) of
    Nothing => putStrLn "Failed to parse commands"
    (Just commands) =>
      case applyCommands commands emptyNode of
        Nothing => putStrLn "Failed to apply commands"
        (Just node) =>
          let wanted = sum . filter (<= 100000) . map size . directories $  node
          in putStrLn $ "Part 1: " ++ show wanted

minBy : Ord b => (a -> b) -> a -> a -> a
minBy f x y =
  let i = f x
      j = f y
  in if i < j then x else y

part2 : String -> IO ()
part2 input =
  case parseCmds (lines input) of
    Nothing => putStrLn "Failed to parse commands"
    (Just commands) =>
      case applyCommands commands emptyNode of
        Nothing => putStrLn "Failed to apply commands"
        (Just node) =>
          let totalSpace = 70000000
              neededSpace = 30000000
              freeSpace = minus totalSpace (size node)
              deficit = minus neededSpace freeSpace
              canidates = filter ((>= deficit) . size) (directories node)
          in case Data.List1.fromList canidates of
               Nothing => putStrLn "Only one canidate?"
               (Just canidates) =>
                 let smallest = foldl1 (minBy size) canidates
                 in putStrLn $ "Part 2: " ++ (show (size smallest))


partial main : IO ()
main =
  do file <- readFile "input"
     case file of
       (Right contents) =>
         do part1 contents
            part2 contents
       (Left err) => printLn err