{-# LANGUAGE TupleSections #-}
module Eval
  (runProgram)
  where
import Syntax
import Typecheck ()

import Control.Monad.State.Lazy
import Control.Monad.Except
import Control.Monad.Reader
data Value = VInt Int
            | VRecord  [(Idnt,Value)]
            | VVariant Idnt Value
            | VPtr Int
            deriving(Eq,Show)
false = VVariant "false" (VRecord [])
true = VVariant "true" (VRecord [])
type Variables = [(Idnt,Value)]
type Heap = (Int,[(Int,Value)])

data ProgramState= ProgramState Variables Heap
  deriving (Show)

type EvalMonad = ReaderT [(Idnt,FuncDecl)] (ExceptT String (State ProgramState))
--type EvalMonad = ExceptT String (State ProgramState)


assertInt :: Value -> EvalMonad Int
assertInt (VInt i) = return i
assertInt v = throwError $ "expected value to be int " ++ show v
assertRec :: Value -> EvalMonad [(Idnt,Value)]
assertRec (VRecord l) = return l
assertRec v = throwError $ "expected value to be int " ++ show v
assertVar :: Value -> EvalMonad (Idnt,Value)
assertVar (VVariant label val) = return (label,val)
assertVar v = throwError $ "expected value to be int " ++ show v
assertPtr :: Value -> EvalMonad Int
assertPtr (VPtr i) = return i
assertPtr v = throwError $ "expected value to be ptr " ++ show v

evalExpr :: Expr ->EvalMonad Value
evalExpr (Var v)= do
  (ProgramState vars _) <- get
  case lookup v vars of
    Just val -> return val
    Nothing -> throwError $ "Variable" ++show v ++ "not defined"
evalExpr (IntLit n) = return $ VInt n
evalExpr (Record l) = do
  fields <- mapM (\(idnt,expr) ->do
                     val <- evalExpr expr
                     return (idnt,val) ) l
  return $ VRecord fields
evalExpr (Variant _ idnt e) = do
  v <- evalExpr e
  return $ VVariant idnt v
evalExpr (Add e1 e2) = do
  v1 <- evalExpr e1
  v2 <- evalExpr e2
  i1 <- assertInt v1
  i2 <- assertInt v2
  return $ VInt $ i1 + i2

evalExpr (Mult e1 e2) = do
  v1 <- evalExpr e1
  v2 <- evalExpr e2
  i1 <- assertInt v1
  i2 <- assertInt v2
  return $ VInt $ i1 + i2
evalExpr (Neg e) = do
  v <- evalExpr e
  i <- assertInt v
  return $ VInt i
evalExpr (Deref e) = do
  v <- evalExpr e
  ptrVal <-assertPtr v
  (ProgramState _ (_,heap)) <- get
  case lookup ptrVal heap of
    Just val -> return val
    Nothing -> throwError $ "invalid pointer" 
evalExpr c@(Match e cases) = do
  v <- evalExpr e
  (label,value) <- assertVar v
  case lookup label (map (\(a,b,c)-> (a,(b,c)) )cases) of
    Just (bind,expr) -> do
      st@(ProgramState vars heap) <- get
      put $ ProgramState ((bind,value):vars) heap
      ret <- evalExpr expr
      put st
      return ret
    Nothing -> throwError $ "Unsupported label"++ show label ++ "in case" ++ show c
evalExpr (Proj label e) = do
  v <- evalExpr e
  l <- assertRec v
  case lookup label l of
    Just v -> return v
    Nothing -> throwError $ "Projection "++ show label ++ "on record" ++ show v
  
evalExpr (Annot e _) = evalExpr e

changeDict ::Eq a=> a -> b -> [(a,b)] -> Maybe [(a,b)]
changeDict _ _ []= Nothing
changeDict var val ((a,b):xs) = if a==var then (Just $ (a,val):xs) else
  do
    ret <- changeDict var val xs
    return $ (a,b):ret
ensureTop :: (Show a, Eq a, Show b) => a -> [(a,b)] -> EvalMonad [(a,b)]
ensureTop var [] = throwError $ "Expected variable " ++ show var ++ " at the end of lexical scope, instead got empty set of variables" 
ensureTop var ((a,b):xs) = if a==var then return xs
  else throwError $ "Expected variable " ++ show var ++ " at the end of lexical scope, instead got" ++ show b
  
 
  

runCom :: Com -> EvalMonad ProgramState
runCom Skip = do
  get
runCom (Seq c1 c2) =do
  runCom c1
  runCom c2
runCom (If e1 c1 c2) = do
  v1 <- evalExpr e1
  if v1 == true  then runCom c1
    else if v1 ==false then runCom c2
         else throwError $ "Not boolean in if "++ show v1
runCom c@(While e1 c1) = do
  v1 <- evalExpr e1
  if v1 == true  then do
      runCom c1
      runCom c
    else if v1 ==false then runCom Skip
         else throwError $ "Not boolean in if "++ show v1
runCom (Asgn var expr) = do
  val <- evalExpr expr
  (ProgramState vars heap) <- get
  case changeDict var val vars of
    Just newVars -> let p = ProgramState newVars heap in
      do
        put p
        return p
    Nothing -> throwError $ "Assigning to not-declared variable " ++ show var
runCom (Decl var expr com) = do
  val <- evalExpr expr
  (ProgramState vars heap) <- get
  put $ ProgramState ((var,val):vars) heap
  runCom com
  --remove new variable
  (ProgramState vars heap) <- get
  vars' <- ensureTop var vars
  put $ ProgramState vars' heap
  get
runCom (Alloc var expr com) = do
  val <- evalExpr expr
  (ProgramState vars (heapCnt,heap)) <- get
  put $ ProgramState ((var,VPtr heapCnt):vars) (heapCnt+1,((heapCnt,val):heap))
  runCom com
  --remove new variable
  (ProgramState vars heap) <- get
  vars' <- ensureTop var vars
  put $ ProgramState vars' heap
  get
runCom (Save var expr) = do
  val <- evalExpr expr
  ptr <- evalExpr (Var var)
  ptrVal <- assertPtr ptr
  (ProgramState vars (heapCnt,heap)) <- get
  heap' <- case changeDict ptrVal val heap of
    Nothing -> throwError $ "invalid pointer"
    Just x -> return x
  put $ ProgramState vars (heapCnt,heap')
  get
runCom c@(SMatch e cases) = do
  v <- evalExpr e
  (label,value) <- assertVar v
  case lookup label (map (\(a,b,c)-> (a,(b,c)) ) cases) of
    Just (bind,com) -> do
      st@(ProgramState vars heap) <- get
      put $ ProgramState ((bind,value):vars) heap
      runCom com
      (ProgramState vars heap) <- get
      vars' <- ensureTop bind vars
      put $ ProgramState vars' heap
      get
    Nothing -> throwError $ "Unsupported label"++ show label ++ "in case" ++ show c
runCom (FunctionCall var fname arguments) = do
   funcs <- ask
   case lookup fname funcs of
     Nothing -> throwError "no such function defined"
     Just (Func _ args locals body ret) ->do
       (ProgramState vars heap) <- get
       vargs <- mapM evalExpr arguments -- evaluate arguments
       --evaluate local variables in new environment
       put $ ProgramState (zip (map fst args) vargs) heap
       vlocals <- mapM (\(var,expr) -> (var,) <$> evalExpr expr) locals
       --evaluate function body and return in newer environment
       (ProgramState vars' heap') <- get
       put $ ProgramState (vlocals++vars') heap'
       runCom body
       vret <- evalExpr ret
       --restore old local variables with new heap
       (ProgramState _ newheap) <- get
       put (ProgramState vars newheap)
       get
       
       
       
       
       
       

runProgram :: Program -> (Either String ProgramState,ProgramState)
runProgram (Program decls locals com) =
   let dec = map ( \d@(Func n _ _ _ _)-> (n,d) ) decls in
   runState ( runExceptT ( runReaderT (do
            env <- mapM (\(var,expr) -> (var,) <$> evalExpr expr) locals
            put (ProgramState env (0,[]))
            runCom com
        ) dec) )(ProgramState [] (0,[]))
  --runState ( runExceptT  (runReaderT (_) decls)) (ProgramState [] [])