{-# LANGUAGE TupleSections #-}

module TypeCheck(typeCheck) where

import Control.Monad
import Data.List
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] | DBType 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) []),
     ("exit", DBFunc Nothing [TInt]),
     ("_builtin_malloc", DBFunc (Just $ TArr TChar Nothing) [TInt]),
     ("char", DBType TChar),
     ("int", DBType TInt)]]

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


typeCheck :: Program -> Error Program
typeCheck (Program tdefs vars funcs) = do
    -- case topologicalSort 
    db <- foldM registerDTypedef emptyDB tdefs
            >>= (\db' -> foldM registerDTypedefResolve db' tdefs)
            >>= (\db' -> foldM registerDVar db' vars)
            >>= (\db' -> foldM registerDFunc db' funcs)
    vars' <- mapM (annotateDVar db) vars
    funcs' <- mapM (annotateDFunc db) funcs
    return $ Program [] vars' funcs'


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

registerDTypedefResolve :: TypeDB -> DTypedef -> Error TypeDB
registerDTypedefResolve db (DTypedef n t) = do
    t' <- resolveType db t
    return $ dbSet db n (DBType t')

registerDVar :: TypeDB -> DVar -> Error TypeDB
registerDVar db (DVar t n _) = case dbFind db n of
    Nothing -> do
        t' <- resolveType db t
        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 -> do
        rt' <- sequence $ fmap (resolveType db) rt
        al' <- forM al $ \(at,an) -> (,an) <$> resolveType db at
        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
    toptype' <- resolveType db toptype
    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
    rettype' <- sequence $ fmap (resolveType db) rettype
    arglist' <- forM arglist $ \(at,an) -> (,an) <$> resolveType db at
    when (name == "main" && rettype' /= Just TInt) $
        Left $ "Function 'main' should return an int"
    db' <- foldM registerArg db arglist'
    block' <- annotateBlock (State dfunc 0) db' block
    return $ DFunc rettype' name arglist' block'
  where
    registerArg :: TypeDB -> (Type, Name) -> Error TypeDB
    registerArg db' (t, n) = dbSet db' n . DBVar <$> resolveType db' 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
    toptype' <- resolveType db toptype
    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
annotateStatement _ db SDebugger = return (db, SDebugger)

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"
    Just (DBType _) -> Left $ "Cannot use type " ++ n ++ " as value 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_ (zip3 as' ats [1 :: Int ..]) $ \(arg, at, num) -> do
                when (isNothing (typeof arg)) $
                        Left "Use of void value in function argument"
                if canCoerce (fromJust $ typeof arg) at
                    then return ()
                    else Left $ "Argument " ++ show num ++ " of " ++ n ++ " has type " ++ pretty at ++
                                " but value of type " ++ pretty (fromJust $ typeof arg) ++
                                " was given"
            return $ ELit (LCall n as') mrt
        Just (DBType t) -> case as of
            [a] -> annotateExpr db (ECast t a)
            _ -> Left $ "Cannot call type " ++ pretty t ++ " as function with " ++
                        show (length as) ++ " arguments"
annotateExpr _ (ELit lit@(LStr s) _) =
    return $ ELit lit (Just $ TArr TChar (Just $ fromIntegral $ length s))
annotateExpr db (ELit (LStruct ms) _) = do
    ms' <- forM ms $ \(n,e) -> (n,) <$> annotateExpr db e
    types <- forM ms' $ \(n,e) -> case typeof e of
                Nothing -> Left $ "Use of void value in struct literal item '" ++ n ++ "'"
                Just t -> return t
    return $ ELit (LStruct ms') (Just $ TStruct $ zip types (map fst ms'))
annotateExpr db (ESubscript arr sub _) = do
    arr' <- annotateExpr db arr
    sub' <- annotateExpr db sub
    when (isNothing (typeof sub')) $
        Left $ "Use of void value as subscripted expression"
    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 (EGet st n _) = do
    st' <- annotateExpr db st
    case typeof st' of
        Nothing -> Left $ "Use of void value as dot-indexed expression"
        Just stt@(TStruct ms) -> case find ((==n) . snd) ms of
            Nothing -> Left $ "Struct of type " ++ pretty stt ++
                              " has no member named '" ++ n ++ "'"
            Just (t, _) -> return $ EGet st' n (Just t)
        Just t -> Left $ "Use of non-struct type " ++ pretty t ++ " as dot-indexed expression"
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"
    Just (DBType _) -> Left $ "Cannot use type " ++ n ++ " as variable 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"
annotateAsExpr db (AEGet ae n _) = do
    ae' <- annotateAsExpr db ae
    case typeof ae' of
        Nothing -> Left $ "Use of void value in dot-indexed assignment expression"
        Just stt@(TStruct ms) -> case find ((==n) . snd) ms of
            Nothing -> Left $ "Struct of type " ++ pretty stt ++ " has no member named '" ++ n ++
                              "' in assignment expression"
            Just (t, _) -> return $ AEGet ae' n (Just t)
        Just t -> Left $ "Use of non-struct type " ++ pretty t ++ " as dot-indexed assignment expression"


resolveType :: TypeDB -> Type -> Error Type
resolveType _ TInt = return TInt
resolveType _ TChar = return TChar
resolveType db (TArr t sz) = liftM (\t' -> TArr t' sz) $ resolveType db t
resolveType db (TStruct ms) = TStruct <$> mapM (\(t,n) -> liftM (,n) $ resolveType db t) ms
resolveType db (TName n) = case dbFind db n of
    Nothing -> Left $ "Type name '" ++ n ++ "' not defined"
    Just (DBType t) -> return t
    Just _ -> Left $ "Name '" ++ n ++ "' used as type is not a type"