module S = Qsp_syntax.S

(** The the Id module, wrap a value in an existencial type with a witness
    associate with. *)
type result = R : { value : 'a; witness : 'a Type.Id.t } -> result

let get : type a. a Type.Id.t -> result -> a option =
 fun typeid (R { value; witness }) ->
  match Type.Id.provably_equal typeid witness with
  | Some Type.Equal -> Some value
  | None -> None

let set : type a. a Type.Id.t -> result -> a -> result option =
 fun typeid (R { witness; _ }) value ->
  match Type.Id.provably_equal typeid witness with
  | Some Type.Equal -> Some (R { witness; value })
  | None -> None

type t = Qsp_syntax.Identifier.t

module type App = sig
  val t : t array
end

open StdLabels

module Helper = struct
  type 'a expr_list = { witness : 'a Type.Id.t; values : 'a list }

  let expr_i : result array list -> 'a Type.Id.t -> int -> 'a expr_list =
   fun args witness i ->
    let result =
      List.fold_left args ~init:{ values = []; witness }
        ~f:(fun (type a) ({ values; witness } : a expr_list) t : a expr_list ->
          match get witness (Array.get t i) with
          | None -> failwith "Does not match"
          | Some value_1 -> { values = value_1 :: values; witness })
    in
    { result with values = result.values }
end

module Make (A : App) = struct
  let identifier = "main_checker"
  let description = "Internal module"
  let is_global = false
  let active = ref false
  let depends = [] (* This modules depends of nothing *)

  type ex = Qsp_syntax.Identifier.t

  type context = result Array.t
  (** We associate each context from the differents test in an array. The
      context for this module is a sort of context of contexts *)

  (* Initialize the modules to check here, at the module level *)
  let checks =
    (* Collect all the dependencies and build the execution order *)
    let graph : ex list =
      let rec build_deps l acc =
        List.fold_left
          (l : ex list)
          ~init:acc
          ~f:(fun
              acc (Qsp_syntax.Identifier.E { module_ = (module S); _ } as ex) ->
            let acc' = ex :: acc in

            build_deps S.depends acc')
      in
      build_deps (Array.to_list A.t) []
    in

    (* Convert the dependenciees using the module identifier only, the
       Tsort.sort function use structural equality comparaison function which
       does not wok with the module embeded in first class module.  *)
    let graph_name =
      List.map graph
        ~f:(fun (Qsp_syntax.Identifier.E { module_ = (module S); _ }) ->
          let deps' =
            List.map
              ~f:(fun (Qsp_syntax.Identifier.E { module_ = (module S); _ }) ->
                S.identifier)
              S.depends
          in
          (S.identifier, deps'))
    in
    match Tsort.sort graph_name with
    | Tsort.Sorted sorted_graph ->
        (* From the identifier, extract the associated check *)
        let _ =
          List.map sorted_graph ~f:(fun name ->
              List.find_map graph
                ~f:(fun
                    (Qsp_syntax.Identifier.E { module_ = (module S); _ } as
                     check)
                  ->
                  match String.equal name S.identifier with
                  | false -> None
                  | true -> Some check)
              |> Option.get)
          (* It’s ok to use unchecked option.get here, because the list was
               created from the same source just before *)
        in

        Array.of_list graph
    | Tsort.ErrorCycle _ ->
        (* This is very unlikely to happen, as it would reflect an error in
           the compilation units order *)
        raise Not_found

  (** Initialize each test, and keep the result in the context. *)
  let initialize : unit -> context =
   fun () ->
    Array.map checks
      ~f:(fun (Qsp_syntax.Identifier.E { module_ = (module S); context; _ }) ->
        let value = S.initialize () in
        R { value; witness = context })

  let finalize : result Array.t -> (string * Qsp_syntax.Report.t) list =
   fun context_array ->
    let _, report =
      Array.fold_left checks ~init:(0, [])
        ~f:(fun
            (i, acc)
            (Qsp_syntax.Identifier.E { module_ = (module S); context; _ })
          ->
          let result = Array.get context_array i in
          let local_context = Option.get (get context result) in
          let reports = S.finalize local_context in
          (i + 1, List.rev_append reports acc))
    in
    report

  (* Global variable for the whole module *)
  let len = Array.length checks

  module Expression : S.Expression with type t' = result array = struct
    type t = result array
    type t' = result array

    let build_ctx : result option array -> Qsp_syntax.S.extract_context =
     fun results ->
      {
        f =
          (fun id ->
            Array.find_map results ~f:(function
              | Some result -> get id result
              | None -> None));
      }

    let literal :
        ctx:Qsp_syntax.S.extract_context ->
        S.pos ->
        t Qsp_syntax.T.literal list ->
        t =
     fun ~ctx pos values ->
      ignore ctx;
      let results = Array.make len None in
      (* Create the new array, filled with None at the begining.

         Then populate the array in place in order to read the previous values
         if requested *)
      (* Extract the result with the given ID from the array *)
      let ctx = build_ctx results in
      let () =
        Array.mapi_inplace results ~f:(fun i _ ->
            let (Qsp_syntax.Identifier.E
                   { module_ = (module S); expr_witness; _ }) =
              Array.get checks i
            in
            (* Map every values to the Checker *)
            let values' =
              List.map values
                ~f:
                  (Qsp_syntax.T.map_litteral ~f:(fun expr ->
                       Option.get (get expr_witness (Array.get expr i))))
            in
            let value = S.Expression.literal ~ctx pos values' in
            Some (R { value; witness = expr_witness }))
      in
      Array.map results ~f:Option.get

    let integer : ctx:Qsp_syntax.S.extract_context -> S.pos -> string -> t =
     fun ~ctx pos value ->
      ignore ctx;
      let results = Array.make len None in
      let ctx = build_ctx results in

      let () =
        Array.mapi_inplace results ~f:(fun i _ ->
            let (Qsp_syntax.Identifier.E
                   { module_ = (module S); expr_witness; _ }) =
              Array.get checks i
            in
            let value = S.Expression.integer ~ctx pos value in
            Some (R { value; witness = expr_witness }))
      in
      Array.map results ~f:Option.get

    (** Unary operator like [-123] or [+'Text']*)
    let uoperator :
        ctx:Qsp_syntax.S.extract_context ->
        S.pos ->
        Qsp_syntax.T.uoperator ->
        t ->
        t =
     fun ~ctx pos op values ->
      ignore ctx;
      let results = Array.make len None in
      let ctx = build_ctx results in

      (* Evaluate the nested expression *)
      let () =
        Array.mapi_inplace results ~f:(fun i _ ->
            let (Qsp_syntax.Identifier.E
                   { module_ = (module S); expr_witness; _ }) =
              Array.get checks i
            in
            let value = Array.get values i in
            match get expr_witness value with
            | None -> failwith "Does not match"
            | Some value ->
                (* Evaluate the single expression *)
                let value = S.Expression.uoperator ~ctx pos op value in
                Some (R { witness = expr_witness; value }))
      in
      Array.map results ~f:Option.get

    (** Basically the same as uoperator, but operate over two operands instead
        of a single one. *)
    let boperator :
        ctx:Qsp_syntax.S.extract_context ->
        S.pos ->
        Qsp_syntax.T.boperator ->
        t ->
        t ->
        t =
     fun ~ctx pos op expr1 expr2 ->
      ignore ctx;
      let results = Array.make len None in
      let ctx = build_ctx results in
      let () =
        Array.mapi_inplace results ~f:(fun i _ ->
            let (E { module_ = (module S); expr_witness; _ }) =
              Array.get checks i
            in
            match
              ( get expr_witness (Array.get expr1 i),
                get expr_witness (Array.get expr2 i) )
            with
            | Some value_1, Some value_2 ->
                let value =
                  S.Expression.boperator ~ctx pos op value_1 value_2
                in
                Some (R { witness = expr_witness; value })
            | _ -> failwith "Does not match")
      in
      Array.map results ~f:Option.get

    (** Call a function. The functions list is hardcoded in lib/lexer.mll *)
    let function_ :
        ctx:Qsp_syntax.S.extract_context ->
        S.pos ->
        Qsp_syntax.T.function_ ->
        t list ->
        t =
     fun ~ctx pos func args ->
      ignore ctx;
      let results = Array.make len None in
      let ctx = build_ctx results in
      let () =
        Array.mapi_inplace results ~f:(fun i _ ->
            let (E { module_ = (module S); expr_witness; _ }) =
              Array.get checks i
            in
            (* Extract the arguments for each module *)
            let args_i = List.rev (Helper.expr_i args expr_witness i).values in
            let value = S.Expression.function_ ~ctx pos func args_i in
            Some (R { witness = expr_witness; value }))
      in
      Array.map results ~f:Option.get

    let ident : ctx:Qsp_syntax.S.extract_context -> (S.pos, t) S.variable -> t =
     fun ~ctx { pos : S.pos; name : string; index : t option } ->
      ignore ctx;
      let results = Array.make len None in
      let ctx = build_ctx results in
      let () =
        Array.mapi_inplace results ~f:(fun i _ ->
            let (E { module_ = (module S); expr_witness; _ }) =
              Array.get checks i
            in

            match index with
            | None ->
                (* Easest case, just return the plain ident *)
                let value =
                  S.Expression.ident ~ctx { pos; name; index = None }
                in
                Some (R { witness = expr_witness; value })
            | Some t -> (
                match get expr_witness (Array.get t i) with
                | None -> failwith "Does not match"
                | Some value_1 ->
                    let value =
                      S.Expression.ident ~ctx
                        { pos; name; index = Some value_1 }
                    in
                    Some (R { witness = expr_witness; value })))
      in
      Array.map results ~f:Option.get

    (** Convert each internal represention for the expression into its external
        representation *)
    let v : t -> t' =
     fun t ->
      let result =
        Array.map2 checks t
          ~f:(fun
              (Qsp_syntax.Identifier.E
                 { module_ = (module S); expr_witness; expr'; _ })
              result
            ->
            match get expr_witness result with
            | None -> failwith "Does not match"
            | Some value ->
                let value = S.Expression.v value in
                R { witness = expr'; value })
      in
      result
  end

  module Instruction :
    S.Instruction
      with type expression = Expression.t'
       and type t' = result array = struct
    type expression = Expression.t'
    type t = result array
    type t' = result array

    let location : S.pos -> string -> t =
     fun pos label ->
      Array.map checks
        ~f:(fun
            (Qsp_syntax.Identifier.E { module_ = (module S); instr_witness; _ })
          ->
          let value = S.Instruction.location pos label in
          R { value; witness = instr_witness })

    let comment : S.pos -> t =
     fun pos ->
      Array.map checks
        ~f:(fun
            (Qsp_syntax.Identifier.E { module_ = (module S); instr_witness; _ })
          ->
          let value = S.Instruction.comment pos in
          R { value; witness = instr_witness })

    let expression : expression -> t =
     fun expr ->
      Array.map2 checks expr
        ~f:(fun
            (Qsp_syntax.Identifier.E
               { module_ = (module S); instr_witness; expr'; _ })
            result
          ->
          match get expr' result with
          | None -> failwith "Does not match"
          | Some value ->
              (* The evaluate the instruction *)
              let value = S.Instruction.expression value in
              R { value; witness = instr_witness })

    let call : S.pos -> Qsp_syntax.T.keywords -> expression list -> t =
     fun pos keyword args ->
      (* The arguments are given like an array of array. Each expression is
         actually the list of each expression in the differents modules. *)
      Array.init len ~f:(fun i ->
          let (E { module_ = (module S); expr'; instr_witness; _ }) =
            Array.get checks i
          in

          let values = List.rev (Helper.expr_i args expr' i).values in

          let value = S.Instruction.call pos keyword values in
          R { witness = instr_witness; value })

    let act : S.pos -> label:expression -> t list -> t =
     fun pos ~label instructions ->
      Array.init len ~f:(fun i ->
          let (E { module_ = (module S); instr_witness; expr'; _ }) =
            Array.get checks i
          in
          let values =
            List.rev (Helper.expr_i instructions instr_witness i).values
          in

          match get expr' (Array.get label i) with
          | None -> failwith "Does not match"
          | Some label_i ->
              let value = S.Instruction.act pos ~label:label_i values in
              R { witness = instr_witness; value })

    (* I think it’s one of the longest module I’ve ever written in OCaml… *)

    let assign :
        S.pos ->
        (S.pos, expression) S.variable ->
        Qsp_syntax.T.assignation_operator ->
        expression ->
        t =
     fun pos { pos = var_pos; name; index } op expression ->
      Array.init len ~f:(fun i ->
          let (E { module_ = (module A); instr_witness; expr'; _ }) =
            Array.get checks i
          in

          let index_i =
            Option.map
              (fun expression ->
                Option.get (get expr' (Array.get expression i)))
              index
          in
          let variable = S.{ pos = var_pos; name; index = index_i } in

          match get expr' (Array.get expression i) with
          | None -> failwith "Does not match"
          | Some value ->
              let value = A.Instruction.assign pos variable op value in

              R { value; witness = instr_witness })

    let rebuild_clause : type a b.
        int ->
        a Type.Id.t ->
        b Type.Id.t ->
        S.pos * result array * result array list ->
        (b, a) S.clause =
     fun i instr_witness expr' clause ->
      let pos_clause, expr_clause, ts = clause in
      match get expr' (Array.get expr_clause i) with
      | None -> failwith "Does not match"
      | Some value ->
          let ts = Helper.expr_i ts instr_witness i in
          let ts = List.rev ts.values in
          let clause = (pos_clause, value, ts) in
          clause

    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_ ->
      (* First, apply the report for all the instructions *)
      let else_ =
        match else_ with
        | None -> None
        | Some (pos, instructions) -> Some (pos, instructions)
      in
      Array.init len ~f:(fun i ->
          let (E { module_ = (module A); instr_witness; expr'; _ }) =
            Array.get checks i
          in

          let clause = rebuild_clause i instr_witness expr' clause
          and elifs = List.map elifs ~f:(rebuild_clause i instr_witness expr')
          and else_ =
            match else_ with
            | None -> None
            | Some (pos, instructions) ->
                let elses = Helper.expr_i instructions instr_witness i in
                Some (pos, List.rev elses.values)
          in

          let value = A.Instruction.if_ pos clause ~elifs ~else_ in
          R { value; witness = instr_witness })

    (** This code is almost a copy/paste from Expression.v but I did not found a
        way to factorize it. *)
    let v : t -> t' =
     fun t ->
      let result =
        Array.map2 checks t
          ~f:(fun
              (Qsp_syntax.Identifier.E
                 { module_ = (module S); instr_witness; instr'; _ })
              result
            ->
            match get instr_witness result with
            | None -> failwith "Does not match"
            | Some value ->
                let value = S.Instruction.v value in
                R { witness = instr'; value })
      in
      result
  end

  module Location :
    S.Location
      with type t = result array
       and type instruction = Instruction.t'
       and type context := context = struct
    type instruction = Instruction.t'
    type t = result array

    let location : context -> S.pos -> instruction list -> t =
     fun local_context pos args ->
      ignore pos;

      let result =
        Array.init len ~f:(fun i ->
            let (E
                   {
                     module_ = (module A);
                     instr';
                     location_witness;
                     context;
                     _;
                   }) =
              Array.get checks i
            in

            let local_context =
              Option.get (get context (Array.get local_context i))
            in

            let instructions = List.rev (Helper.expr_i args instr' i).values in
            let value = A.Location.location local_context pos instructions in
            R { value; witness = location_witness })
      in
      result

    let v : t -> Qsp_syntax.Report.t list =
     fun args ->
      let report = ref [] in
      let () =
        Array.iteri args ~f:(fun i result ->
            let (E { module_ = (module A); location_witness; _ }) =
              Array.get checks i
            in
            match get location_witness result with
            | None -> failwith "Does not match"
            | Some value ->
                let re = A.Location.v value in
                report := List.rev_append re !report)
      in
      !report
  end
end