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 System.IO
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
    | RVName Name
  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 l] = case l of
    a : _ -> return a
    _ -> throw "Empty list in 'car'"
vmRunBuiltin "cdr" [RVList l] = case l of
    _ : a -> return (RVList a)
    _ -> throw "Empty list in 'cdr'"
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]"
printshow (RVName name) = name

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 (VName name) = RVName name
toRunValue _ = undefined

throw :: String -> IO a
throw s = hPutStrLn stderr s >> IO.ioError (IO.userError kErrorExit)