import System.File.ReadWrite
import Data.Vect
import Data.String
import Data.Fin

data Crate = C Char

Show Crate where
  show (C c) = pack ['[', c, ']']

parseCrate : String -> Maybe Crate
parseCrate str = let str = unpack str in
                 case str of
                   ('[' :: x :: ']' :: _) => Just (C x)
                   _ => Nothing

data Stacks : Nat -> Type where
  Stks : { n : Nat } -> Vect n (List Crate) -> Stacks n

Show (Stacks n) where
  show (Stks {n} stacks) =
    let longestStack = foldl max 0 . map length $ stacks
        stringStacks = map (extendFront longestStack "   ") . map (map show) $ stacks
        lastLine = joinBy " " . map (\x => " " ++ (show x) ++ " ") $ rangeFromTo 1 n
        stackLines = map (joinBy " " . toList) . transpose $ stringStacks
    in joinBy "\n" . toList $ (snoc stackLines lastLine)
    where extendFront : (m : Nat) -> a -> List a -> Vect m a
          extendFront m x xs =
            let extended = if length xs < m
                           then replicate (minus m (length xs)) x ++ xs
                           else take m xs
            in case toVect m extended of
                Just vect => vect
                Nothing => replicate m x
          transpose : { a : _ } -> Vect b (Vect a e) -> Vect a (Vect b e)
          transpose [] = replicate _ []
          transpose (x :: xs) = let xsTrans = transpose xs in
                                zipWith (::) x xsTrans

parseStacks : {n : Nat} -> List String -> Stacks n
parseStacks [] = Stks $ replicate _ []
-- Special logic for skipping the last line, since that's just the stack numbers
parseStacks (x :: []) = Stks $ replicate _ []
parseStacks {n} (x :: xs) =
  let (Stks bottom) = parseStacks xs
      row : Vect n (Maybe Crate)
      row = map parseCrate . map pack . chunk 4 [] . unpack $ x
  in Stks $ zipWith parseStacksHelper row bottom
  where chunk : {n : Nat} -> (m: Nat) -> List a -> List a -> Vect n (List a)
        chunk {n = 0} m defaultValue xs = []
        chunk {n = (S k)} m defaultValue [] = replicate _ defaultValue
        chunk {n = (S k)} m defaultValue xs =
          let (head, tail) = splitAt m xs in
          head :: chunk {n = k} m defaultValue tail
        parseStacksHelper : Maybe Crate -> List Crate -> List Crate
        parseStacksHelper (Just crate) xs = crate :: xs
        parseStacksHelper Nothing xs = xs

-- Get the tops of each stack
tops : Stacks n -> String
tops (Stks xss) = pack . map (\(C c) => c) . catMaybes . toList . map head' $ xss

data Command : Nat -> Type where
  Cmd : {n : Nat} -> Nat -> Fin n -> Fin n -> Command n

Show (Command n) where
  show (Cmd x y z) = (show x) ++ " " ++ (show ((finToNat y) + 1)) ++ "->" ++ (show ((finToNat z) + 1))

parseCommand : {n : Nat} -> String -> Maybe (Command n)
parseCommand str =
  let nums = toVect 3 . catMaybes . map parsePositive . words $ str
  in do [x, y, z] <- nums
        from <- natToFin (minus y 1) n
        to <- natToFin (minus z 1) n
        Just (Cmd x from to)

-- Apply a command to a collection of stacks
applyCommand : {n : Nat} -> Bool -> Command n -> Stacks n -> Stacks n
applyCommand multiple (Cmd count from to) stacks =
  let (crates, removed) = extractCrates count from stacks
  in addCrates (if multiple then crates else reverse crates) to removed
  where extractCrates : Nat -> Fin n -> Stacks n -> ((List Crate), Stacks n)
        extractCrates count from (Stks stacks) =
          let stack = index from stacks
              (head, tail) = splitAt count stack
          in (head, Stks $ replaceAt from tail stacks)
        addCrates : List Crate -> Fin n -> Stacks n -> Stacks n
        addCrates crates to (Stks stacks) =
          let stack = index to stacks
          in Stks $ replaceAt to (crates ++ stack) stacks

-- Apply a list of commands to a stacks
applyCommands : {n : Nat} -> Bool -> List (Command n) -> Stacks n -> Stacks n
applyCommands multiple commands stacks = foldl (\stack, command => applyCommand multiple command stack) stacks commands

parseProblem : String -> Maybe (n ** (Stacks n, List (Command n)))
parseProblem content =
  case last' . takeWhile(/= "") . lines $ content of
    Just columnsLine =>
      let size = foldl max 0 . catMaybes . map parsePositive . words $ columnsLine
          stacks = parseStacks {n = size} . takeWhile (/= "") . lines $ content
          commands = catMaybes . map (parseCommand {n = size}) . drop 1 . dropWhile (/= "") . lines $ content
      in Just (size ** (stacks, commands))
    Nothing => Nothing


doPart : {n : Nat} -> Bool -> Stacks n -> List (Command n) -> String
doPart multiple stacks commands = tops $ applyCommands multiple commands stacks


main : IO ()
main =
  do file <- readFile "input"
     case file of
       Right content =>
         case parseProblem content of
           Just (size ** (stacks, commands)) =>
             let part1 = doPart False stacks commands
                 part2 = doPart True stacks commands
             in do putStrLn "Input:"
                   printLn stacks
                   putStr "\n\nPart 1: "
                   putStrLn part1
                   putStr "Part 2: "
                   putStrLn part2
           Nothing => putStrLn "Failed to parse problem"
       Left err => printLn err