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 -> (
      match VariableMap.find_opt x context with
        None -> Error (`AbsentAssignment ("The variable " ^ x ^ " is not defined."))
      | Some t -> Ok t
    )
  | Function (xs, typef, fbody) -> (
      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
            match (typefbody = tout) with
              (false) -> Error (`WrongTypeSpecification "Function does not return specified type.")
            | (true) -> Ok typef
        )
      | _ -> Error (`WrongTypeSpecification "Specification of function is not a function type.")
    )
  | Application (f, xs) -> (
      let* evalf = evaluate_type f context in
      match evalf with
        FunctionType (tin, tout) -> (
          let rec helper params typeparams =
            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) ->
    let* typex = evaluate_type xval context in
    evaluate_type rest (VariableMap.add x typex context)
  | LetFun (f, xs, typef, fbody, rest) ->
    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.")