-- Author: Yuchen Pei (me@ypei.me)
import Data.List.Split (splitOn)
import Data.Map (Map)
import qualified Data.Map as Map

exec :: [[Char]] -> Int -> Int -> Map Char Int -> Int
exec rom n addr regs
  | addr >= length rom || addr < 0 = n
  | op == "set" = exec rom n (addr + 1) (Map.insert dest val regs)
  | op == "sub" = exec rom n (addr + 1) (Map.insert dest (dval - val) regs)
  | op == "mul" = exec rom (n + 1) (addr + 1) (Map.insert dest (dval * val) regs)
  | op == "jnz" = exec rom n (if dval /= 0 then addr + val else addr + 1) regs
      where ins = rom !! addr
            op:[dest]:xs = splitOn " " ins
            dval = getVal [dest]
            val = getVal $ head xs
            getVal xs = if head xs `elem` alphabet then regs Map.! (head xs) else read xs

isPrime :: Int -> Bool
isPrime x = and $ ((/=0) . rem x) <$> [2 .. floor $ sqrt $ fromIntegral x]

solve2 = length $ filter (==False) $ isPrime <$> [105700, 105717 .. 122700] -- The last number, 122700 should not be included because according to the assembly code when b == c the program terminates without checking primality of b. But the AOC website only accepted the answer when it is included.

alphabet = ['a'..'h']

initRegs :: Map Char Int
initRegs = Map.fromList $ zip ['a'..'h'] (cycle [0])


solve1 :: [Char] -> Int
solve1 xs = exec (lines xs) 0 0 initRegs


input = "set b 57\nset c b\njnz a 2\njnz 1 5\nmul b 100\nsub b -100000\nset c b\nsub c -17000\nset f 1\nset d 2\nset e 2\nset g d\nmul g e\nsub g b\njnz g 2\nset f 0\nsub e -1\nset g e\nsub g b\njnz g -8\nsub d -1\nset g d\nsub g b\njnz g -13\njnz f 2\nsub h -1\nset g b\nsub g c\njnz g 2\njnz 1 3\nsub b -17\njnz 1 -23"