module Utility = Utility
open Types;;

Random.self_init ()

let (let*) = Result.bind

let rec evaluate_type (program: t_exp) (context: ftype VariableMap.t) : (ftype, 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
    )
  | Function (xs, typef, fbody) -> (
      (* first check that the function has the right specified type then check
         that the number of inputs are the right number, finally eval the type
         of the body using the bindings for the inputs *)
      match typef with
        FunctionType (tin, tout) -> (
          if List.length xs <> List.length tin then
            Error (`WrongTypeSpecification
                     ("Type specification for function has wrong arity."))
          else
            let context1 = List.fold_left2
                (fun acc x t -> VariableMap.add x t acc)
                context
                xs
                tin
            in
            let* typefbody = evaluate_type fbody context1 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, xs) -> (
      (* check that the type is actually a function, then checks that the
         supplied inputs are of the right type, returns the return type if all
         inputs are supplied, otherwise a function from the remaining inputs to
         the output types *)
      let* evalf = evaluate_type f context in
      match evalf with
        FunctionType (tin, tout) -> (
          let rec helper (params: t_exp list) (typeparams: ftype list) =
            (* consumes until params are exausted *)
            match (params, typeparams) with
              ([], _) -> Ok typeparams
            | (_, []) -> Error (`WrongArity ("Function application has arity " ^
                                            (List.length tin |> string_of_int) ^
                                            ", but was applied to " ^
                                            (List.length xs |> string_of_int) ^
                                            " parameters"))
            | (p::tlparams, v::tltypeparams) ->
              if evaluate_type p context = Ok v then
                helper tlparams tltypeparams
              else
                Error (`WrongType "Argument with wrong type.")
          in
          let* typesremaining = helper xs tin in
          match typesremaining with
            [] -> Ok tout
          | t -> Ok (FunctionType (t, tout))
        )
      | _ -> 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.")
    )
  | 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, xs, typef, fbody, rest) ->
    (* like with the function type, but also add f itself to the bindings *)
    match typef with
      FunctionType (tin, tout) -> (
        if List.length xs <> List.length tin then
          Error (`WrongArity "Type specification for function has wrong arity.")
        else
          let context1 = VariableMap.add f typef context in
          let context2 = List.fold_left2
              (fun acc x t -> VariableMap.add x t acc)
              context1
              xs
              tin
          in
          let* typefbody = evaluate_type fbody context2 in
          let* typerest = evaluate_type rest context1 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, error) result =
  let* typeprogram = evaluate_type program VariableMap.empty in
  match typeprogram with
    FunctionType ([IntegerType], IntegerType) -> (
      Ok (FunctionType ([IntegerType], IntegerType))
    )
  | _ -> Error (`WrongType "Program is not a function from int to int.")