module Main where

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set

import qualified IntCode as IC


type Name = String

data Ref = Imm Int | Var Name
  deriving (Show)

data Stmt
    = Declare Name
    | Set Name Ref
    | Add Name Ref Ref
    | Mul Name Ref Ref
    | Input Name
    | Output Ref
    | Clt Name Ref Ref
    | Ceq Name Ref Ref
    | If Ref [Stmt] [Stmt]
    | While Ref [Stmt]
    | Halt
  deriving (Show)

assemble :: [Stmt] -> [Integer]
assemble stmts =
    IC.unparse (IC.Jnz (IC.Imm 1) (IC.mkImm (3 + nvars)))
    ++ replicate nvars 0
    ++ concatMap IC.unparse (snd (goL nvars stmts))
  where
    vars = Set.unions (map collectDeclares stmts)
    var2idx = Map.fromList (zip (Set.toList vars) [3..])
    nvars = Set.size vars

    goL :: Int -> [Stmt] -> (Int, [IC.IC])
    goL = (fmap concat .) . foldProduce go

    go :: Int -> Stmt -> ([IC.IC], Int)
    go off (Declare _) = ([], off)
    go off (Set n a) = go off (Add n a (Imm 0))
    go off (Add n a b) = ([IC.Add (ref a) (ref b) (ref (Var n))], off + 4)
    go off (Mul n a b) = ([IC.Mul (ref a) (ref b) (ref (Var n))], off + 4)
    go off (Input n) = ([IC.Inp (ref (Var n))], off + 2)
    go off (Output a) = ([IC.Out (ref a)], off + 2)
    go off (Clt n a b) = ([IC.Clt (ref a) (ref b) (ref (Var n))], off + 4)
    go off (Ceq n a b) = ([IC.Clt (ref a) (ref b) (ref (Var n))], off + 4)
    go off (If a [] s2) =
        let (off', ics2) = goL (off + 3) s2
        in (IC.Jnz (ref a) (IC.mkImm off') : ics2
           , off')
    go off (If a s1 s2) =
        let (offAfterThen, ics2) = goL (off + 3) s2
            (offAfterElse, ics1) = goL (offAfterThen + 3) s1
        in ([IC.Jez (ref a) (IC.mkImm (offAfterThen + 3))]
            ++ ics2
            ++ [IC.Jnz (IC.Imm 1) (IC.mkImm offAfterElse)]
            ++ ics1
           , offAfterElse)
    go off (While a s) =
        let (offAfterBody, ics) = goL (off + 3) s
        in ([IC.Jez (ref a) (IC.mkImm (offAfterBody + 3))]
            ++ ics
            ++ [IC.Jnz (ref a) (IC.mkImm (off + 3))]
           , offAfterBody + 3)
    go off Halt = ([IC.Hlt], off + 1)

    ref (Imm n) = IC.mkImm n
    ref (Var n) = case Map.lookup n var2idx of
                      Just idx -> IC.Addr idx
                      Nothing  -> error $ "Undeclared variable '" ++ n ++ "'"

collectDeclares :: Stmt -> Set.Set Name
collectDeclares (Declare n) = Set.singleton n
collectDeclares (If _ a b) = Set.unions (map collectDeclares (a ++ b))
collectDeclares (While _ a) = Set.unions (map collectDeclares a)
collectDeclares _ = Set.empty

foldProduce :: (s -> a -> (b, s)) -> s -> [a] -> (s, [b])
foldProduce _ s [] = (s, [])
foldProduce f s (x:xs) = let (y, s') = f s x in fmap (y :) (foldProduce f s' xs)


main :: IO ()
main = do
    let prog = [Declare "a"
               ,Declare "c"
               ,Declare "out"
               ,Set "out" (Imm 1)
               ,Input "a"
               ,Clt "c" (Imm 0) (Var "a")
               ,While (Var "c")
                    [Mul "out" (Var "out") (Var "a")
                    ,Add "a" (Var "a") (Imm (-1))
                    ,Clt "c" (Imm 0) (Var "a")
                    ]
               ,Output (Var "out")
               ,Halt
               ]
    let res = assemble prog
    print res
    print (snd (IC.run res [10]))