module Interpreter where

import Control.Monad
import Control.Exception
import System.Exit

import Data.Map (Map)
import qualified Data.Map.Strict as Map
import Data.List
import Data.Maybe

import CMM.Abs
import CMM.Print
import CMM.ErrM

data Val = D Double | I Integer | B Bool | Void deriving (Eq)
instance Show Val where
  show (D n) = show n
  show (I n) = show n
  show (B n) = show n
  show Void = "void"

type Env = (Fns, [Context]) -- functions and context stack
type Fns = Map Id Def
type Context = Map Id (Maybe Val) -- variables with their values

printEnv :: Env -> String
printEnv (_, ctx) = show ctx

emptyEnv :: Env
emptyEnv = (Map.empty, [Map.empty])

extendFn :: Env -> Def -> Env
extendFn env (DFun t ident args stms) = (Map.insert ident (DFun t ident args stms) (fst env), snd env)

lookupFn :: Env -> Id -> Maybe Def
lookupFn (fns, _) ident = Map.lookup ident fns

newBlock :: Env -> Env
newBlock env = (fst env, Map.empty:snd env)

exitBlock :: Env -> Env
exitBlock env = (fst env, tail $ snd env)

clean :: Env -> Env
clean (fns, _) = (fns, [Map.empty])

extendVar :: Env -> Id -> Maybe Val -> Env
extendVar (fns, ctxs) i v = (fns, Map.insert i v (head ctxs):tail ctxs)

modifyVar :: Env -> Id -> Maybe Val -> Env
modifyVar (fns, ctxs) i v = case findIndex (Map.member i) ctxs of
  Nothing -> (fns, ctxs) -- this should never happen due to type checker
  Just idx -> (fns, modifyNth ctxs idx (Map.insert i v))

modifyNth :: [a] -> Int -> (a -> a) -> [a]
modifyNth (x:xs) 0 fn = (fn x):xs
modifyNth (x:xs) n fn = x:(modifyNth xs (n-1) fn)

lookupVar :: Env -> Id -> Maybe Val
lookupVar e i = case find (Map.member i) (snd e) of
  Nothing -> Nothing
  Just ctx -> fromJust $ Map.lookup i ctx

interpret :: Program -> IO ()
interpret (PDefs defs) = do
  -- print defs -- uncomment to print AST
  let env = foldl extendFn emptyEnv defs
  case lookupFn env (Id "main") of
    Nothing -> interpreterErr env "no main function"
    Just (DFun _ _ _ stms) ->  do
      evalStms env stms
      return ()

evalStms :: Env -> [Stm] -> IO (Val, Env)
evalStms env [] = return (Void, env)
evalStms env (s:ss) = do
  -- putStrLn $ printTree s
  case s of
    SDecl _ iden -> evalStms (extendVar env iden Nothing) ss
    SDecls _ iden idens -> do
      let env' = foldl (\e i -> extendVar e i Nothing) env (iden:idens)
      evalStms env' ss
    SInit _ ident exp -> do
      (v, env') <- evalExp env exp
      evalStms (extendVar env' ident (Just v)) ss
    SReturn exp -> evalExp env exp
    SBlock stms -> do
      (v, env') <- evalStms (newBlock env) stms
      if v == Void
        then evalStms (exitBlock env') ss
        else return (v, exitBlock env')
    SWhile exp stm -> do
      (v, env') <- evalExp env exp
      case v of
        B False -> evalStms env' ss
        B True -> do
          (v, env'') <- evalStms (newBlock env') [stm]
          if v == Void
            then evalStms (exitBlock env'') (s:ss)
            else return (v, env'')
    SIfElse exp stm1 stm2 -> do
      (v, env') <- evalExp env exp
      case v of
        B True -> do
          (v', env'') <- evalStms (newBlock env') [stm1]
          if v' == Void
            then evalStms (exitBlock env'') ss
            else return (v', env'')
        B False -> do
          (v', env'') <- evalStms (newBlock env') [stm2]
          if v' == Void
            then evalStms (exitBlock env'') ss
            else return (v', env'')
    SExp exp -> do
      (v, env') <- evalExp env exp
      evalStms env' ss

evalExp :: Env -> Exp -> IO (Val, Env)
evalExp env exp = do
  -- putStrLn $ show exp
  case exp of
    EInt n -> return (I n, env)
    EDouble n -> return (D n, env)
    ETrue -> return (B True, env)
    EFalse -> return (B False, env)
    EId ident -> case lookupVar env ident of
      Nothing -> interpreterErr env $ "uninitialized variable " ++ show exp
      Just v -> do
        -- putStrLn $ show v
        return (v, env)
    ECall (Id "printInt") [exp] -> do
      (v, env') <- evalExp env exp
      print $ v
      return (Void, env')
    ECall (Id "printDouble") [exp] -> do
      (v, env') <- evalExp env exp
      print $ v
      return (Void, env')
    ECall (Id "readInt") [] -> do
      line <- getLine
      let v = (read :: String -> Integer) line
      return (I v, env)
    ECall (Id "readDouble") [] -> do
      line <- getLine
      let v = (read :: String -> Double) line
      return (D v, env)
    ECall iden exps -> do
      (vs, env') <- evalExps env exps
      case lookupFn env iden of
        Nothing -> interpreterErr env $ "uninitialized func " ++ printTree iden
        Just (DFun _ _ args stms) -> do
          let env'' = foldl bindArg (clean env) (zip vs args)
          (v, _) <- evalStms env'' stms
          return (v, env')
    EPIncr ident ->
      case lookupVar env ident of
        Nothing -> interpreterErr env $ "uninitialized variable " ++ printTree ident
        Just v -> do
          let env2 = modifyVar env ident (Just (inc v))
          return (v, env2)
    EPDecr ident ->
      case lookupVar env ident of
        Nothing -> interpreterErr env $ "uninitialized variable " ++ printTree ident
        Just v -> do
          let env2 = modifyVar env ident (Just (dec v))
          return (v, env2)
    EIncr ident ->
      case lookupVar env ident of
        Nothing -> interpreterErr env $ "uninitialized variable " ++ printTree ident
        Just v -> do
          let v2 = inc v
          let env2 = modifyVar env ident (Just v2)
          return (v2, env2)
    EDecr ident ->
      case lookupVar env ident of
        Nothing -> interpreterErr env $ "uninitialized variable " ++ printTree ident
        Just v -> do
          let v2 = dec v
          let env2 = modifyVar env ident (Just v2)
          return (v2, env2)
    EMul exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (mul v1 v2, env'')
    EDiv exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (divide v1 v2, env'')
    EAdd exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (add v1 v2, env'')
    ESub exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (sub v1 v2, env'')
    ELt exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (less v1 v2, env'')
    ELEq exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (leq v1 v2, env'')
    EGt exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (greater v1 v2, env'')
    EGEq exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (geq v1 v2, env'')
    EEq exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (B (v1 == v2), env'')
    ENEq exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      (v2, env'') <- evalExp env' exp2
      return (B (not (v1 == v2)), env'')
    EAnd exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      if v1 == B False
        then return (v1, env')
        else evalExp env' exp2
    EOr exp1 exp2 -> do
      (v1, env') <- evalExp env exp1
      if v1 == B True
        then return (v1, env')
        else evalExp env' exp2
    EAss iden exp -> do
      (v, env') <- evalExp env exp
      return (v, modifyVar env' iden (Just v))

bindArg :: Env -> (Val, Arg) -> Env
bindArg env (v, (ADecl _ iden)) = extendVar env iden (Just v)

evalExps :: Env -> [Exp] -> IO ([Val], Env)
evalExps env [] = return ([], env)
evalExps env (e:es) = do
  (v, env') <- evalExp env e
  (vs, env'') <- evalExps env' es
  return (v:vs, env'')

interpreterErr :: Env -> String -> IO a
interpreterErr env s = do
  putStrLn "INTERPRETER ERROR"
  putStrLn $ printEnv env
  putStrLn s
  exitFailure

conj :: Val -> Val -> Val
conj (B n1) (B n2) = B (n1 && n2)

disj :: Val -> Val -> Val
disj (B n1) (B n2) = B (n1 || n2)

add :: Val -> Val -> Val
add (I n1) (I n2) = I (n1 + n2)
add (D n1) (D n2) = D (n1 + n2)
add (D n1) (I n2) = D (n1 + (fromInteger n2))
add (I n1) (D n2) = D ((fromInteger n1) + n2)

mul :: Val -> Val -> Val
mul (I n1) (I n2) = I (n1 * n2)
mul (D n1) (D n2) = D (n1 * n2)
mul (D n1) (I n2) = D (n1 * (fromInteger n2))
mul (I n1) (D n2) = D ((fromInteger n1) * n2)

divide :: Val -> Val -> Val
divide (I n1) (I n2) = I (div n1 n2)
divide (D n1) (D n2) = D (n1 / n2)
divide (D n1) (I n2) = D (n1 / (fromInteger n2))
divide (I n1) (D n2) = D ((fromInteger n1) / n2)

sub :: Val -> Val -> Val
sub (I n1) (I n2) = I (n1 - n2)
sub (D n1) (D n2) = D (n1 - n2)
sub (D n1) (I n2) = D (n1 - (fromInteger n2))
sub (I n1) (D n2) = D ((fromInteger n1) - n2)

inc :: Val -> Val
inc (I n) = I (n+1)
inc (D n) = D (n+1)

dec :: Val -> Val
dec (I n) = I (n-1)
dec (D n) = D (n-1)

less :: Val -> Val -> Val
less (I n1) (I n2) = B (n1 < n2)
less (D n1) (D n2) = B (n1 < n2)
less (I n1) (D n2) = B ((fromInteger n1) < n2)
less (D n1) (I n2) = B (n1 < (fromInteger n2))

leq :: Val -> Val -> Val
leq (I n1) (I n2) = B (n1 <= n2)
leq (D n1) (D n2) = B (n1 <= n2)
leq (I n1) (D n2) = B ((fromInteger n1) <= n2)
leq (D n1) (I n2) = B (n1 <= (fromInteger n2))

geq :: Val -> Val -> Val
geq (I n1) (I n2) = B (n1 >= n2)
geq (D n1) (D n2) = B (n1 >= n2)
geq (I n1) (D n2) = B ((fromInteger n1) >= n2)
geq (D n1) (I n2) = B (n1 >= (fromInteger n2))

greater :: Val -> Val -> Val
greater (I n1) (I n2) = B (n1 > n2)
greater (D n1) (D n2) = B (n1 > n2)
greater (I n1) (D n2) = B ((fromInteger n1) > n2)
greater (D n1) (I n2) = B (n1 > (fromInteger n2))