{-# LANGUAGE TupleSections, LambdaCase #-}
module VM(vmRun) where

import qualified Data.Array.IO as A
import Control.Monad
import Data.Char
import Data.List
import qualified Data.Map.Strict as Map
import qualified Data.IntMap.Strict as IMap
import System.IO
import qualified System.IO.Error as IO
-- import Debug.Trace

import AST
import qualified DString as DS
import Intermediate


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

type TempMap = A.IOArray Int RunValue

data State = State
    { sTempMap :: TempMap
    , sArgs :: [RunValue] {- current arguments -}
    , sCloVals :: [RunValue] {- current closure -}
    , sStack :: [RunValue] {- IPush/IPop stack -}
    , sHandles :: IMap.IntMap Handle
    , sUniq :: Int
    }

-- TODO: are more constructors from Value needed?
data RunValue
    = RVClosure Name [RunValue]
    | RVList [RunValue]
    | RVNum Int
    | RVString DS.DString
    | RVQuoted RunValue
    | RVName Name
  deriving Show

kErrorExit :: String
kErrorExit = "VM:exit"

vmRun :: IRProgram -> IO ()
vmRun irprogram@(IRProgram bbs gfds datas) = do
    let alltemps = onlyTemporaries (allRefs irprogram)
    tmap <- A.newArray (minimum alltemps, maximum alltemps) (RVNum 0)
    let bbmap = Map.fromList [(bidOf bb, bb) | bb <- bbs]
        info = Info bbmap gfds datas
        state = State tmap [] [] [] (IMap.fromList [(-1, stdin), (-2, stdout), (-3, stderr)]) 0
    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 (dest, instr) = case instr of
    IAssign ref -> do
        findRef (sTempMap state) ref >>= assignRef state dest
        return state
    IParam i -> do
        let args = sArgs state
        if i < length args then assignRef state dest (args !! i)
            else error $ show args ++ ", " ++ show i ++ ", param-out-of-range"
        return state
    IClosure i -> do
        let closure = sCloVals state
        if i < length closure then assignRef state dest (closure !! i)
            else error $ show closure ++ ", " ++ show i ++ ", closure-out-of-range"
        return state
    IData i -> do
        if i < length datas then assignRef state dest (toRunValue (datas !! i))
            else error "data-out-of-range"
        return state
    ICallC cl as -> do
        (rv, state') <- callClosure info state cl as
        assignRef state' dest rv
        return state'
    IAllocClo name clrefs -> do
        clovals <- mapM (findRef (sTempMap state)) clrefs
        assignRef state dest (RVClosure name clovals)
        return state
    IDiscard ref -> do
        case ref of
            RTemp _ -> assignRef state ref (RVNum 0)
            _ -> return ()
        return state
    IPush refs -> do
        values <- mapM (findRef (sTempMap state)) refs
        return (state { sStack = values ++ sStack state })
    IPop refs -> do
        when (length (sStack state) < length refs) $
            error "VM: IPop on too-small stack"
        let (values, newStack) = splitAt (length refs) (sStack state)
            state' = state { sStack = newStack }
        mapM_ (uncurry (assignRef state')) (zip refs values)
        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 -> do
        val <- findRef tmap ref
        vmRunBB info state . (bbmap Map.!) $ if truthy val then b1 else b2
    IJmp b -> vmRunBB info state (bbmap Map.! b)
    IRet ref -> (,state) <$> findRef tmap ref
    ITailC cl as -> callClosure info state cl as
    IExit -> IO.ioError (IO.userError kErrorExit)
    IUnknown -> undefined

findRef :: TempMap -> Ref -> IO RunValue
findRef _ (RConst n) = return (RVNum n)
findRef tmap (RTemp i) = A.readArray tmap i
findRef _ (RSClo name) = return (RVClosure name [])
findRef _ RNone = error "VM: None ref used"

assignRef :: State -> Ref -> RunValue -> IO ()
assignRef state (RTemp i) rv = A.writeArray (sTempMap state) i rv
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 =
    findRef tmap cl >>= \case
        RVClosure clname clvals -> case Map.lookup clname gfds of
            Just (GlobFuncDef b _ _) -> do
                args <- mapM (findRef tmap) as
                (rv, state') <- vmRunBB info (state { sArgs = args, sCloVals = clvals }) (bbmap Map.! b)
                return (rv, state' { sArgs = sArgs state, sCloVals = sCloVals state })
            Nothing -> do
                -- Take 'tail as' to skip the first self-link argument
                args <- mapM (findRef tmap) (tail as)
                vmRunBuiltin state clname args
        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] = case equalOp a b of
    Left err -> throw err
    Right True -> return (RVNum 1, state)
    Right False -> return (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 "cons" [val, RVList l] = return (RVList (val : l), state)
vmRunBuiltin state "sys-open-file" [RVNum modenum, RVString path] = do
    let mode = [ReadMode, WriteMode] !! modenum
        fid = sUniq state
    handle <- openFile (DS.unpack path) mode
    return (RVNum fid, state { sHandles = IMap.insert fid handle (sHandles state), sUniq = fid + 1 })
vmRunBuiltin state "sys-close-file" [RVNum fid] = do
    hClose (sHandles state IMap.! fid)
    return (RVList [], state { sHandles = IMap.delete fid (sHandles state) })
vmRunBuiltin state "sys-get-char" [RVNum fid] = do
    let h = sHandles state IMap.! fid
    eof <- hIsEOF h
    if eof
        then return (RVList [], state)
        else hGetChar h >>= \ch -> return (RVString (DS.singleton ch), state)
vmRunBuiltin state "sys-put-string" [RVNum fid, RVString str] = do
    DS.hPutStr (sHandles state IMap.! fid) str
    return (RVList [], state)
vmRunBuiltin state "sys-flush" [RVNum fid] = do
    hFlush (sHandles state IMap.! 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 (DS.length str), state)
vmRunBuiltin state "substr" [RVNum idx, RVNum len, RVString str] =
    let s = (if len >= 0 then DS.take len else id) (DS.drop idx str)
    in return (RVString s, state)
vmRunBuiltin state "ord" [RVString str] = return (RVNum (if DS.null str then 0 else ord (DS.head str)), state)
vmRunBuiltin state "chr" [RVNum num] = return (RVString (DS.singleton (chr num)), state)
vmRunBuiltin state "concat" values
    | Just strings <- sequence (map fromRVString values) = return (RVString ({-# SCC builtin_string_concat #-} DS.concat strings), state)
    | otherwise = throw "Non-string arguments to 'concat'"
vmRunBuiltin state "type-list?" [value]   = return (RVNum (case value of { RVList _   -> 1; _ -> 0 }), state)
vmRunBuiltin state "type-number?" [value] = return (RVNum (case value of { RVNum _    -> 1; _ -> 0 }), state)
vmRunBuiltin state "type-string?" [value] = return (RVNum (case value of { RVString _ -> 1; _ -> 0 }), state)
vmRunBuiltin state "type-quoted?" [value] = return (RVNum (case value of { RVQuoted _ -> 1; _ -> 0 }), state)
vmRunBuiltin state "type-symbol?" [value] = return (RVNum (case value of { RVName _   -> 1; _ -> 0 }), state)
vmRunBuiltin _ "error" values = throw ("error: " ++ intercalate " " (map show values))
vmRunBuiltin _ name args = error (name ++ " " ++ show args)

equalOp :: RunValue -> RunValue -> Either String Bool
equalOp (RVClosure _ _) _ = Left "Cannot compare closures in '='"
equalOp _ (RVClosure _ _) = Left "Cannot compare closures in '='"
equalOp (RVList vs) (RVList ws)
    | length vs == length ws = all id <$> sequence (zipWith equalOp vs ws)
    | otherwise = Right False
equalOp (RVNum a) (RVNum b) = Right (a == b)
equalOp (RVString s) (RVString t) = Right (s == t)
equalOp (RVQuoted v) (RVQuoted w) = equalOp v w
equalOp (RVName n) (RVName m) = Right (n == m)
equalOp _ _ = Right False

printshow :: RunValue -> String
printshow (RVString str) = DS.unpack 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 (DS.pack 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 DS.DString
fromRVString (RVString str) = Just str
fromRVString _ = Nothing