module TypeChecker
open Syntax

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

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 && thenExpr = elseExpr 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 paramType = checkExpr arg env
    match name with
    | (Var funName) -> if (getFunArgType funName env) = paramType then paramType else failwith "Illegal parm type"
    | _             -> failwith "Syntax error: expected (Var name)"

let checkFunDef name typeDef expr env =
    let updatedEnv = Map.add name typeDef env
    checkExpr expr updatedEnv

let checkProgram program =
    let env =