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 -})

-- TODO: are more constructors from Value needed?
data RunValue
    = RVClosure Name [RunValue]
    | RVList [RunValue]
    | RVNum Int
    | RVString String
    | RVQuoted RunValue
  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 ->
        if i < length args then return (assignRef state dest (args !! i))
            else error $ show closure ++ ", " ++ show i ++ ", param-out-of-range"
    IClosure i ->
        if i < length closure then return (assignRef state dest (closure !! i))
            else error $ show closure ++ ", " ++ show i ++ ", closure-out-of-range"
    IData i ->
        if i < length datas then return (assignRef state dest (toRunValue (datas !! i)))
            else error "data-out-of-range"
    ICallC cl as ->
        -- trace ("callc " ++ show (findRef tmap cl) ++ "  " ++ show (map (findRef tmap) as)) $
        case findRef tmap cl of
            RVClosure 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))
            obj -> error $ "VM: Cannot call non-closure object: " ++ show obj
    IAllocClo name clrefs -> return (assignRef state dest (RVClosure 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) = RVNum n
findRef tmap (RTemp i) = fromJust (Map.lookup i tmap)
findRef _ (RSClo name) = RVClosure 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 (RVList [])
vmRunBuiltin "<=" [RVNum a, RVNum b] = return (RVNum (fromEnum (a <= b)))
vmRunBuiltin "+" [RVNum a, RVNum b] = return (RVNum (a + b))
vmRunBuiltin "-" [RVNum a, RVNum b] = return (RVNum (a - b))
vmRunBuiltin "car" [RVList (a:_)] = return a
vmRunBuiltin "cdr" [RVList (_:a)] = return (RVList a)
vmRunBuiltin "list" values = return (RVList values)
vmRunBuiltin name args = error (name ++ " " ++ show args)

printshow :: RunValue -> String
printshow (RVString str) = str
printshow (RVList values) = show values
printshow (RVNum i) = show i
printshow (RVQuoted value) = '\'' : show value
printshow (RVClosure _ _) = "[closure]"

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

toRunValue :: Value -> RunValue
toRunValue (VList values) = RVList (map toRunValue values)
toRunValue (VNum i) = RVNum i
toRunValue (VString s) = RVString s
toRunValue (VQuoted value) = RVQuoted (toRunValue value)
toRunValue _ = undefined