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 Expression.t' = 'b
            and type Instruction.t = 'c
            and type Instruction.t' = 'd
            and type Location.t = 'e);
      expr_witness : 'a Id.typeid;
      expr' : ('b * Report.t list) Id.typeid;
      instr_witness : 'c Id.typeid;
      instr' : ('d * Report.t list) Id.typeid;
      location_witness : 'e Id.typeid;
    }
      -> t

let build :
    (module S.Analyzer
       with type Expression.t = _
        and type Expression.t' = _
        and type Instruction.t = _
        and type Instruction.t' = _
        and type Location.t = 'a) ->
    'a Id.typeid * t =
 fun module_ ->
  let expr_witness = Id.newtype ()
  and expr' = Id.newtype ()
  and instr_witness = Id.newtype ()
  and instr' = Id.newtype ()
  and location_witness = Id.newtype () in
  let t =
    E { module_; expr_witness; expr'; instr_witness; instr'; location_witness }
  in
  (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 witness
        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 _ -> value_1) :: values; witness })
    in
    { result with values = result.values }

  type 'a expr_list = { witness : 'a Id.typeid; values : 'a list }

  let expr_i : result array list -> 'a Id.typeid -> 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 }

  let map_args report args = List.map args ~f:(fun v -> v report)
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 =
     fun pos value ->
      Array.map A.t ~f:(fun (E { module_ = (module S); expr_witness; _ }) ->
          let value = S.Expression.literal pos value in
          R { value; witness = expr_witness })

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

    (** Unary operator like [-123] or [+'Text']*)
    let uoperator : S.pos -> T.uoperator -> t -> t =
     fun pos op values ->
      (* Evaluate the nested expression *)
      let results = values 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 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 = S.Expression.uoperator pos op value in
                R { witness = expr_witness; value })
      in
      results

    (** 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 -> t -> t =
     fun pos op expr1 expr2 ->
      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 = S.Expression.boperator pos op value_1 value_2 in
                R { witness = expr_witness; value }
            | _ -> failwith "Does not match")
      in

      take_arg expr1 expr2

    (** Call a function. The functions list is hardcoded in lib/lexer.mll *)
    let function_ : S.pos -> T.function_ -> t list -> t =
     fun pos func args ->
      let 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 = List.rev (Helper.expr_i args expr_witness i).values in
            let value = S.Expression.function_ pos func args_i in
            R { witness = expr_witness; value })
      in
      result

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

    (** Convert each internal represention for the expression into its external
        representation *)
    let v : t -> t' * Report.t list =
     fun t ->
      let result =
        Array.map2 A.t t
          ~f:(fun (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' * Report.t list
       and type t' = result array = struct
    type expression = Expression.t' * Report.t list
    type t = result array
    type t' = result array

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

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

    let expression : expression -> t S.repr =
     fun expr report ->
      let expr, _report = expr in
      let results =
        Array.map2 A.t expr
          ~f:(fun
              (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 report in
                R { value; witness = instr_witness })
      in
      results

    let call : S.pos -> T.keywords -> expression list -> t S.repr =
     fun pos keyword args report ->
      (* The arguments are given like an array of array. Each expression is
         actually the list of each expression in the differents modules. *)

      (* Accumulate the results *)
      let report, args =
        List.fold_left_map args ~init:report ~f:(fun report (v, r) ->
            (r @ report, v))
      in

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

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

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

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

      let result =
        Array.init len ~f:(fun i ->
            let (E { module_ = (module S); instr_witness; expr'; _ }) =
              Array.get A.t i
            in
            let values =
              List.rev (Helper.args_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 report
                in
                R { witness = instr_witness; value })
      in

      result

    (* 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 in
      let report = ref report and len = Array.length A.t in

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

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

            let index_i =
              Option.map
                (fun expression ->
                  match get expr' (Array.get expression i) with
                  | None -> failwith "Does not match"
                  | Some value -> value)
                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 !report
                in

                R { value; witness = instr_witness })
      in

      result

    (** 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 in
      let t =
        List.map t ~f:(fun t ->
            let t = t report in
            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' f 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.args_i ts instr_witness i in
          let ts = List.rev ts.values 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_:(S.pos * t S.repr list) option ->
        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_ =
        match else_ with
        | None -> (report, None)
        | Some (pos, instructions) ->
            let instructions = Helper.map_args report instructions in
            (report, Some (pos, instructions))
      in
      let len = Array.length A.t in

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

            (* This function helps to build the expression in the clauses *)
            let f = Fun.id in

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

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

      result

    (** This code is almost a copy/paste from Expression.v but I did not found
        a way to factorize it. *)
    let v : t -> t' * Report.t list =
     fun t ->
      let result =
        Array.map2 A.t t
          ~f:(fun
              (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' * Report.t list) S.repr = struct
    type instruction = (Instruction.t' * Report.t list) S.repr
    type t = result array

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

      (* Extract the instructions and accumulate the result *)
      let instructions = Helper.map_args report instructions in

      let report, args =
        List.fold_left_map instructions ~init:report ~f:(fun report (v, r) ->
            (r @ report, v))
      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'; location_witness; _ }) =
              Array.get A.t i
            in

            let instructions = List.rev (Helper.args_i args instr' i).values in
            let value, re = A.Location.location pos instructions !report in
            report := re;
            R { value; witness = location_witness })
      in
      (result, !report)
  end
end