module TypeChecker
open Syntax
open Util

let getFunParmamType funName env =
    let funType = Map.find funName env
    match funType with
    | FuncType(left, right) -> left
    | BoolType              -> BoolType
    | IntType               -> IntType

let addParamType env param =
    let (Parameter(name, paramType)) = param
    Map.add name paramType env

let addFuncType env funcDef =
    match funcDef with
    | FuncDef (name, _, typeDef, _) -> Map.add name typeDef env

let rec checkExpr expr env =
    match expr with
    | ConstBool b                       -> BoolType
    | ConstInt i                        -> IntType
    | Var name                          -> Map.find name env
    | PrimOp ("+", lOper, rOper)        -> checkIntOp lOper rOper env
    | PrimOp ("-", lOper, rOper)        -> checkIntOp lOper rOper env
    | PrimOp ("*", lOper, rOper)        -> checkIntOp lOper rOper env
    | PrimOp ("/", lOper, rOper)        -> checkIntOp lOper rOper env
    | PrimOp (">", lOper, rOper)        -> checkIntOp lOper rOper env
    | PrimOp ("<", lOper, rOper)        -> checkIntOp lOper rOper env
    | PrimOp ("=", lOper, rOper)        -> checkBinOp lOper rOper env
    | If (ifExpr, thenExpr, elseExpr )  -> checkIf ifExpr thenExpr elseExpr env
    | Let(name, valExpr, restExpr)      -> checkLet name valExpr restExpr env
    | FunCall (name, arg)               -> checkFunCall name arg env
    | _                                 -> failwith "Syntax error"

and checkLet name valExpr restExpr env =
    let valType = checkExpr valExpr env
    let updatedEnv = Map.add name valType env
    checkExpr restExpr updatedEnv

and checkIf ifExpr thenExpr elseExpr env =
    let ifType = checkExpr ifExpr env
    let thenType = checkExpr thenExpr env
    let elseType = checkExpr elseExpr env
    if ifType = BoolType || ifType = IntType && thenType = elseType then thenType else failwith "Error"

and checkIntOp lOper rOper env =
    let lType = checkExpr lOper env
    let rType = checkExpr rOper env
    if lType = rType && rType = IntType then lType else failwith "Error"

and checkBinOp lOper rOper env =
    let lType = checkExpr lOper env
    let rType = checkExpr rOper env
    if (lType = IntType || lType = BoolType) && lType = rType then rType else failwith "Error"

and checkFunCall name arg env =
    let argType = checkExpr arg env
    match name with
    | (Var funName) -> if (getFunParmamType funName env) = argType then argType else failwith "Illegal param type"
    | _             -> failwith "Syntax error: expected (Var name)"

let checkFunc funcDef env =
    let (FuncDef(name, paramList, typeDef, expr)) = funcDef
    let paramTypes = getParamTypes paramList typeDef
    let updatedEnv = List.fold addParamType env paramTypes
    let checkType = checkExpr expr updatedEnv
    checkType = IntType || checkType = BoolType

let checkProgram program =
    match program with
    | Program (funcDefs, expr) -> let env = List.fold addFuncType Map.empty funcDefs
                                  let funcsValid = List.forall (fun funcDef -> checkFunc funcDef env) funcDefs
                                  let mainType = checkExpr expr env
                                  let mainValid = mainType = IntType || mainType = BoolType
                                  mainValid && funcsValid