module VM(vmRun) where

import Control.Monad
import Data.List
import Data.Maybe
import qualified Data.Map.Strict as Map
import Data.Map.Strict ((!))
import qualified System.IO.Error as IO
import Debug.Trace

import AST
import Intermediate


data Info =
    Info (Map.Map Int BB)  -- basic blocks
         (Map.Map Name GlobFuncDef)  -- global functions
         [Value]  -- data table

type TempMap = Map.Map Int RunValue

data State = State TempMap ([RunValue] {- current arguments -}, [RunValue] {- current closure -})

data RunValue = RClosure Name [RunValue] | RValue Value
  deriving Show

kErrorExit :: String
kErrorExit = "VM:exit"

vmRun :: IRProgram -> IO ()
vmRun (IRProgram bbs gfds datas) =
    let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
        info = Info bbmap gfds datas
        state = State Map.empty ([], [])
    in IO.catchIOError (void $ vmRunBB info state (head bbs)) vmErrorHandler

vmErrorHandler :: IOError -> IO ()
vmErrorHandler e =
    if IO.isUserError e && IO.ioeGetErrorString e == kErrorExit then return () else IO.ioError e

vmRunBB :: Info -> State -> BB -> IO (RunValue, State)
vmRunBB info state (BB _ inss term) = do
    state' <- foldM (vmRunInstr info) state inss
    vmRunTerm info state' term

vmRunInstr :: Info -> State -> Instruction -> IO State
-- vmRunInstr _ _ ins | traceShow ins False = undefined
vmRunInstr info@(Info bbmap gfds datas) state@(State tmap (args, closure)) (dest, instr) = case instr of
    IAssign ref -> return (assignRef state dest (findRef tmap ref))
    IParam i -> return (assignRef state dest (args !! i))
    IClosure i -> return (assignRef state dest (closure !! i))
    IData i -> return (assignRef state dest (RValue (datas !! i)))
    ICallC cl as -> case findRef tmap cl of
        RClosure clname clvals -> case Map.lookup clname gfds of
            Just (GlobFuncDef b _ _) ->
                let Just bb = Map.lookup b bbmap
                in do
                    -- traceM ("call " ++ show cl ++ " with arguments " ++ show (map (findRef tmap) as))
                    (rv, _) <- vmRunBB info (State tmap (map (findRef tmap) as, clvals)) bb
                    return (assignRef state dest rv)
            Nothing -> liftM (assignRef state dest) (vmRunBuiltin clname (map (findRef tmap) as))
        _ -> error "VM: Cannot call non-closure object"
    IAllocClo name clrefs -> return (assignRef state dest (RClosure name (map (findRef tmap) clrefs)))
    IDiscard _ -> return state

vmRunTerm :: Info -> State -> Terminator -> IO (RunValue, State)
vmRunTerm info@(Info bbmap _ _) state@(State tmap _) term = case term of
    IBr ref b1 b2 -> vmRunBB info state . (bbmap !) $ if truthy (findRef tmap ref) then b1 else b2
    IJmp b -> vmRunBB info state (bbmap ! b)
    IRet ref -> return (findRef tmap ref, state)
    IExit -> IO.ioError (IO.userError kErrorExit)
    IUnknown -> undefined

findRef :: TempMap -> Ref -> RunValue
findRef _ (RConst n) = RValue (VNum n)
findRef tmap (RTemp i) = fromJust (Map.lookup i tmap)
findRef _ (RSClo name) = RClosure name []
findRef _ RNone = error "VM: None ref used"

assignRef :: State -> Ref -> RunValue -> State
assignRef (State tmap pair) (RTemp i) rv = State (Map.insert i rv tmap) pair
assignRef _ ref _ = error $ "VM: Unexpected ref " ++ show ref ++ " assigned"

vmRunBuiltin :: Name -> [RunValue] -> IO RunValue
-- vmRunBuiltin name args | trace (name ++ " " ++ show args) False = undefined
vmRunBuiltin "print" as = putStrLn (intercalate ", " (map printshow as)) >> return (RValue (VList []))
vmRunBuiltin "<=" [RValue (VNum a), RValue (VNum b)] = return (RValue (VNum (fromEnum (a <= b))))
vmRunBuiltin "+" [RValue (VNum a), RValue (VNum b)] = return (RValue (VNum (a + b)))
vmRunBuiltin "-" [RValue (VNum a), RValue (VNum b)] = return (RValue (VNum (a - b)))
vmRunBuiltin "car" [RValue (VList (a:_))] = return (RValue a)
vmRunBuiltin "cdr" [RValue (VList (_:a))] = return (RValue (VList a))
vmRunBuiltin name args = error (name ++ " " ++ show args)

printshow :: RunValue -> String
printshow (RValue (VString str)) = str
printshow (RValue value) = show value
printshow (RClosure _ _) = "[closure]"

truthy :: RunValue -> Bool
truthy (RValue (VNum n)) = n /= 0
truthy _ = True