module VM(vmRun) where

import Control.Monad
import Data.Char
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
    { sTempMap :: TempMap
    , sArgs :: [RunValue] {- current arguments -}
    , sCloVals :: [RunValue] {- current closure -}
    , sHandles :: Map.Map Int Handle
    , sUniq :: Int
    }

-- 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 [] [] (Map.fromList [(-1, stdin), (-2, stdout), (-3, stderr)]) 0
    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 _ _ datas) state@(State { sTempMap = tmap, sArgs = args, sCloVals = 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 -> do
        (rv, state') <- callClosure info state cl as
        return (assignRef state' dest rv)
    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 { sTempMap = 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)
    ITailC cl as -> callClosure info state cl as
    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 (RTemp i) rv = state { sTempMap = Map.insert i rv (sTempMap state) }
assignRef _ ref _ = error $ "VM: Unexpected ref " ++ show ref ++ " assigned"

callClosure :: Info -> State -> Ref -> [Ref] -> IO (RunValue, State)
callClosure info@(Info bbmap gfds _) state@(State { sTempMap = tmap }) cl as =
    case findRef tmap cl of
        RVClosure clname clvals -> case Map.lookup clname gfds of
            Just (GlobFuncDef b _ _) -> do
                (rv, state') <- vmRunBB info (state { sArgs = map (findRef tmap) as, sCloVals = clvals })
                                        (bbmap Map.! b)
                return (rv, state' { sArgs = sArgs state, sCloVals = sCloVals state })
            Nothing ->
                -- Take 'tail as' to skip the first self-link argument
                vmRunBuiltin state clname (map (findRef tmap) (tail as))
        obj -> error $ "VM: Cannot call non-closure object: " ++ show obj

vmRunBuiltin :: State -> Name -> [RunValue] -> IO (RunValue, State)
-- vmRunBuiltin _ name args | trace (name ++ " " ++ show args) False = undefined
vmRunBuiltin state "print" as = putStrLn (intercalate ", " (map printshow as)) >> return (RVList [], state)
vmRunBuiltin state "=" [a, b] = return (if equalOp a b then RVNum 1 else RVNum 0, state)
vmRunBuiltin state "<=" [RVNum a, RVNum b] = return (RVNum (fromEnum (a <= b)), state)
vmRunBuiltin state "+" [RVNum a, RVNum b] = return (RVNum (a + b), state)
vmRunBuiltin state "-" [RVNum a, RVNum b] = return (RVNum (a - b), state)
vmRunBuiltin state "*" [RVNum a, RVNum b] = return (RVNum (a * b), state)
vmRunBuiltin state "/" [RVNum a, RVNum b] = return (RVNum (a `div` b), state)
vmRunBuiltin state "mod" [RVNum a, RVNum b] = return (RVNum (a `mod` b), state)
vmRunBuiltin state "null?" [v] = return (RVNum (case v of { RVList [] -> 1; _ -> 0 }), state)
vmRunBuiltin state "car" [RVList l] = case l of
    a : _ -> return (a, state)
    _ -> throw "Empty list in 'car'"
vmRunBuiltin state "cdr" [RVList l] = case l of
    _ : a -> return (RVList a, state)
    _ -> throw "Empty list in 'cdr'"
vmRunBuiltin state "list" values = return (RVList values, state)
vmRunBuiltin state "sys-open-file" [RVNum modenum, RVString path] = do
    let mode = [ReadMode, WriteMode] !! modenum
        fid = sUniq state
    handle <- openFile path mode
    return (RVNum fid, state { sHandles = Map.insert fid handle (sHandles state), sUniq = fid + 1 })
vmRunBuiltin state "sys-close-file" [RVNum fid] = do
    hClose (sHandles state ! fid)
    return (RVList [], state { sHandles = Map.delete fid (sHandles state) })
vmRunBuiltin state "sys-get-char" [RVNum fid] = do
    let h = sHandles state ! fid
    eof <- hIsEOF h
    if eof
        then return (RVList [], state)
        else hGetChar h >>= \ch -> return (RVString [ch], state)
vmRunBuiltin state "sys-put-string" [RVNum fid, RVString str] = do
    hPutStr (sHandles state ! fid) str
    return (RVList [], state)
vmRunBuiltin state "sys-flush" [RVNum fid] = do
    hFlush (sHandles state ! fid)
    return (RVList [], state)
vmRunBuiltin state "sys-stdin" [] = return (RVNum (-1), state)
vmRunBuiltin state "sys-stdout" [] = return (RVNum (-2), state)
vmRunBuiltin state "sys-stderr" [] = return (RVNum (-3), state)
vmRunBuiltin state "length" [RVString str] = return (RVNum (length str), state)
vmRunBuiltin state "substr" [RVString str, RVNum idx, RVNum len] =
    return (RVString (take len (drop idx str)), state)
vmRunBuiltin state "ord" [RVString str] = return (RVNum (case str of { "" -> 0; c:_ -> ord c }), state)
vmRunBuiltin state "chr" [RVNum num] = return (RVString [chr num], state)
vmRunBuiltin state "concat" values
    | Just strings <- sequence (map fromRVString values) = return (RVString (concat strings), state)
    | otherwise = throw "Non-string arguments to 'concat'"
vmRunBuiltin _ name args = error (name ++ " " ++ show args)

equalOp :: RunValue -> RunValue -> Bool
equalOp (RVClosure _ _) _ = error "Cannot compare closures in '='"
equalOp _ (RVClosure _ _) = error "Cannot compare closures in '='"
equalOp (RVList vs) (RVList ws) = length vs == length ws && all id (zipWith equalOp vs ws)
equalOp (RVNum a) (RVNum b) = a == b
equalOp (RVString s) (RVString t) = s == t
equalOp (RVQuoted v) (RVQuoted w) = equalOp v w
equalOp (RVName n) (RVName m) = n == m
equalOp _ _ = False

printshow :: RunValue -> String
printshow (RVString str) = str
printshow (RVList values) = "[" ++ intercalate "," (map printshow values) ++ "]"
printshow (RVNum i) = show i
printshow (RVQuoted value) = '\'' : printshow 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)

fromRVString :: RunValue -> Maybe String
fromRVString (RVString str) = Just str
fromRVString _ = Nothing