module TypeCheck where

import Expr
import Control.Monad

typeOf :: Value -> Type
typeOf (BoolV {}) = BoolT
typeOf (IntV {}) = IntT
typeOf (ListV (v : _)) = ListT $ typeOf v
typeOf (ListV []) = ListT AnyT
typeOf (TreeV (TNode { treeLeft, treeRoot, treeRight })) = TreeT $ typeOf treeRoot
typeOf (TreeV (TLeaf v)) = TreeT $ typeOf v

isBool = (== BoolT) . typeOf
isInt = (== IntT) . typeOf
isList x | ListT {} <- typeOf x = True
         | otherwise = False
isTree x | TreeT {} <- typeOf x = True
         | otherwise = False

simpleUnify :: Type -> Type -> Maybe Type
simpleUnify AnyT u = Just u
simpleUnify t AnyT = Just t
simpleUnify _ _ = Nothing

checkType :: TypeConf -> Expr -> Maybe Type
checkType conf (left :&&: right) = do BoolT <- checkType conf left
                                      BoolT <- checkType conf right
                                      return BoolT
checkType conf (left :||: right) = do BoolT <- checkType conf left
                                      BoolT <- checkType conf right
                                      return BoolT
checkType conf (NotE e) = do BoolT <- checkType conf e
                             return BoolT
checkType conf (left :=: right) = do leftT <- checkType conf left
                                     rightT <- checkType conf right
                                     guard $ leftT == rightT
                                     return BoolT
checkType conf (Leq0 e) = do IntT <- checkType conf e
                             return BoolT
checkType conf (IsEmptyE e) = do ListT _ <- checkType conf e
                                 return BoolT
checkType conf (left :+: right) = do IntT <- checkType conf left
                                     IntT <- checkType conf right
                                     return IntT
checkType conf (left :-: right) = do IntT <- checkType conf left
                                     IntT <- checkType conf right
                                     return IntT
checkType conf (IncE e) = do IntT <- checkType conf e
                             return IntT
checkType conf (DecE e) = do IntT <- checkType conf e
                             return IntT
checkType conf ZeroE = return IntT
checkType conf (Div2E e) = do IntT <- checkType conf e
                              return IntT
checkType conf (TailE e) = do ListT t <- checkType conf e
                              guard $ t /= AnyT -- NOTE: can't take tail from list that is provem empty (?)
                                                -- TODO: check that AnyT is only for empty list & no bad consequences
                              return $ ListT t
checkType conf (HeadE e) = do ListT t <- checkType conf e
                              guard $ t /= AnyT -- NOTE: can't take elem from list that is provem empty (?)
                                                -- TODO: check that AnyT is only for empty list & no bad consequences
                              return t
checkType conf (left :++: right) = do ListT t <- checkType conf left
                                      ListT u <- checkType conf right
                                      w <- simpleUnify t u
                                      return $ ListT w
checkType conf (left ::: right) = do t <- checkType conf left
                                     ListT u <- checkType conf right
                                     w <- simpleUnify t u
                                     return $ ListT w
checkType conf EmptyListE = return $ ListT AnyT -- NOTE: ListT AnyT - type of generic empty list
checkType conf (IsLeafE e) = do TreeT _ <- checkType conf e
                                return BoolT
checkType conf (TreeValE e) = do TreeT t <- checkType conf e
                                 return t
checkType conf (TreeLeftE e) = do TreeT t <- checkType conf e
                                  return $ TreeT t
checkType conf (TreeRightE e) = do TreeT t <- checkType conf e
                                   return $ TreeT t
checkType conf (CreateNodeE {nodeLeft, nodeRoot, nodeRight}) = do TreeT t <- checkType conf nodeLeft
                                                                  u <- checkType conf nodeRoot
                                                                  guard $ t == u
                                                                  TreeT w <- checkType conf nodeRight
                                                                  guard $ t == w
                                                                  return $ TreeT t
checkType conf (CreateLeafE e) = do t <- checkType conf e
                                    return $ TreeT t
checkType conf (IfE {ifCond, ifDoThen, ifDoElse}) = do BoolT <- checkType conf ifCond
                                                       leftT <- checkType conf ifDoThen
                                                       rightT <- checkType conf ifDoElse
                                                       guard $ leftT == rightT
                                                       return leftT
checkType conf (SelfE es) = do let ts = typeConfInput conf
                               guard $ length ts == length es
                               guard $ and $ zipWith (\t e -> checkType conf e == Just t) ts es
                               return $ typeConfOutput conf
checkType conf (InputE i) = Just $ typeConfInput conf !! i
checkType _ Hole = Nothing
-- checkType _ _ = Nothing