{-# LANGUAGE TupleSections #-}
module Typecheck
(inferExpr,
 checkExpr,
 wellFormed,
 Environment,
 )
  where
import Syntax
--import Control.Monad
import Control.Monad.Except
type Environment = [(Idnt,Typ)]

type Error = Except String

assertInt :: Typ -> Expr -> Error ()
assertInt TInt _ = return ()
assertInt t e = throwError $ "Expected type of expression "++ show e ++ " to be (TInt), instead got" ++ show t 
  
assertRec :: Typ -> Expr -> Error TypeDict
assertRec (TRecord ts) _ = return ts
assertRec t e = throwError $ "Expected type of expression "++ show e ++ " to be (TRecord _), instead got" ++ show t 

assertVar :: Typ -> Expr -> Error TypeDict
assertVar (TVariant ts) _ = return ts
assertVar t e = throwError $ "Expected type of expression "++ show e ++ " to be (TVariant _), instead got" ++ show t 

assertPtr :: Typ -> Expr -> Error Typ
assertPtr (TPtr t) _ = return t
assertPtr t e = throwError $ "Expected type of expression "++ show e ++ " to be (TPtr _), instead got" ++ show t 

envlookup :: Idnt -> TypeDict -> Expr -> Error Typ
envlookup name env expr = case lookup name env of
  Just t -> return t
  Nothing -> throwError $ "Variable "++ name ++ "not defined in expr" ++ show expr

insert :: Idnt -> Typ -> Environment -> TypeDict
insert v t l = (v,t):l

subset ::Show a=> [Idnt] -> [Idnt]-> a -> Error ()
subset [] _ _= return ()
subset (x:xs) l a= if x `elem` l then subset xs l a else throwError $ "Case "++ show x ++ "not checked in case/match" ++ show a

inferExpr :: Expr -> Environment -> Error Typ
inferExpr e@(Var name) env = envlookup name env e
inferExpr (IntLit _) _ = return TInt
inferExpr (Record ts) env = do
  tts <- mapM (\(v,e) -> (v,) <$> inferExpr e env) ts
  return $ TRecord tts
inferExpr ee@(Variant t v e) env = do
  ts <- assertVar t e
  tt <- envlookup v ts ee
  checkExpr e env tt
  return t
inferExpr (Add e1 e2) env = do
  checkExpr e1 env TInt
  checkExpr e2 env TInt
  return TInt
inferExpr (Mult e1 e2) env = do
  checkExpr e1 env TInt
  checkExpr e2 env TInt
  return TInt
inferExpr (Neg e) env = do
  checkExpr e env TInt
  return TInt
inferExpr (Deref e1) env = do
  t <- inferExpr e1 env
  tt <- assertPtr t e1
  return tt
inferExpr c@(Match e cases) env = do
  varType <- inferExpr e env
  variants <- assertVar varType e
  caseTypes <- mapM (\(v,bind,e) ->(v,) <$> do
                        tt <- envlookup v variants e
                        inferExpr e (insert bind tt env)) cases
  subset (map fst variants) (map (\(_,v,_)->v) cases) c
  ret <- allSame caseTypes
  return $ snd ret
  where
    allSame :: [(Idnt,Typ)] -> Error (Idnt,Typ)
    allSame (x:[]) = return x
    allSame [] = throwError $ "Empty case "++ show c
    allSame ((xx,x):xs) = do
      (v,t) <- allSame xs
      if t==x then return (v,t) else
        throwError $ "Different types in branches " ++ show v ++ "," ++ show xx ++ "of expression" ++ show c
inferExpr (Proj name e) env = do
  t <- inferExpr e env
  tt <- assertRec t e
  case lookup name tt of
    Just typ -> return typ
    Nothing -> throwError $ "Record " ++show e ++ "doesn't have field" ++ name

checkExpr :: Expr -> Environment -> Typ -> Error ()
checkExpr e@(Add e1 e2) env t = do
  assertInt t e
  checkExpr e1 env TInt
  checkExpr e2 env TInt
checkExpr e@(Mult e1 e2) env t = do
  assertInt t e
  checkExpr e1 env TInt
  checkExpr e2 env TInt
checkExpr e@(Neg ee) env t = do
  assertInt t e
  checkExpr ee env TInt
checkExpr (Deref e1) env t = do
 checkExpr e1 env (TPtr t)
checkExpr e env t = do
  t' <- inferExpr e env
  if t == t' then return ()
    else throwError $ "Expected type of expression " ++
         show e ++ " to be"++ show t  ++" instead got" ++ show t'

bool = TVariant [("True",TRecord []),("False", TRecord [])]
  
cenvlookup :: Idnt -> TypeDict -> Com -> Error Typ
cenvlookup name env c = case lookup name env of
  Just t -> return t
  Nothing -> throwError $ "Variable "++ name ++ "not defined in expr" ++ show c

wellFormed :: Com ->Environment -> Error ()
wellFormed Skip _= return ()
wellFormed (Seq s1 s2) env = do
  wellFormed s1 env
  wellFormed s2 env
wellFormed (If e c1 c2) env = do
  checkExpr e env bool
  wellFormed c1 env
  wellFormed c2 env
wellFormed (While e c) env = do
  checkExpr e env bool
  wellFormed c env
wellFormed c@(Asgn v e) env = do
  vt <- cenvlookup v env c
  et <- inferExpr e env
  if vt==et then return ()
    else throwError $ "Variable " ++ show v ++ "::" ++ show vt ++ " was assigned expression " ++ show e ++ "::" ++ show et
wellFormed (Decl v e c) env= do
  et <- inferExpr e env
  wellFormed c (insert v et env)
wellFormed (Alloc v e c) env = do
  et <- inferExpr e env
  wellFormed c (insert v (TPtr et) env)
wellFormed c@(Save v e) env = do
  vt <- cenvlookup v env c
  et <- inferExpr e env
  if vt==TPtr et then return ()
    else throwError $ "Expression" ++ show e ++ "::" ++ show et ++ " was saved to pointer "++show v ++ " of type" ++ show vt
wellFormed c@(SMatch e cases) env = do
  varType <- inferExpr e env
  variants <- assertVar varType e
  mapM (\(v,bind,c) -> do
                        tt <- envlookup v variants e
                        wellFormed c (insert bind tt env)) cases
  subset (map fst variants) (map (\(_,v,_)->v) cases) c
  return ()