{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
module IntCode (
    Arg(..), IC(..),
    parse, decode, run, unparse,
    runInterruptible, runContinue, Continuation()
) where

import Control.Monad.ST
import qualified Data.Array as A
import qualified Data.Array.ST as SA
import qualified Data.Array.IArray as IA


data Arg = Imm Int | Addr Int
  deriving (Show)

data IC
    = Add Arg Arg Arg
    | Mul Arg Arg Arg
    | Inp Arg
    | Out Arg
    | Jnz Arg Arg
    | Jez Arg Arg
    | Clt Arg Arg Arg
    | Ceq Arg Arg Arg
    | Hlt
  deriving (Show)

parse :: String -> [Int]
parse = map read . splitOn ','
  where
    splitOn c s = case break (== c) s of
                      (pre, _ : post) -> pre : splitOn c post
                      _ -> [s]

unparse :: IC -> [Int]
unparse ic = case ic of
                 Add a b c -> go 1 [a,b,c]
                 Mul a b c -> go 2 [a,b,c]
                 Inp a     -> go 3 [a    ]
                 Out a     -> go 4 [a    ]
                 Jnz a b   -> go 5 [a,b  ]
                 Jez a b   -> go 6 [a,b  ]
                 Clt a b c -> go 7 [a,b,c]
                 Ceq a b c -> go 8 [a,b,c]
                 Hlt       -> [99]
  where
    go code as = (100 * mode as + code) : map bare as
      where bare (Imm n) = n
            bare (Addr n) = n
    mode [] = 0
    mode (Imm _ : as) = 1 + 10 * mode as
    mode (Addr _ : as) = 10 * mode as

decode :: [Int] -> (IC, Int)
decode [] = error "IC: Execution fell off end of program"
decode (ins : rest) =
    let (code, modes) = insModeSplit ins
    in (decodeBase code modes rest, 1 + insNArgs code)

class Monad m => ReadArray m a i e where
    readArray :: a -> i -> m e

instance IA.Ix i => ReadArray (ST s) (SA.STUArray s i Int) i Int where
    readArray = SA.readArray

decode' :: ReadArray m a Int Int => a -> Int -> m (IC, Int)
decode' arr ptr = do
    ins <- readArray arr ptr
    let (code, modes) = insModeSplit ins
        nargs = insNArgs code
    params <- sequence [readArray arr (ptr + i) | i <- [1..nargs]]
    return (decodeBase code modes params, 1 + nargs)

insModeSplit :: Int -> (Int, [Int])
insModeSplit n = (n `mod` 100, map (`mod` 10) (iterate (`div` 10) (n `div` 100)))

insNArgs :: Int -> Int
insNArgs n = case n of { 1 -> 3; 2 -> 3; 3 -> 1; 4 -> 1;
                         5 -> 2; 6 -> 2; 7 -> 3; 8 -> 3;
                         99 -> 0;
                         _ -> error $ "IC: Unknown instruction " ++ show n }

decodeBase :: Int -> [Int] -> [Int] -> IC
decodeBase  1 (am:bm:cm:_) (a:b:c:_) = Add (insArg am a) (insArg bm b) (insArg cm c)
decodeBase  2 (am:bm:cm:_) (a:b:c:_) = Mul (insArg am a) (insArg bm b) (insArg cm c)
decodeBase  3 (am:      _) (a:    _) = Inp (insArg am a)
decodeBase  4 (am:      _) (a:    _) = Out (insArg am a)
decodeBase  5 (am:bm:   _) (a:b:  _) = Jnz (insArg am a) (insArg bm b)
decodeBase  6 (am:bm:   _) (a:b:  _) = Jez (insArg am a) (insArg bm b)
decodeBase  7 (am:bm:cm:_) (a:b:c:_) = Clt (insArg am a) (insArg bm b) (insArg cm c)
decodeBase  8 (am:bm:cm:_) (a:b:c:_) = Ceq (insArg am a) (insArg bm b) (insArg cm c)
decodeBase 99 (         _) (      _) = Hlt
decodeBase ic _            _         = error $ "IC: Unknown instruction " ++ show ic

insArg :: Int -> Int -> Arg
insArg 0 n = Addr n
insArg 1 n = Imm n
insArg m n = error $ "IC: Unknown parameter mode " ++ show m ++ " for parameter " ++ show n

run :: [Int] -> [Int] -> ([Int], [Int])
run initMem input =
    case runInterruptible initMem input of
        Left _ -> error "IC: Not enough input"
        Right res -> res

data Continuation = Continuation (A.Array Int Int) Int

runInterruptible :: [Int] -> [Int] -> Either (Continuation, [Int]) ([Int], [Int])
runInterruptible initMem input = runST $ do
    arr <- SA.newListArray (0, length initMem - 1) initMem
    res <- runArray arr 0 input
    case fst res of
        Left cont -> return (Left (cont, snd res))
        Right mem -> return (Right (mem, snd res))

runContinue :: Continuation -> [Int] -> Either (Continuation, [Int]) ([Int], [Int])
runContinue (Continuation frozen ptr) moreinput = runST $ do
    arr <- SA.thaw frozen
    res <- runArray arr ptr moreinput
    case fst res of
        Left cont -> return (Left (cont, snd res))
        Right mem -> return (Right (mem, snd res))

runArray :: SA.STUArray s Int Int -> Int -> [Int] -> ST s (Either Continuation [Int], [Int])
runArray arr ptr inp = do
    (ic, len) <- decode' arr ptr
    let continueInp inp' = runArray arr (ptr + len) inp'
        continue = continueInp inp
        jump ptr' = runArray arr ptr' inp
    case ic of
        Add a b (Addr c) -> (+) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
        Mul a b (Addr c) -> (*) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
        Inp (Addr a) -> case inp of
                     val : rest -> SA.writeArray arr a val >> continueInp rest
                     [] -> do
                         frozen <- SA.freeze arr
                         return (Left (Continuation frozen ptr), [])
        Out a -> resolve a >>= \val -> fmap (val :) <$> continue
        Jnz a b -> resolve a >>= \val -> if val /= 0 then resolve b >>= jump else continue
        Jez a b -> resolve a >>= \val -> if val == 0 then resolve b >>= jump else continue
        Clt a b (Addr c) -> ((fromEnum .) . (<)) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
        Ceq a b (Addr c) -> ((fromEnum .) . (==)) <$> resolve a <*> resolve b >>= SA.writeArray arr c >> continue
        Hlt -> do
            mem <- SA.getElems arr
            return (Right mem, [])
        ins -> error $ "IC: Invalid instruction " ++ show ins
  where
    resolve (Imm n) = return n
    resolve (Addr n) = SA.readArray arr n