module Utility = Utility
open Types;;

Random.self_init ()

let (let*) = Result.bind

let rec evaluate_type (program: t_exp) (context: ftype VariableMap.t) : (ftype, [> typechecking_error]) result =
  match program with
    Integer _ -> Ok IntegerType
  | Boolean _ -> Ok BooleanType
  | Variable x -> ( (* check for the type in the context *)
      match VariableMap.find_opt x context with
        None -> Error (`AbsentAssignment
                         ("The variable " ^ x ^ " is not defined."))
      | Some t -> Ok t
    )
  | Tuple (x, y) -> (
      let* xtype = evaluate_type x context in
      let* ytype = evaluate_type y context in
      Ok (TupleType (xtype, ytype))
    )
  | Function (x, typef, fbody) -> (
      (* first check that the function has the right specified type then check
         the type of the body using the bindings for the input *)
      match typef with
        FunctionType (tin, tout) -> (
          let* typefbody = evaluate_type fbody (VariableMap.add x tin context) in
          if (typefbody = tout) then
            Ok typef
          else
            Error (`WrongTypeSpecification
                     ("Function does not return specified type."))
        )
      | _ -> Error (`WrongTypeSpecification
                      ("Specification of function is not a function type."))
    )
  | Application (f, x) -> (
      let* evalf = evaluate_type f context in
      let* evalx = evaluate_type x context in
      match evalf with
        FunctionType (tin, tout) -> (
          if tin = evalx then
            Ok tout
          else
            Error (`WrongType "Appling function with wrong input type to value")
        )
      | _ -> Error (`WrongType "Applying to a non function type")
    )
  | Plus (x, y)
  | Minus (x, y)
  | Times (x, y)
  | Division (x, y)
  | Modulo (x, y)
  | Power (x, y) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      match typex, typey with
      | (IntegerType, IntegerType) -> Ok IntegerType
      | (IntegerType, _) -> Error (`WrongType "Second term is not an integer.")
      | (_, _) -> Error (`WrongType "First term is not an integer.")
    )
  | PowerMod (x, y, z) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      let* typez = evaluate_type z context in
      match typex, typey, typez with
      | (IntegerType, IntegerType, IntegerType) -> Ok IntegerType
      | (IntegerType, IntegerType, _) -> Error (`WrongType ("Third term is " ^
                                                            "not an integer."))
      | (IntegerType, _, _) -> Error (`WrongType
                                        ("Second term is not an integer."))
      | (_, _, _) -> Error (`WrongType "First term is not an integer.")
    )
  | Rand (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (IntegerType) -> Ok IntegerType
      | (_) -> Error (`WrongType "Term is not an integer.")
    )
  | BAnd (x, y)
  | BOr (x, y) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      match typex, typey with
      | (BooleanType, BooleanType) -> Ok BooleanType
      | (BooleanType, _) -> Error (`WrongType "Second term is not a boolean.")
      | (_, _) -> Error (`WrongType "First term is not a boolean.")
    )
  | BNot (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (BooleanType) -> Ok BooleanType
      | (_) -> Error (`WrongType "Term is not a boolean.")
    )
  | First (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (TupleType (x, _)) -> Ok x
      | (_) -> Error (`WrongType "Term is not a tuple.")
    )
  | Second (x) -> (
      let* typex = evaluate_type x context in
      match typex with
      | (TupleType (_, x)) -> Ok x
      | (_) -> Error (`WrongType "Term is not a tuple.")
    )
  | Cmp (x, y)
  | CmpLess (x, y)
  | CmpLessEq (x, y)
  | CmpGreater (x, y)
  | CmpGreaterEq (x, y) -> (
      let* typex = evaluate_type x context in
      let* typey = evaluate_type y context in
      match typex, typey with
      | (IntegerType, IntegerType) -> Ok BooleanType
      | (IntegerType, _) -> Error (`WrongType "Second term is not an integer.")
      | (_, _) -> Error (`WrongType "First term is not an integer.")
    )
  | IfThenElse (guard, if_exp, else_exp) -> (
      let* typeguard = evaluate_type guard context in
      let* typeif_exp = evaluate_type if_exp context in
      let* typeelse_exp = evaluate_type else_exp context in
      match typeguard, typeif_exp, typeelse_exp with
        (BooleanType, t1, t2) -> (
          if t1 = t2 then
            Ok t1
          else
            Error (`WrongType "If branches do not have the same type.")
        )
      | (_, _, _) -> Error (`WrongType "If guard is not a boolean.")
    )
  | LetIn (x, xval, rest) ->
    (* bind the type to the variable name in the context *)
    let* typex = evaluate_type xval context in
    evaluate_type rest (VariableMap.add x typex context)
  | LetFun (f, x, typef, fbody, rest) ->
    (* like with the function case, but also add f itself to the bindings *)
    match typef with
      FunctionType (tin, tout) -> (
        let newcontext = VariableMap.add f typef context in
        let newcontextwithx = VariableMap.add x tin newcontext in
        let* typefbody = evaluate_type fbody newcontextwithx in
        let* typerest = evaluate_type rest newcontext in
        match (typefbody = tout, typerest) with
          (false, _) -> Error (`WrongTypeSpecification
                                 "Function does not return specified type.")
        | (true, t) -> Ok t
      )
    | _ -> Error (`WrongTypeSpecification
                    "Specification of function is not a function type.")

let typecheck (program: t_exp) : (ftype, [> typechecking_error]) result =
  let* typeprogram = evaluate_type program VariableMap.empty in
  match typeprogram with
    FunctionType (IntegerType, IntegerType) -> (
      Ok (typeprogram)
    )
  | _ -> Error (`WrongType "Program is not a function from int to int.")