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

data Operation = Noop | Addx Int

Show Operation where
  show Noop = "Noop"
  show (Addx i) = "Addx " ++ show i

cycles : Operation -> Nat
cycles Noop = 0
cycles (Addx i) = 1

parseOp : String -> Maybe Operation
parseOp str =
  case map trim . words $ str of
    ["noop"] => Just Noop
    ["addx", val] => do val <- parseInteger val
                        pure (Addx val)
    _ => Nothing

parseOps : String -> Maybe (List Operation)
parseOps str = traverse parseOp (lines str)

data Pixel = On | Off

Show Pixel where
  show On = "#"
  show Off = " "

record State where
  constructor MkState
  cycle : Nat
  xRegister : Int
  waitCycles : Nat
  waitingOp : Maybe Operation

Show State where
  show (MkState cycle xRegister waitCycles waitingOp) =
    "Cycle: " ++ show cycle ++ " X: " ++ show xRegister ++ " waitingOp: " ++ show waitingOp

startState : State
startState = MkState 0 1 0 Nothing

isWaiting : State -> Bool
isWaiting (MkState cycle xRegister waitCycles Nothing) = False
isWaiting (MkState cycle xRegister waitCycles (Just x)) = True

-- Tick the CPU, taking an operation off the stack if needed
tick : State -> List Operation -> (State, List Operation)
tick (MkState cycle xRegister 0 Nothing) [] =
  (MkState (S cycle) xRegister 0 Nothing, [])
tick (MkState cycle xRegister 0 (Just Noop)) [] =
  (MkState (S cycle) xRegister 0 Nothing, [])
tick (MkState cycle xRegister 0 (Just (Addx i))) [] =
  (MkState (S cycle) (xRegister + i) 0 Nothing, [])
tick (MkState cycle xRegister (S k) waitingOp) [] =
  (MkState (S cycle) xRegister k waitingOp, [])
tick (MkState cycle xRegister 0 Nothing) (x :: xs) =
  (MkState (S cycle) xRegister (cycles x) (Just x), xs)
tick (MkState cycle xRegister 0 (Just Noop)) (x :: xs) =
  (MkState (S cycle) xRegister (cycles x) (Just x), xs)
tick (MkState cycle xRegister 0 (Just (Addx i))) (x :: xs) =
  (MkState (S cycle) (xRegister + i) (cycles x) (Just x), xs)
tick (MkState cycle xRegister (S k) waitingOp) rest@(x :: xs) =
  (MkState (S cycle) xRegister k waitingOp, rest)

tickMultiple : State -> List Operation -> List State
tickMultiple state [] =
  if isWaiting state
    then let (newState, rest) = tick state []
             tail = tickMultiple newState rest
         in state :: tail
    else [state]
tickMultiple state ops@(x :: xs) =
  let (newState, rest) = tick state ops
      tail = tickMultiple newState rest
  in state :: tail

data CRT : Type where
  MkCrt : Vect 6 (Vect 40 Pixel) -> CRT

Show CRT where
  show (MkCrt xs) =
    let rows = map (concatMap show) xs
    in joinBy "\n" (toList rows)

emptyCrt : CRT
emptyCrt = MkCrt (replicate _ (replicate _ Off))

applyState : CRT -> State -> CRT
applyState (MkCrt xs) (MkState 0 xRegister waitCycles waitingOp) = MkCrt xs
applyState (MkCrt xs) (MkState cycle@(S k) xRegister waitCycles waitingOp) =
  let crtPos = k `mod` 40
      crtRow = k `div` 40
      xDiff = xRegister - (cast crtPos)
  in case (natToFin crtRow 6, natToFin crtPos 40) of
       (Just row, Just col) =>
         if abs xDiff <= 1
           then let oldRow = index row xs
                    newRow = replaceAt col On oldRow
                    newXs = replaceAt row newRow xs
                in MkCrt newXs
           else MkCrt xs
       _ => MkCrt xs

indexes : List Nat -> List a -> Maybe (List (Nat, a))
indexes [] xs = Just []
indexes (x :: ys) xs =
  do idx <- natToFin x (length xs)
     let value = index' xs idx
     rest <- indexes ys xs
     pure ((x, value) :: rest)

cycle : (start : Nat) -> (inc : Nat) -> (end : Nat) -> List Nat
cycle start inc end =
  if start >= end
    then []
    else start :: cycle (assert_smaller start (start + inc)) inc end

part1 : List Operation -> Maybe Int
part1 ops =
  let states = tickMultiple startState ops in
  do selectedStates <- indexes (cycle 20 40 (length states)) states
     let totalSignal = sum . map (\x => (xRegister x) * (cast (cycle x))) . map snd $ selectedStates
     pure totalSignal

part2 : List Operation -> CRT
part2 ops =
  let states = tickMultiple startState ops
  in foldl applyState emptyCrt states

main : IO ()
main =
  do Right file <- readFile "input"
       | Left err => printLn err
     Just ops <- pure (parseOps file)
       | Nothing => putStrLn "Error parsing ops"
     Just part1Result <- pure (part1 ops)
       | Nothing => putStrLn "Error in part1"
     putStrLn ("Part 1: " ++ show part1Result)
     let part2Result = part2 ops
     putStrLn "\nPart 2:"
     printLn part2Result