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

data Operation = Noop | Addx Int
%name Operation op

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
%name Pixel pixel

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

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

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

extraSimple = """
  noop
  addx 3
  addx -5
"""

simple = """
  addx 15
  addx -11
  addx 6
  addx -3
  addx 5
  addx -1
  addx -8
  addx 13
  addx 4
  noop
  addx -1
  addx 5
  addx -1
  addx 5
  addx -1
  addx 5
  addx -1
  addx 5
  addx -1
  addx -35
  addx 1
  addx 24
  addx -19
  addx 1
  addx 16
  addx -11
  noop
  noop
  addx 21
  addx -15
  noop
  noop
  addx -3
  addx 9
  addx 1
  addx -3
  addx 8
  addx 1
  addx 5
  noop
  noop
  noop
  noop
  noop
  addx -36
  noop
  addx 1
  addx 7
  noop
  noop
  noop
  addx 2
  addx 6
  noop
  noop
  noop
  noop
  noop
  addx 1
  noop
  noop
  addx 7
  addx 1
  noop
  addx -13
  addx 13
  addx 7
  noop
  addx 1
  addx -33
  noop
  noop
  noop
  addx 2
  noop
  noop
  noop
  addx 8
  noop
  addx -1
  addx 2
  addx 1
  noop
  addx 17
  addx -9
  addx 1
  addx 1
  addx -3
  addx 11
  noop
  noop
  addx 1
  noop
  addx 1
  noop
  noop
  addx -13
  addx -19
  addx 1
  addx 3
  addx 26
  addx -30
  addx 12
  addx -1
  addx 3
  addx 1
  noop
  noop
  noop
  addx -9
  addx 18
  addx 1
  addx 2
  noop
  noop
  addx 9
  noop
  noop
  noop
  addx -1
  addx 2
  addx -37
  addx 1
  addx 3
  noop
  addx 15
  addx -21
  addx 22
  addx -6
  addx 1
  noop
  addx 2
  addx 1
  noop
  addx -10
  noop
  noop
  addx 20
  addx 1
  addx 2
  addx 2
  addx -6
  addx -11
  noop
  noop
  noop
  """

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

testPart1 : IO ()
testPart1 =
  do Just ops <- pure $ parseOps simple
       | Nothing => putStrLn "Error parsing ops"
     let states = tickMultiple startState ops
     let idxs = cycle 20 40 (length states)
     Just selectedStates <- pure . map (map snd) $ indexes idxs states
       | Nothing => putStrLn "Not enough cycles"
     putStrLn "Selected States:"
     traverse_ printLn selectedStates
     let totalSignal = sum . map (\x => (xRegister x) * (cast (cycle x))) $ selectedStates
     putStrLn ("Total Signal Strength: " ++ show totalSignal)

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

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)