{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}

module Days.D05 where

import Common
import Parse (Parser, digit, number, parsecDay, someLines)

import Data.Vector (Vector, (!), (//))
import Data.Vector qualified as V
import Text.Megaparsec (MonadParsec (try), anySingle, between, endBy, sepBy, skipManyTill, (<|>))
import Text.Megaparsec.Char (char, letterChar, newline, string)

type State = Vector String
type Instruction = (Int, Int, Int)
type Intermediate = (State, [Instruction])

data CrateMoverModel = CM9000 | CM9001

stackOn :: String -> [String] -> [String]
stackOn row stacks = zipWith (\c l -> if c == ' ' then l else c : l) paddedRow stacks
  where
    paddedRow =
        if length row < length stacks
            then row ++ replicate (length stacks - length row) ' '
            else row

buildStacks :: [String] -> Vector String
buildStacks rows = V.fromList $ foldr stackOn emptyStacks rows
  where
    numStacks = maximum (map length rows)
    emptyStacks = replicate numStacks ""

parser :: Parser Intermediate
parser = do
    stackMatrix <- try crate `sepBy` char ' ' `endBy` newline
    _ <- skipManyTill anySingle newline *> newline
    instructions <- someLines instruction
    return (buildStacks stackMatrix, instructions)
  where
    crate =
        between (char '[') (char ']') letterChar
            <|> between (char ' ') (char ' ') (char ' ')

    instruction :: Parser Instruction
    instruction =
        (,,)
            <$ string "move "
            <*> number
            <* string " from "
            <*> (pred <$> digit)
            <* string " to "
            <*> (pred <$> digit)

doReverse :: CrateMoverModel -> [a] -> [a]
doReverse CM9000 = reverse
doReverse CM9001 = id

move :: CrateMoverModel -> State -> Instruction -> State
move model s (c, src, dst) =
    let
        (transfer, newSrc) = splitAt c (s ! src)
        newDst = doReverse model transfer ++ s ! dst
     in
        s // [(src, newSrc), (dst, newDst)]

part1 :: Intermediate -> Vector Char
part1 = uncurry ((fmap head .) . foldl (move CM9000))

part2 :: Intermediate -> Vector Char
part2 = uncurry ((fmap head .) . foldl (move CM9001))

day :: Day
day = parsecDay parser (definitive . part1, definitive . part2)

fast :: Day
fast = day