open StdLabels

(** This module provide a way to create new Id dynamically in the runtime,
    and some fonctions for comparing them. *)
module Id : sig
  type 'a typeid
  (** The type created on-the-fly. *)

  val newtype : unit -> 'a typeid
  (** Create a new instance of a dynamic type *)

  type ('a, 'b) eq = Eq : ('a, 'a) eq

  val try_cast : 'a typeid -> 'b typeid -> ('a, 'b) eq option
  (** Compare two types using the Eq pattern *)
end = struct
  type 'a witness = ..

  module type Witness = sig
    type t
    type _ witness += Id : t witness
  end

  type 'a typeid = (module Witness with type t = 'a)
  type ('a, 'b) eq = Eq : ('a, 'a) eq

  let try_cast : type a b. a typeid -> b typeid -> (a, b) eq option =
   fun x y ->
    let module X : Witness with type t = a = (val x) in
    let module Y : Witness with type t = b = (val y) in
    match X.Id with Y.Id -> Some Eq | _ -> None

  let newtype (type u) () =
    (* The extensible type need to be extended in a module, it is not possible
       to declare a type in a function. That’s why we need to pack a module
       here *)
    let module Witness = struct
      type t = u
      type _ witness += Id : t witness
    end in
    (module Witness : Witness with type t = u)
end

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

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

type t =
  | E : {
      module_ :
        (module S.Analyzer
           with type Expression.t = 'a
            and type Instruction.t = 'b
            and type Location.t = 'c);
      expr_witness : 'a Id.typeid;
      instr_witness : 'b Id.typeid;
      location_witness : 'c Id.typeid;
    }
      -> t

let build :
    (module S.Analyzer
       with type Expression.t = 'a
        and type Instruction.t = 'b
        and type Location.t = 'c) ->
    'a Id.typeid * 'b Id.typeid * 'c Id.typeid * t =
 fun module_ ->
  let expr_witness = Id.newtype ()
  and instr_witness = Id.newtype ()
  and location_witness = Id.newtype () in
  let t = E { module_; expr_witness; instr_witness; location_witness } in
  (expr_witness, instr_witness, location_witness, t)

module type App = sig
  val t : t array
end

module Helper = struct
  type 'a args_list = { witness : 'a Id.typeid; values : 'a S.repr list }
  (** This types helps the compiler to know which kind of arguments are hold
        inside the list. This is just a list with the additionnal witnesse
        information *)

  (** Extract all the lines from the given module 

      **Beware** The values are reversed. You should apply a List.rev if you
      want to keep them in the same order than the modules to apply.
    *)
  let args_i : result array list -> 'a Id.typeid -> int -> 'a args_list =
   fun args witness i ->
    let result =
      List.fold_left args ~init:{ values = []; witness }
        ~f:(fun (type a) ({ values; witness } : a args_list) t : a args_list ->
          match get witness (Array.get t i) with
          | None -> failwith "Does not match"
          | Some value_1 ->
              { values = (fun r -> (value_1, r)) :: values; witness })
    in
    { result with values = result.values }

  let map_args report args =
    List.fold_left_map args ~init:report ~f:(fun report v ->
        let v, result = v report in
        (result, v))
end

module Make (A : App) = struct
  module Expression : S.Expression with type t' = result array = struct
    type t = result array
    type t' = result array

    let literal : S.pos -> string -> t S.repr =
     fun pos value report ->
      let report, values =
        Array.fold_left_map A.t ~init:report
          ~f:(fun report (E { module_ = (module S); expr_witness; _ }) ->
            let value, report = S.Expression.literal pos value report in
            (report, R { value; witness = expr_witness }))
      in
      (values, report)

    let integer : S.pos -> string -> t S.repr =
     fun pos value report ->
      let report, values =
        Array.fold_left_map A.t ~init:report
          ~f:(fun report (E { module_ = (module S); expr_witness; _ }) ->
            let value, report = S.Expression.integer pos value report in
            (report, R { value; witness = expr_witness }))
      in
      (values, report)

    (** Unary operator like [-123] or [+'Text']*)
    let uoperator : S.pos -> T.uoperator -> t S.repr -> t S.repr =
     fun pos op values report ->
      (* Evaluate the nested expression *)
      let results, report = values report in

      (* Now evaluate the remaining expression.

         Traverse both the module the apply, and the matching expression already
         evaluated.

         It’s easer to use [map] and declare [report] as reference instead of
         [fold_left2] and accumulate the report inside the closure, because I
         don’t manage the order of the results.
      *)
      let report = ref report in
      let results =
        Array.map2 A.t results
          ~f:(fun (E { module_ = (module S); expr_witness; _ }) value ->
            match get expr_witness value with
            | None -> failwith "Does not match"
            | Some value ->
                (* Evaluate the single expression *)
                let value, report' =
                  S.Expression.uoperator pos op (fun r -> (value, r)) !report
                in
                report := report';
                R { witness = expr_witness; value })
      in
      (results, !report)

    (** Basically the same as uoperator, but operate over two operands instead
        of a single one. 

        In order to operate over the values (application, op1, op2) I’ve
        written a function [take_arg] which works like a [Array.map3] *)
    let boperator : S.pos -> T.boperator -> t S.repr -> t S.repr -> t S.repr =
     fun pos op expr1 expr2 report ->
      let expr1, report = expr1 report in
      let expr2, report = expr2 report in

      let report = ref report in

      let take_arg : result array -> result array -> result array =
       fun expr1 expr2 ->
        let len = Array.length A.t in
        Array.init len ~f:(fun i ->
            let (E { module_ = (module S); expr_witness; _ }) =
              Array.get A.t 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, r =
                  S.Expression.boperator pos op
                    (fun r -> (value_1, r))
                    (fun r -> (value_2, r))
                    !report
                in
                report := r;
                R { witness = expr_witness; value }
            | _ -> failwith "Does not match")
      in

      let results = take_arg expr1 expr2 in
      (results, !report)

    (** Call a function. The functions list is hardcoded in lib/lexer.mll *)
    let function_ : S.pos -> T.function_ -> t S.repr list -> t S.repr =
     fun pos func args report ->
      let report, args = Helper.map_args report args in
      let report = ref report and len = Array.length A.t in
      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module S); expr_witness; _ }) =
              Array.get A.t i
            in
            (* Extract the arguments for each module *)
            let args_i = Helper.args_i args expr_witness i in

            let value, r =
              S.Expression.function_ pos func (List.rev args_i.values) !report
            in
            report := r;
            R { witness = expr_witness; value })
      in
      (result, !report)

    let ident : (S.pos, t S.repr) S.variable -> t S.repr =
     fun { pos : S.pos; name : string; index : t S.repr option } report ->
      let len = Array.length A.t in

      let report = ref report in
      let index =
        Option.map
          (fun v ->
            let v, r = v !report in
            report := r;
            v)
          index
      in

      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module S); expr_witness; _ }) =
              Array.get A.t i
            in

            match index with
            | None ->
                (* Easest case, just return the plain ident *)
                let value, r =
                  S.Expression.ident { pos; name; index = None } !report
                in
                report := r;
                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, r =
                      S.Expression.ident
                        { pos; name; index = Some (fun r -> (value_1, r)) }
                        !report
                    in
                    report := r;
                    R { witness = expr_witness; value }))
      in
      (result, !report)

    let v : t * Report.t list -> t' * Report.t list = fun t -> t
  end

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

    let location : S.pos -> string -> t S.repr =
     fun pos label report ->
      let report, values =
        Array.fold_left_map A.t ~init:report
          ~f:(fun report (E { module_ = (module S); instr_witness; _ }) ->
            let value, report = S.Instruction.location pos label report in

            (report, R { value; witness = instr_witness }))
      in
      (values, report)

    let comment : S.pos -> t S.repr =
     fun pos report ->
      let report, values =
        Array.fold_left_map A.t ~init:report
          ~f:(fun report (E { module_ = (module S); instr_witness; _ }) ->
            let value, report = S.Instruction.comment pos report in

            (report, R { value; witness = instr_witness }))
      in
      (values, report)

    let expression : expression -> t S.repr =
     fun expr report ->
      let expr, report = expr report in
      let report = ref report in
      let results =
        Array.map2 A.t expr
          ~f:(fun
              (E { module_ = (module S); instr_witness; expr_witness; _ })
              (R { value; witness })
            ->
            match Id.try_cast witness expr_witness with
            | None -> failwith "Does not match"
            | Some Eq ->
                (* The evaluate the instruction *)
                let value, r =
                  S.Instruction.expression
                    (fun r -> S.Expression.v (value, r))
                    !report
                in
                report := r;
                R { value; witness = instr_witness })
      in
      (results, !report)

    let call : S.pos -> T.keywords -> expression list -> t S.repr =
     fun pos keyword args report ->
      let report, args = Helper.map_args report args in
      let report = ref report and len = Array.length A.t in
      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module S); expr_witness; instr_witness; _ }) =
              Array.get A.t i
            in
            let args_i = Helper.args_i args expr_witness i in
            let values =
              List.rev_map args_i.values ~f:(fun value r ->
                  S.Expression.v (value r))
            in

            let value, r = S.Instruction.call pos keyword values !report in
            report := r;
            R { witness = instr_witness; value })
      in
      (result, !report)

    let act : S.pos -> label:expression -> t S.repr list -> t S.repr =
     fun pos ~label instructions report ->
      let label, report = label report in
      let report, instructions = Helper.map_args report instructions in
      let report = ref report and len = Array.length A.t in

      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module S); instr_witness; expr_witness; _ }) =
              Array.get A.t i
            in
            let args_i = Helper.args_i instructions instr_witness i in
            let values =
              List.rev_map args_i.values ~f:(fun value r -> value r)
            in

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

      (result, !report)

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

    let assign :
        S.pos ->
        (S.pos, expression) S.variable ->
        T.assignation_operator ->
        expression ->
        t S.repr =
     fun pos { pos = var_pos; name; index } op expression report ->
      let expression, report = expression report in
      let report = ref report and len = Array.length A.t in

      let index =
        Option.map
          (fun v ->
            let v, r = v !report in
            report := r;
            v)
          index
      in

      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module A); instr_witness; expr_witness; _ }) =
              Array.get A.t i
            in

            let index_i =
              Option.map
                (fun expression ->
                  match get expr_witness (Array.get expression i) with
                  | None -> failwith "Does not match"
                  | Some value ->
                      let value r = A.Expression.v (value, r) in
                      value)
                index
            in
            let variable = S.{ pos = var_pos; name; index = index_i } in

            match get expr_witness (Array.get expression i) with
            | None -> failwith "Does not match"
            | Some value ->
                let value, r =
                  A.Instruction.assign pos variable op
                    (fun r -> A.Expression.v (value, r))
                    !report
                in
                report := r;

                R { value; witness = instr_witness })
      in

      (result, !report)

    (** Helper function used to prepare the clauses *)
    let map_clause :
        Report.t list ->
        (expression, t) S.clause ->
        Report.t list * (S.pos * Expression.t' * t list) =
     fun report clause ->
      let clause_pos, expression, t = clause in
      let expression, report = expression report in
      let report, t =
        List.fold_left_map t ~init:report ~f:(fun report t ->
            let t, report = t report in
            (report, t))
      in
      let clause = (clause_pos, expression, t) in
      (report, clause)

    let rebuild_clause :
        type a b.
        int ->
        a Id.typeid ->
        b Id.typeid ->
        (b -> 'c) ->
        S.pos * result array * result array list ->
        ('c, a) S.clause =
     fun i instr_witness expr_witness f clause ->
      let pos_clause, expr_clause, ts = clause in
      match get expr_witness (Array.get expr_clause i) with
      | None -> failwith "Does not match"
      | Some value ->
          let ts = Helper.args_i ts instr_witness i in
          let ts = List.rev_map ts.values ~f:(fun value r -> value r) in
          let clause = (pos_clause, f value, ts) in
          clause

    let if_ :
        S.pos ->
        (expression, t) S.clause ->
        elifs:(expression, t) S.clause list ->
        else_:t S.repr list ->
        t S.repr =
     fun pos clause ~elifs ~else_ report ->
      (* First, apply the report for all the instructions *)
      let report, clause = map_clause report clause in
      let report, elifs = List.fold_left_map elifs ~init:report ~f:map_clause in
      let report, else_ = Helper.map_args report else_ in
      let report = ref report and len = Array.length A.t in

      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module A); instr_witness; expr_witness; _ }) =
              Array.get A.t i
            in

            (* This function helps to build the expression in the clauses *)
            let f v r = A.Expression.v (v, r) in

            let clause = rebuild_clause i instr_witness expr_witness f clause
            and elifs =
              List.map elifs ~f:(rebuild_clause i instr_witness expr_witness f)
            and elses = Helper.args_i else_ instr_witness i in

            let else_ = List.rev elses.values in

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

      (result, !report)

    let v : t * Report.t list -> t' * Report.t list = fun t -> t
  end

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

    let location : S.pos -> instruction S.repr list -> t S.repr =
     fun pos instructions report ->
      ignore pos;

      let report, instructions = Helper.map_args report instructions in

      let report = ref report and len = Array.length A.t in
      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module A); instr_witness; location_witness; _ })
                =
              Array.get A.t i
            in

            let instructions_i : A.Instruction.t Helper.args_list =
              Helper.args_i instructions instr_witness i
            in
            let inst : A.Instruction.t S.repr list = instructions_i.values in
            let instructions : A.Instruction.t' S.repr list =
              List.rev_map inst ~f:(fun value report ->
                  let value, report = value report in
                  A.Instruction.v (value, report))
            in
            let value, re = A.Location.location pos instructions !report in
            report := re;
            R { value; witness = location_witness })
      in
      (result, !report)
  end
end