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

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 = Stks (List (List Crate))

Show Stacks where
  show (Stks stacks) =
    let longestStack = foldl max 0 . map length $ stacks
        stringStacks = map (extendFront longestStack "   ") . map (map show) $ stacks
        lastLine = unwords . map (\x => " " ++ x ++ " ") . map show $ rangeFromTo 1 (length stacks)
        lineStacks = filter (/= "") . map unwords . transposeList longestStack $ stringStacks
    in joinBy "\n" $ (lineStacks ++ [lastLine])
    where extendFront : Nat -> a -> List a -> List a
          extendFront k x xs =
            if length xs < k
            then replicate (minus k (length xs)) x ++ xs
            else xs
          transposeList : Nat -> List (List a) -> List (List a)
          transposeList k xs =
            if any isNil xs
            then replicate k []
            else let firsts = map (take 1) xs
                     rests = map (drop 1) xs
                 in foldl (++) [] firsts :: transposeList k (assert_smaller xs rests)

-- Chunk up a list
chunk : Nat -> List a -> List (List a)
chunk k [] = []
chunk k xs = let (head, tail) = splitAt k xs in
             head :: chunk k tail

-- Ensure a list of maybes is at least a specific length
ensureLen : Nat -> List (Maybe a) -> List (Maybe a)
ensureLen k [] = replicate k Nothing
ensureLen 0 xs = xs
ensureLen (S k) (x :: xs) = x :: ensureLen k xs

-- Parse a list of stacks, this should be passed a list of lines including the line of numbers
parseStacks : Nat -> List String -> Stacks
parseStacks k [] = Stks $ replicate k []
parseStacks k (x :: []) = parseStacks k []
parseStacks k (x :: xs) = let (Stks bottom) = parseStacks k xs
                              row = ensureLen k . map parseCrate . map pack . chunk 4 . unpack $ x
                          in Stks $ zipWith parseStacksHelper row bottom
  where 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 -> String
tops (Stks xss) = pack . map (\(C c) => c) . catMaybes . map head' $ xss

simple : String
simple = """
      [D]
  [N] [C]
  [Z] [M] [P]
  1   2   3

  move 1 from 2 to 1
  move 3 from 1 to 3
  move 2 from 2 to 1
  move 1 from 1 to 2
  """
data Command = Cmd Nat Nat Nat

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

-- Parse a command
parseCommand : String -> Maybe Command
parseCommand str =
  let nums : List Nat
      nums = catMaybes . map parsePositive . words $ str
  in case nums of
        [x, y, z] => Just (Cmd x (minus y 1) (minus z 1))
        _ => Nothing

-- Apply a command to the stacks
applyCommand : Command -> Bool -> Stacks -> Maybe Stacks
applyCommand (Cmd count from to) multiple stacks =
  do (crates, removed) <- extractCrates count from stacks
     addCrates (if multiple then crates else reverse crates) to removed
  where extractCrates : Nat -> Nat -> Stacks -> Maybe ((List Crate), Stacks)
        extractCrates count from (Stks []) = Nothing
        extractCrates count 0 (Stks (x :: xs)) =
          if count > length x
          then Nothing
          else let (taken, rest) = splitAt count x in
               Just (taken, Stks (rest :: xs))
        extractCrates count (S k) (Stks (x :: xs)) =
          do (c, (Stks s)) <- extractCrates count k (Stks xs)
             Just (c, (Stks (x :: s)))
        addCrates : List Crate -> Nat -> Stacks -> Maybe Stacks
        addCrates xs to (Stks []) = Nothing
        addCrates xs 0 (Stks (x :: ys)) = Just (Stks ((xs ++ x) :: ys))
        addCrates xs (S k) (Stks (x :: ys)) =
          do (Stks added) <- addCrates xs k (Stks ys)
             Just (Stks (x :: added))

data CmdResult = CR Bool (Maybe Command) Stacks

Show CmdResult where
  show (CR bool command stacks) =
    unlines ["Command applied: " ++ (show bool), "Command: " ++ (show command), show stacks]

-- Apply a list of commands to a stacks
applyCommands : List Command -> Bool -> Stacks -> List CmdResult
applyCommands [] multiple x = [CR False Nothing x]
applyCommands (y :: xs) multiple x =
  case applyCommand y multiple x of
    Just stacks => CR True (Just y) stacks :: applyCommands xs multiple stacks
    Nothing => [CR False (Just y) x]

mapM : Monad m => (a -> m b) -> List a -> m (List b)
mapM f [] = pure []
mapM f (x :: xs) =
  do x <- f x
     rest <- mapM f xs
     pure $ x :: rest

partial unwrap : Maybe a -> a
unwrap (Just x) = x

partial main : IO ()
main =
 do file <- readFile "input"
    case file of
      Right content =>
        let stacks = parseStacks 9 . takeWhile (/= "") . lines $ content
            commands = catMaybes . map parseCommand . drop 1 . dropWhile (/= "") . lines $ content
            part1 = unwrap . map (\(CR _ _ stacks) => tops stacks) . last' . applyCommands commands False $ stacks
            part2 = unwrap . map (\(CR _ _ stacks) => tops stacks) . last' . applyCommands commands True $ stacks
        in do putStrLn "Input: "
              printLn stacks
              putStrLn " "
              putStr "Part 1: "
              putStrLn part1
              putStr "Part 2: "
              putStrLn part2
      Left err => printLn err