open StdLabels
module S = Qsp_syntax.S
module T = Qsp_syntax.T
module Report = Qsp_syntax.Report

let identifier = "type_check"
let description = "Ensure all the expression are correctly typed"
let is_global = false
let active = ref true

type context = unit

let initialize = Fun.id
let finalize () = []

module Helper = struct
  type argument_repr = { pos : S.pos; t : Get_type.t }

  module DynType = struct
    type nonrec t = Get_type.t -> Get_type.t
    (** Dynamic type is a type unknown during the code.

        For example, the equality operator accept either Integer or String, but
        we expect that both sides of the equality uses the same type.*)

    (** Build a new dynamic type *)
    let t : unit -> t =
     fun () ->
      let stored = ref None in
      fun t ->
        match !stored with
        | None ->
            stored := Some t;
            t
        | Some t -> t
  end

  (** Declare an argument for a function.

      - Either we already know the type and we just have to compare.
      - Either the type shall constrained by another one
      - Or we have a variable number of arguments. *)
  type argument =
    | Fixed of Get_type.type_of
    | Dynamic of DynType.t
    | Variable of argument

  let compare :
      ?level:Report.level ->
      strict:bool ->
      Get_type.type_of ->
      argument_repr ->
      Report.t list ->
      Report.t list =
   fun ?(level = Report.Warn) ~strict expected actual report ->
    let equal =
      match (expected, actual.t) with
      (* Strict equality for this ones, always true *)
      | String, Variable String
      | String, Raw String
      | String, Variable NumericString
      | String, Raw NumericString
      | Integer, Variable Integer
      | Integer, Raw Integer
      | NumericString, Variable NumericString
      | NumericString, Raw NumericString
      | Bool, Raw Bool
      | Bool, Variable Bool
      (* Also include the conversion between bool and integer *)
      | Integer, Raw Bool
      | Integer, Variable Bool
      (* The type NumericString can be used as a generic type in input *)
      | _, Variable NumericString
      | NumericString, Raw String
      | NumericString, Variable String
      | NumericString, Raw Integer
      | NumericString, Variable Integer ->
          true
      | Bool, Variable Integer
      | Bool, Raw Integer
      | String, Variable Integer
      | String, Raw Bool
      | String, Variable Bool
      | Integer, Variable String
      | Integer, Raw NumericString ->
          not strict
      (* Explicit rejected cases  *)
      | String, Raw Integer | Integer, Raw String -> false
      | _, _ -> false
    in
    if equal then report
    else
      let result_type = match actual.t with Variable v -> v | Raw r -> r in
      let message =
        Format.asprintf "The type %a is expected but got %a" Get_type.pp_type_of
          expected Get_type.pp_type_of result_type
      in
      Report.message level actual.pos message :: report

  let rec compare_parameter :
      strict:bool ->
      ?level:Report.level ->
      argument ->
      argument_repr ->
      Report.t list ->
      Report.t list =
   fun ~strict ?(level = Report.Warn) expected param report ->
    match expected with
    | Fixed t -> compare ~strict ~level t param report
    | Dynamic d ->
        let type_ = match d param.t with Raw r -> r | Variable v -> v in
        compare ~strict ~level type_ param report
    | Variable c -> compare_parameter ~level ~strict c param report

  (** Compare the arguments one by one *)
  let compare_args :
      ?strict:bool ->
      ?level:Report.level ->
      S.pos ->
      argument list ->
      argument_repr list ->
      Report.t list ->
      Report.t list =
   fun ?(strict = false) ?(level = Report.Warn) pos expected actuals report ->
    let tl, report =
      List.fold_left actuals ~init:(expected, report)
        ~f:(fun (expected, report) param ->
          match expected with
          | (Variable _ as hd) :: _ ->
              let check = compare_parameter ~strict ~level hd param report in
              (expected, check)
          | hd :: tl ->
              let check = compare_parameter ~strict ~level hd param report in
              (tl, check)
          | [] ->
              let msg = Report.error param.pos "Unexpected argument" in
              ([], msg :: report))
    in
    match tl with
    | [] | Variable _ :: _ -> report
    | _ ->
        let msg = Report.error pos "Not enougth arguments given" in
        msg :: report
end

type t' = { result : Get_type.t Lazy.t; pos : S.pos }

let arg_of_repr : Get_type.t Lazy.t -> S.pos -> Helper.argument_repr =
 fun type_of pos -> { pos; t = Lazy.force type_of }

module TypedExpression = struct
  type nonrec t' = t' * Report.t list
  type state = { pos : S.pos }
  type t = state * Report.t list

  let v : Get_type.t Lazy.t * t -> t' =
   fun (type_of, (t, r)) -> ({ result = type_of; pos = t.pos }, r)

  (** The variable has type string when starting with a '$' *)
  let ident :
      (S.pos, Get_type.t Lazy.t * t) S.variable -> Get_type.t Lazy.t -> t =
   fun var _type_of ->
    (* Extract the error from the index *)
    let report =
      match var.index with
      | None -> []
      | Some (_, expr) ->
          let _, r = expr in
          r
    in
    ({ pos = var.pos }, report)

  let integer : S.pos -> string -> Get_type.t Lazy.t -> t =
   fun pos value _type_of ->
    let int_value = int_of_string_opt value in

    let report =
      match int_value with
      | Some 0 -> []
      | Some _ -> []
      | None -> Report.error pos "Invalid integer value" :: []
    in

    ({ pos }, report)

  let literal :
      S.pos -> (Get_type.t Lazy.t * t) T.literal list -> Get_type.t Lazy.t -> t
      =
   fun pos values type_of ->
    ignore type_of;
    let init = [] in
    let report =
      List.fold_left values ~init ~f:(fun report -> function
        | T.Text _ -> report
        | T.Expression (_, t) ->
            let report = List.rev_append (snd t) report in
            report)
    in
    ({ pos }, report)

  let function_ :
      S.pos ->
      T.function_ ->
      (Get_type.t Lazy.t * t) list ->
      Get_type.t Lazy.t ->
      t =
   fun pos function_ params _type_of ->
    (* Accumulate the expressions and get the results, the report is given in
       the differents arguments, and we build a list with the type of the
       parameters. *)
    let types, report =
      List.fold_left params ~init:([], [])
        ~f:(fun (types, report) (type_of, param) ->
          ignore type_of;
          let t, r = param in
          let arg = arg_of_repr type_of t.pos in
          (arg :: types, r @ report))
    in
    let types = List.rev types and default = { pos } in

    match function_ with
    | Arrcomp | Arrpos | Arrsize | Countobj | Desc | Dyneval | Getobj | Instr
    | Isplay ->
        (default, report)
    | Desc' | Dyneval' | Getobj' -> (default, report)
    | Func | Func' -> (default, report)
    | Iif | Iif' ->
        let d = Helper.DynType.t () in
        let expected = Helper.[ Fixed Bool; Dynamic d; Dynamic d ] in
        let report = Helper.compare_args pos expected types report in
        (* Extract the type for the expression *)
        ({ pos }, report)
    | Input | Input' ->
        (* Input should check the result if the variable is a num and raise a
           message in this case.*)
        let expected = Helper.[ Fixed String ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Isnum ->
        let expected = Helper.[ Fixed String ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Lcase | Lcase' | Ucase | Ucase' ->
        let expected = Helper.[ Fixed String ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Len ->
        let expected = Helper.[ Fixed NumericString ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Loc ->
        let expected = Helper.[ Fixed String ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Max | Max' | Min | Min' ->
        let d = Helper.DynType.t () in
        (* All the arguments must have the same type *)
        let expected = Helper.[ Variable (Dynamic d) ] in
        let report = Helper.compare_args pos expected types report in
        ({ pos }, report)
    | Mid | Mid' ->
        let expected = Helper.[ Fixed String; Variable (Fixed Integer) ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Msecscount -> (default, report)
    | Rand ->
        let expected = Helper.[ Variable (Fixed Integer) ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Replace -> (default, report)
    | Replace' -> (default, report)
    | Rgb -> (default, report)
    | Rnd ->
        (* No arg *)
        let report = Helper.compare_args pos [] types report in
        (default, report)
    | Selact -> (default, report)
    | Str | Str' ->
        let expected = Helper.[ Variable (Fixed Integer) ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)
    | Strcomp -> (default, report)
    | Strfind -> (default, report)
    | Strfind' -> (default, report)
    | Strpos -> (default, report)
    | Trim -> (default, report)
    | Trim' -> (default, report)
    | Val ->
        let expected = Helper.[ Fixed NumericString ] in
        let report = Helper.compare_args pos expected types report in
        (default, report)

  (** Unary operator like [-123] or [+'Text']*)
  let uoperator :
      S.pos -> T.uoperator -> Get_type.t Lazy.t * t -> Get_type.t Lazy.t -> t =
   fun pos operator t1 type_of ->
    ignore type_of;
    let type_of, (t, report) = t1 in
    match operator with
    | Add -> (t, report)
    | Neg | No ->
        let types = [ arg_of_repr type_of t.pos ] in
        let expected = Helper.[ Fixed Integer ] in
        let report = Helper.compare_args pos expected types report in
        ({ pos }, report)

  let boperator :
      S.pos ->
      T.boperator ->
      Get_type.t Lazy.t * t ->
      Get_type.t Lazy.t * t ->
      Get_type.t Lazy.t ->
      t =
   fun pos operator (type_1, t1) (type_2, t2) type_of ->
    ignore type_of;
    let t1, report1 = t1 in
    let t2, report2 = t2 in

    let report = report1 @ report2 in

    let types = [ arg_of_repr type_1 t1.pos; arg_of_repr type_2 t2.pos ] in

    match operator with
    | T.Plus ->
        (* We cannot really much here, because the (+) function can be used to
           concatenate string or add numbers.

           When concatenating, it’s allowed to add an integer and a number.
        *)
        ({ pos }, report)
    | T.Eq | T.Neq | Lt | Gte | Lte | Gt ->
        (* If the expression is '' or 0, we accept the comparaison as if
            instead of raising a warning *)
        let d = Helper.(Dynamic (DynType.t ())) in
        let expected = [ d; d ] in
        (* Compare and report as error if the types are incompatible. If no
           error is reported, try in strict mode, and report as a warning. *)
        let report =
          match
            Helper.compare_args ~level:Error pos expected (List.rev types)
              report
          with
          | [] ->
              Helper.compare_args ~strict:true pos expected (List.rev types)
                report
          | report -> report
        in
        ({ pos }, report)
    | T.Mod | T.Minus | T.Product | T.Div ->
        (* Operation over number *)
        let expected = Helper.[ Fixed Integer; Fixed Integer ] in
        let report = Helper.compare_args pos expected types report in
        ({ pos }, report)
    | T.And | T.Or ->
        (* Operation over booleans *)
        let expected = Helper.[ Fixed Bool; Fixed Bool ] in
        let report = Helper.compare_args pos expected types report in
        ({ pos }, report)
end

module Expression = Compose.TypeBuilder.Make (TypedExpression)

module Instruction = struct
  type t = Report.t list
  type t' = Report.t list

  let v : t -> t' = fun local_report -> local_report

  type expression = Expression.t'

  (** Call for an instruction like [GT] or [*CLR] *)
  let call : S.pos -> T.keywords -> expression list -> t =
   fun _pos _ expressions ->
    List.fold_left expressions ~init:[] ~f:(fun acc a ->
        let _, report = a in
        (List.rev_append report) acc)

  let location : S.pos -> string -> t = fun _pos _ -> []

  (** Comment *)
  let comment : S.pos -> t = fun _pos -> []

  (** Raw expression *)
  let expression : expression -> t = fun expression -> snd expression

  (** Helper function used in the [if_] function. *)
  let fold_clause : t -> (expression, t) S.clause -> t =
   fun report (_pos, expr, instructions) ->
    let result, r = expr in

    let r2 =
      Helper.compare ~strict:false Get_type.Bool
        (arg_of_repr result.result result.pos)
        []
    in

    List.fold_left instructions
      ~init:(r @ r2 @ report)
      ~f:(fun acc a ->
        let report = a in
        (List.rev_append report) acc)

  let if_ :
      S.pos ->
      (expression, t) S.clause ->
      elifs:(expression, t) S.clause list ->
      else_:(S.pos * t list) option ->
      t =
   fun _pos clause ~elifs ~else_ ->
    (* Traverse the whole block recursively *)
    let report = fold_clause [] clause in
    let report = List.fold_left elifs ~f:fold_clause ~init:report in

    match else_ with
    | None -> report
    | Some (_, instructions) ->
        List.fold_left instructions ~init:report ~f:(fun acc a ->
            let report = a in
            (List.rev_append report) acc)

  let act : S.pos -> label:expression -> t list -> t =
   fun _pos ~label instructions ->
    let result, report = label in
    let report =
      Helper.compare ~strict:false Get_type.String
        (arg_of_repr result.result result.pos)
        report
    in

    List.fold_left instructions ~init:report ~f:(fun acc a ->
        let report = a in
        (List.rev_append report) acc)

  let assign :
      S.pos ->
      (S.pos, expression) S.variable ->
      T.assignation_operator ->
      expression ->
      t =
   fun pos variable op expression ->
    let right_expression, report = expression in

    let report' = Option.map snd variable.index |> Option.value ~default:[] in

    let report = List.rev_append report' report in

    match (op, Get_type.get_type (Lazy.force right_expression.result)) with
    | T.Eq', Get_type.Integer ->
        (* Assigning an intger is allowed in a string variable, but raise a
           warning. *)
        let var_type = Lazy.from_val (Get_type.ident variable) in
        let op1 = arg_of_repr var_type variable.pos in
        let expected = Helper.[ Fixed Integer ] in
        Helper.compare_args ~strict:true ~level:Report.Warn pos expected [ op1 ]
          report
    | _, _ -> (
        let var_type = Lazy.from_val (Get_type.ident variable) in
        let op1 = arg_of_repr var_type variable.pos in
        let op2 = arg_of_repr right_expression.result right_expression.pos in

        let d = Helper.DynType.t () in
        (* Every part of the assignation should be the same type *)
        let expected = Helper.[ Dynamic d; Dynamic d ] in

        match
          Helper.compare_args ~strict:false ~level:Report.Error pos expected
            [ op1; op2 ] []
        with
        | [] ->
            Helper.compare_args ~strict:true ~level:Report.Warn pos expected
              [ op1; op2 ] report
        | reports -> reports @ report)
end

module Location = struct
  type t = Report.t list
  type instruction = Instruction.t'

  let v = Fun.id

  let location : unit -> S.pos -> instruction list -> t =
   fun () _pos instructions ->
    let report =
      List.fold_left instructions ~init:[] ~f:(fun report instruction ->
          let report' = instruction in
          report' @ report)
    in
    report
end