#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple vector containers])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-local-binds #-}
{-# LANGUAGE OverloadedStrings #-}

import Text.Pretty.Simple
-- import Data.List.NonEmpty (NonEmpty(..))
-- import qualified Data.List.NonEmpty as NE
import Data.Vector (Vector)
import Debug.Trace (trace)
import Data.List (unfoldr, zip4, unzip4)
import Data.IntMap.Strict (IntMap)
import qualified Data.Vector as V
import qualified Data.IntMap.Strict as M

data Op = OpNop !Int
        | OpJmp !Int
        | OpAcc !Int
  deriving (Show, Ord, Eq)

data VMState = VMRun | VMHalt
  deriving (Show, Ord, Eq)

data VirtualMachine = VM
  { pgmCounter      :: !Int
  , pgmAccumulator  :: !Int
  , pgmInstructions :: Vector Op
  , vmState         :: VMState
  }
  deriving (Show, Ord, Eq)

initialVM :: Vector Op -> VirtualMachine
initialVM ins = VM 1 0 ins VMRun

stepVM :: VirtualMachine -> Maybe (VirtualMachine, VirtualMachine)
stepVM current@(VM c a ins VMRun) = Just (current, checkHalt next)
  where
    next = go (ins V.! (c-1))

    go (OpNop _)        = VM (c+1) a     ins VMRun
    go (OpJmp n)        = VM (c+n) a     ins VMRun
    go (OpAcc n)        = VM (c+1) (a+n) ins VMRun

    checkHalt vm0@(VM c0 a0 _ _) | c0 == lastCount + 1 =
      trace ("halted, acc="<>show a0) $ vm0 { vmState = VMHalt }
    checkHalt vm0@(VM _ _ _ _)  | otherwise =
      vm0

    lastCount = V.length ins
stepVM (VM _ _ _ VMHalt) = Nothing

exampleData :: [String]
exampleData =
  [ "nop +0"
  , "acc +1"
  , "jmp +4"
  , "acc +3"
  , "jmp -3"
  , "acc -99"
  , "acc +1"
  , "jmp -4"
  , "acc +6"
  ]

vmFromInput :: [String] -> VirtualMachine
vmFromInput input = initialVM (ins input)
  where
    ins = V.fromList . map parseOp

parseOp :: String -> Op
parseOp = go . (\(x,y)-> (x,tail y)) . break (== ' ')
  where
    go ("nop",'+':s) = OpNop (read s :: Int)
    go ("acc",'+':s) = OpAcc (read s :: Int)
    go ("jmp",'+':s) = OpJmp (read s :: Int)
    go ("nop",s) = OpNop (read s :: Int)
    go ("acc",s) = OpAcc (read s :: Int)
    go ("jmp",s) = OpJmp (read s :: Int)
    go s@_ = trace (show s) $ undefined

-- Find smallest looping prefix in list of Int, returns associated pgmAccumulator
solvePart1 :: VirtualMachine -> Int
solvePart1 vm0 =
  snd . last $ shortestCycleOn $ (\vm -> (pgmCounter vm, pgmAccumulator vm)) <$> unfoldr stepVM vm0

-- Shortest cycle in [Int]
shortestCycle :: [Int] -> [Int]
shortestCycle ix = reverse $ go M.empty ix []
  where
    go :: (IntMap ()) -> [Int] -> [Int] -> [Int]
    go mem (x:_ ) ys | (mem M.!? x) /= Nothing = ys
    go mem (x:xs) ys | otherwise               = go (M.insert x () mem) xs (x:ys)
    go _ [] ys                                 = ys

shortestCycleOn :: [(Int,Int)] -> [(Int,Int)]
shortestCycleOn tx = zip (shortestCycle cx) ax
  where
    (cx, ax) = unzip tx

shortestCycleVm :: [VirtualMachine] -> [VirtualMachine]
shortestCycleVm vx = (\(a,b,c,d) -> VM a b c d) <$> zip4 (shortestCycle cx) ax ins state
  where
    (cx, ax, ins, state) = unzip4 $ (\(VM a b c d) -> (a,b,c,d)) <$> vx

-- Try lots of permutations of OpNop -> OpJmp, or OpJmp -> OpNop
-- until the VirtualMachine halts, and read the value of the accumulator
solvePart2 :: Vector Op -> Vector VMState
solvePart2 = fmap (vmState . last . shortestCycleVm . unfoldr stepVM . initialVM) . swapInstructions

swapInstructions :: Vector Op -> Vector (Vector Op)
swapInstructions ins = V.imap swap ins
  where
    swap i (OpNop v) = ins V.// [(i,OpJmp v)]
    swap i (OpJmp v) = ins V.// [(i,OpNop v)]
    swap _ (OpAcc _) = ins

main :: IO ()
main = do
  putStrLn "Day 8 - Part 1"
  putStrLn ":: Tests"
  print exampleData
  print $ parseOp <$> exampleData
  putStrLn ":: Solving part 1"
  pPrint $ solvePart1 (vmFromInput exampleData)
  input <- lines <$> readFile "day8/input"
  pPrint $ solvePart1 (vmFromInput input)
  putStrLn ":: Tests"
  let miniIns = V.fromList [OpJmp 2, OpNop 0, OpAcc (-8)]
  pPrint $ length $ swapInstructions $ miniIns
  let miniVM = initialVM $ V.fromList [ OpNop 0 ]
  pPrint $ miniVM { vmState = VMHalt }
  pPrint $ (unfoldr stepVM) $ miniVM
  putStrLn ":: Solving part 2"
  pPrint $ solvePart2 $ V.fromList $ parseOp <$> exampleData
  pPrint $ solvePart2 $ V.fromList $ parseOp <$> input