data Instruction
  = Jmp Int
  | Acc Int
  | Nop Int
  deriving (Eq, Show)

data Execution
  = Halt
  | Continue
  | Loop
  deriving (Show, Eq)

type Visited = S.Set Int
type Accumulator = Int
type Cpu = (Top :>> [Instruction] :>> Instruction, Accumulator, Visited)

day8 :: IO ()
day8 = do
  code <- fromWithin traverse . zipper . parse <$> input "day8.txt"
  print $ view _1 <$> traverse runState [part1, part2] (code, 0, mempty)
  where
    parse =
      fromRight (error "impossible")
      . A.eitherResult
      . A.parse parseInstructions
      . T.unlines

    movePointer :: Int -> State Cpu Execution
    movePointer i = do
      (z, a, v) <- get
      let nextPosition = focalPoint z + i
      if S.member nextPosition v then
        pure Loop
      else
        case moveTo nextPosition z of
          Just z' -> do
            put (z', a, S.insert nextPosition v)
            pure Continue
          Nothing -> do
            pure Halt

    part1 :: State Cpu (Execution, Int)
    part1 = do
      s <- get
      case s ^. _1 . focus of
        (Jmp i) -> do
          r <- movePointer i
          case r of
            Continue -> part1
            x        -> pure (x, s ^. _2)
        (Acc i) -> do
          modify (_2 %~ (+ i))
          movePointer 1
          part1
        (Nop _) -> do
          movePointer 1
          part1

    explore :: Int -> Instruction -> State Cpu (Execution, Int)
    explore i branch = do
      s <- get
      let (e, a) = evalState part1 (s & _1 %~ (focus .~ branch))
      case e of
        Halt -> do
          pure (e, a)
        _ -> do
          movePointer i
          part2

    part2 :: State Cpu (Execution, Int)
    part2 = do
      s <- get
      case s ^. _1 . focus of
        (Jmp i) -> do
          explore i (Nop i)
        (Nop 0) -> do
          movePointer 1
          part2
        (Nop i) -> do
          explore 1 (Jmp i)
        (Acc i) -> do
          modify (_2 %~ (+ i))
          movePointer 1
          part2

    parseInstructions =
      parseInstruction `A.sepBy` A.space

    parseInstruction = do
      i <- A.choice ["acc" $> Acc, "jmp" $> Jmp, "nop" $> Nop]
      A.skipSpace
      v <- A.signed A.decimal
      pure $ i v