module TypeCheck(typeCheck) where

import Control.Monad
import Data.Maybe
import qualified Data.Map.Strict as Map

import AST
import Defs
import Pretty
import TypeRules


data DBItem = DBVar Type | DBFunc (Maybe Type) [Type]

type TypeDB = [Map.Map Name DBItem]

dbFind :: TypeDB -> Name -> Maybe DBItem
dbFind db name = findJust $ map (Map.lookup name) db
  where findJust [] = Nothing
        findJust (Just x:_) = Just x
        findJust (Nothing:l) = findJust l

dbFindTop :: TypeDB -> Name -> Maybe DBItem
dbFindTop [] _ = error "dbFindTop on empty scope stack" 
dbFindTop (m:_) name = Map.lookup name m

dbSet :: TypeDB -> Name -> DBItem -> TypeDB
dbSet [] _ _ = error "dbSet on empty scope stack"
dbSet (m:ms) name val = Map.insert name val m : ms

emptyDB :: TypeDB
emptyDB = [Map.fromList
    [("putc", DBFunc Nothing [TChar]),
     ("putint", DBFunc Nothing [TInt]),
     ("getc", DBFunc (Just TInt) []),
     ("_builtin_malloc", DBFunc (Just $ TArr TChar Nothing) [TInt])]]

withScope :: TypeDB -> (TypeDB -> a) -> a
withScope db f = f (Map.empty : db)


typeCheck :: Program -> Error Program
typeCheck (Program vars funcs) = do
    db <- foldM registerDVar emptyDB vars
            >>= \db' -> foldM registerDFunc db' funcs
    vars' <- mapM (annotateDVar db) vars
    funcs' <- mapM (annotateDFunc db) funcs
    return $ Program vars' funcs'


registerDVar :: TypeDB -> DVar -> Error TypeDB
registerDVar db (DVar t n _) = case dbFind db n of
    Nothing -> return $ dbSet db n (DBVar t)
    Just _ -> Left $ "Duplicate top-level name '" ++ n ++ "'"

registerDFunc :: TypeDB -> DFunc -> Error TypeDB
registerDFunc db (DFunc rt n al _) = case dbFind db n of
    Nothing -> return $ dbSet db n (DBFunc rt (map fst al))
    Just _ -> Left $ "Duplicate top-level name '" ++ n ++ "'"


annotateDVar :: TypeDB -> DVar -> Error DVar
annotateDVar db (DVar toptype name expr) = do
    expr' <- annotateExpr db expr
    when (isNothing (typeof expr')) $
        Left $ "Cannot assign void value in global declaration of " ++ name
    let typ = fromJust $ typeof expr'
    if canCoerce typ toptype
        then return $ DVar toptype name expr'
        else Left $ "Cannot assign a value of type " ++ pretty typ ++
                    " to a variable of type " ++ pretty toptype

data State = State {stDfunc :: DFunc, stLoopDepth :: Int}
  deriving Show

annotateDFunc :: TypeDB -> DFunc -> Error DFunc
annotateDFunc db dfunc@(DFunc rettype name arglist block) = do
    when (name == "main" && rettype /= Just TInt) $
        Left $ "Function 'main' should return an int"
    let db' = foldl registerArg db arglist
    block' <- annotateBlock (State dfunc 0) db' block
    return $ DFunc rettype name arglist block'
  where
    registerArg :: TypeDB -> (Type, Name) -> TypeDB
    registerArg db' (t, n) = dbSet db' n (DBVar t)

annotateBlock :: State -> TypeDB -> Block -> Error Block
annotateBlock state db (Block sts) =
    Block . snd <$> foldM (\(db', l) st ->
                               (\(db'', st') -> (db'', l ++ [st'])) <$> annotateStatement state db' st)
                          (db, []) sts

annotateStatement :: State -> TypeDB -> Statement -> Error (TypeDB, Statement)
annotateStatement _ db (SDecl toptype name expr) = do
    expr' <- annotateExpr db expr
    when (isNothing (typeof expr')) $
        Left $ "Cannot assign void value in declaration of " ++ name
    when (isJust (dbFindTop db name)) $
        Left $ "Duplicate declaration of variable " ++ name
    let typ = fromJust $ typeof expr'
    if canCoerce typ toptype
        then return $ (dbSet db name (DBVar toptype), SDecl toptype name expr')
        else Left $ "Cannot assign a value of type " ++ pretty typ ++
                    " to a variable of type " ++ pretty toptype
annotateStatement _ db (SAs ae expr) = do
    ae' <- annotateAsExpr db ae
    expr' <- annotateExpr db expr
    when (isNothing (typeof expr')) $
        Left $ "Cannot assign a void value in assignment of " ++ pretty ae
    let typ = fromJust $ typeof expr'
    let aetyp = fromJust $ typeof ae'
    when (not $ isBasicType aetyp) $
        Left $ "Cannot assign to a location of type " ++ pretty aetyp
    when (not $ canCoerce typ aetyp) $
        Left $ "Cannot assign a value of type " ++ pretty typ ++
               " to a location of type " ++ pretty aetyp
    return (db, SAs ae' expr')
annotateStatement st db (SIf expr bl1 bl2) = do
    expr' <- annotateExpr db expr
    when (isNothing (typeof expr')) $
        Left $ "Cannot use void value in 'if' condition"
    when (not $ canCoerce (fromJust (typeof expr')) TInt) $
        Left $ "Cannot use type " ++ pretty (fromJust (typeof expr')) ++ " in 'if' condition"
    bl1' <- withScope db $ flip (annotateBlock st) bl1
    bl2' <- withScope db $ flip (annotateBlock st) bl2
    return (db, SIf expr' bl1' bl2')
annotateStatement (State dfunc ld) db (SWhile expr bl) = do
    expr' <- annotateExpr db expr
    when (isNothing (typeof expr')) $
        Left $ "Cannot use void value in 'while' condition"
    when (not $ canCoerce (fromJust (typeof expr')) TInt) $
        Left $ "Cannot use type " ++ pretty (fromJust (typeof expr')) ++ " in 'while' condition"
    bl' <- withScope db $ flip (annotateBlock (State dfunc (ld+1))) bl
    return (db, SWhile expr' bl')
annotateStatement (State {stLoopDepth = ld}) db (SBreak n) =
    if n > ld
        then Left $ "'break " ++ show n ++ "' while in only " ++ show ld ++ " loops"
        else return (db, SBreak n)
annotateStatement (State {stDfunc = DFunc Nothing _ _ _}) db (SReturn Nothing) =
    return (db, SReturn Nothing)
annotateStatement (State {stDfunc = DFunc (Just _) _ _ _}) _ (SReturn Nothing) =
    Left "Cannot return void value from non-void function"
annotateStatement (State {stDfunc = DFunc mrt _ _ _}) db (SReturn (Just expr)) = do
    expr' <- annotateExpr db expr
    case mrt of
        Nothing -> Left $ "Cannot return non-void value from void function"
        Just rt -> do
            when (isNothing (typeof expr')) $
                Left $ "Cannot use void value in 'return'"
            when (not $ canCoerce (fromJust (typeof expr')) rt) $
                Left $ "Cannot coerce type " ++ pretty (fromJust (typeof expr')) ++ " to " ++ pretty rt ++
                       " in 'return'"
            return (db, SReturn (Just expr'))
annotateStatement _ db (SExpr expr) = (\expr' -> (db, SExpr expr')) <$> annotateExpr db expr

annotateExpr :: TypeDB -> Expression -> Error Expression
annotateExpr db (EBin bo e1 e2 _) = do
    e1' <- annotateExpr db e1
    e2' <- annotateExpr db e2
    when (isNothing (typeof e1')) $ Left $ "Use of void value in expression: " ++ show e1'
    when (isNothing (typeof e2')) $ Left $ "Use of void value in expression: " ++ show e2'
    let t1 = fromJust $ typeof e1'
        t2 = fromJust $ typeof e2'
    rt <- let errval = Left $ "Operator " ++ pretty bo ++ " doesn't take" ++
                              " arguments of types " ++ pretty t1 ++ " and " ++
                              pretty t2
          in maybe errval return $ retTypeBO bo t1 t2
    return $ EBin bo e1' e2' (Just rt)
annotateExpr db (EUn uo e _) = do
    e' <- annotateExpr db e
    when (isNothing (typeof e')) $ Left "Use of void value in expression"
    let t = fromJust $ typeof e'
    rt <- let errval = Left $ "Unary operator " ++ pretty uo ++ " doesn't take" ++
                              " an argument of type " ++ pretty t
          in maybe errval return $ retTypeUO uo t
    return $ EUn uo e' (Just rt)
annotateExpr _ (ELit lit@(LInt _) _) = return $ ELit lit (Just TInt)
annotateExpr _ (ELit lit@(LChar _) _) = return $ ELit lit (Just TChar)
annotateExpr db (ELit lit@(LVar n) _) = case dbFind db n of
    Nothing -> Left $ "Use of undeclared variable " ++ n ++ " in expression"
    Just (DBVar t) -> return $ ELit lit (Just t)
    Just (DBFunc _ _) -> Left $ "Cannot use function " ++ n ++ " in expression"
annotateExpr db (ELit (LCall n as) _) = do
    as' <- mapM (annotateExpr db) as
    case dbFind db n of
        Nothing -> Left $ "Use of undeclared function " ++ n
        Just (DBVar t) -> Left $ "Cannot call variable of type " ++ pretty t
        Just (DBFunc mrt ats) -> do
            when (length as' /= length ats) $
                Left $ "Function '" ++ n ++ "' expected " ++ show (length ats) ++
                       " arguments but got " ++ show (length as')
            forM_ (zip as' ats) $ \(arg, at) -> do
                when (isNothing (typeof arg)) $
                        Left "Use of void value in function argument"
                if canCoerce (fromJust $ typeof arg) at
                    then return ()
                    else Left $ "Argument of " ++ n ++ " has type " ++ pretty at ++
                                " but value of type " ++ pretty (fromJust $ typeof arg) ++
                                " was given"
            return $ ELit (LCall n as') mrt
annotateExpr db (ESubscript arr sub _) = do
    arr' <- annotateExpr db arr
    sub' <- annotateExpr db sub
    let subtyp = fromJust (typeof sub')
    when (subtyp /= TInt) $
        Left $ "Type of array subscript should be int, but is " ++ pretty subtyp
    case fromJust (typeof arr') of
        TArr et _ -> return $ ESubscript arr' sub' (Just et)
        _ -> Left $ "Subscripted expression is not an array: " ++ pretty arr
annotateExpr db (ECast t e) = do
    e' <- annotateExpr db e
    let typ = fromJust (typeof e')
    if canCast typ t
        then return $ ECast t e'
        else Left $ "Cannot cast value of type " ++ pretty typ ++ " to type " ++ pretty t
annotateExpr db e@(ENew t sze) = do
    sze' <- annotateExpr db sze
    case typeof sze' of
        Nothing -> Left $ "Use of void value in array length in 'new' expression: " ++ pretty e
        Just TInt -> return ()
        Just szet -> Left $ "Type of array length in 'new' expression should be int, is " ++ pretty szet
    if isBasicType t
        then return $ ENew t sze'
        else Left $ "Can only allocate arrays of basic types using 'new': " ++ pretty e

annotateAsExpr :: TypeDB -> AsExpression -> Error AsExpression
annotateAsExpr db (AEVar n _) = case dbFind db n of
    Nothing -> Left $ "Use of undeclared variable " ++ n ++ " in assignment expression"
    Just (DBVar t) -> return $ AEVar n (Just t)
    Just (DBFunc _ _) -> Left $ "Cannot use function " ++ n ++ " in assignment expression"
annotateAsExpr db (AESubscript ae expr _) = do
    ae' <- annotateAsExpr db ae
    expr' <- annotateExpr db expr
    case typeof expr' of
        Nothing -> Left $ "Use of void value in array index in assignment expression"
        Just TInt -> return ()
        Just t -> Left $ "Use of non-int type " ++ pretty t ++ " in array index in assignment expression"
    case fromJust (typeof ae') of
        TArr t _ -> return $ AESubscript ae' expr' (Just t)
        t -> Left $ "Indexed expression '" ++ pretty ae ++ "' has non-array type " ++ pretty t ++
                    " in assignment expression"