open Std_internal
module Variant_and_record_intf = Variant_and_record_intf

module Helper (A : Variant_and_record_intf.S) (B : Variant_and_record_intf.S) = struct
  type map = { map : 'a. 'a A.t -> 'a B.t }

  let map_variant (type variant) { map } (variant : variant A.Variant.t) =
    let map_create = function
      | A.Tag.Args fct -> B.Tag_internal.Args fct
      | A.Tag.Const k -> B.Tag_internal.Const k
    in
    let map_tag tag =
      match tag with
      | A.Variant.Tag tag ->
        let label = A.Tag.label tag in
        let rep = map (A.Tag.traverse tag) in
        let arity = A.Tag.arity tag in
        let args_labels = A.Tag.args_labels tag in
        let index = A.Tag.index tag in
        let ocaml_repr = A.Tag.ocaml_repr tag in
        let tyid = A.Tag.tyid tag in
        let create = map_create (A.Tag.create tag) in
        B.Variant_internal.Tag
          (B.Tag.internal_use_only
             { B.Tag_internal.label
             ; rep
             ; arity
             ; args_labels
             ; index
             ; ocaml_repr
             ; tyid
             ; create
             })
    in
    let typename = A.Variant.typename_of_t variant in
    let polymorphic = A.Variant.is_polymorphic variant in
    let tags =
      Array.init (A.Variant.length variant) (fun index ->
        map_tag (A.Variant.tag variant index))
    in
    let value (a : variant) =
      match A.Variant.value variant a with
      | A.Variant.Value (atag, a) ->
        (fun (type args) (atag : (variant, args) A.Tag.t) (a : args) ->
          let (B.Variant_internal.Tag btag) = tags.(A.Tag.index atag) in
          (fun (type ex) (btag : (variant, ex) B.Tag.t) ->
            let Type_equal.T =
              Typename.same_witness_exn (A.Tag.tyid atag) (B.Tag.tyid btag)
            in
            let btag = (btag : (variant, args) B.Tag.t) in
            B.Variant_internal.Value (btag, a))
            btag)
          atag
          a
    in
    B.Variant.internal_use_only { B.Variant_internal.typename; tags; polymorphic; value }
  ;;

  let map_record (type record) { map } (record : record A.Record.t) =
    let map_field field =
      match field with
      | A.Record.Field field ->
        let label = A.Field.label field in
        let rep = map (A.Field.traverse field) in
        let index = A.Field.index field in
        let is_mutable = A.Field.is_mutable field in
        let tyid = A.Field.tyid field in
        let get = A.Field.get field in
        B.Record_internal.Field
          (B.Field.internal_use_only
             { B.Field_internal.label; rep; index; is_mutable; tyid; get })
    in
    let typename = A.Record.typename_of_t record in
    let has_double_array_tag = A.Record.has_double_array_tag record in
    let fields =
      Array.init (A.Record.length record) (fun index ->
        map_field (A.Record.field record index))
    in
    let create { B.Record_internal.get } =
      let get (type a) (afield : (_, a) A.Field.t) =
        match fields.(A.Field.index afield) with
        | B.Record_internal.Field bfield ->
          (fun (type ex) (bfield : (record, ex) B.Field.t) ->
            let Type_equal.T =
              Typename.same_witness_exn (A.Field.tyid afield) (B.Field.tyid bfield)
            in
            let bfield = (bfield : (record, a) B.Field.t) in
            get bfield)
            bfield
      in
      A.Record.create record { A.Record.get }
    in
    B.Record.internal_use_only
      { B.Record_internal.typename; fields; has_double_array_tag; create }
  ;;
end

module type Named = sig
  type 'a computation

  module Context : sig
    type t

    val create : unit -> t
  end

  type 'a t

  val init : Context.t -> 'a Typename.t -> 'a t
  val get_wip_computation : 'a t -> 'a computation
  val set_final_computation : 'a t -> 'a computation -> 'a computation
  val share : _ Typerep.t -> bool
end

module type Computation = sig
  type 'a t

  include Variant_and_record_intf.S with type 'a t := 'a t

  val int : int t
  val int32 : int32 t
  val int64 : int64 t
  val nativeint : nativeint t
  val char : char t
  val float : float t
  val string : string t
  val bytes : bytes t
  val bool : bool t
  val unit : unit t
  val option : 'a t -> 'a option t
  val list : 'a t -> 'a list t
  val array : 'a t -> 'a array t
  val lazy_t : 'a t -> 'a lazy_t t
  val ref_ : 'a t -> 'a ref t
  val function_ : 'a t -> 'b t -> ('a -> 'b) t
  val tuple2 : 'a t -> 'b t -> ('a * 'b) t
  val tuple3 : 'a t -> 'b t -> 'c t -> ('a * 'b * 'c) t
  val tuple4 : 'a t -> 'b t -> 'c t -> 'd t -> ('a * 'b * 'c * 'd) t
  val tuple5 : 'a t -> 'b t -> 'c t -> 'd t -> 'e t -> ('a * 'b * 'c * 'd * 'e) t
  val record : 'a Record.t -> 'a t
  val variant : 'a Variant.t -> 'a t

  module Named : Named with type 'a computation := 'a t
end

(* special functor application for computation as closure of the form [a -> b] *)
module Make_named_for_closure (X : sig
  type 'a input
  type 'a output
  type 'a t = 'a input -> 'a output
end) =
struct
  module Context = struct
    type t = unit

    let create = ignore
  end

  type 'a t =
    { runtime_dereference : 'a X.t
    ; runtime_reference : 'a X.t ref
    ; compiletime_dereference : 'a X.t option ref
    }

  exception Undefined of string

  let init () name =
    let path = Typename.Uid.name (Typename.uid name) in
    let r = ref (fun _ -> raise (Undefined path)) in
    { runtime_dereference = (fun input -> !r input)
    ; runtime_reference = r
    ; compiletime_dereference = ref None
    }
  ;;

  let get_wip_computation shared =
    match shared.compiletime_dereference.contents with
    | Some clos -> clos
    | None -> shared.runtime_dereference
  ;;

  let set_final_computation shared computation =
    let compiletime_dereference = shared.compiletime_dereference in
    match compiletime_dereference.contents with
    | Some _ -> assert false
    | None ->
      if Base.phys_equal shared.runtime_dereference computation then assert false;
      compiletime_dereference := Some computation;
      shared.runtime_reference := computation;
      computation
  ;;

  let share _ = true
end

module Ident = struct
  type t =
    { name : string
    ; implements : Typename.Uid.t -> bool
    }

  exception Broken_dependency of string

  let check_dependencies name required =
    match required with
    | [] -> fun _ -> ()
    | _ ->
      fun uid ->
        List.iter
          (fun { name = name'; implements } ->
            if not (implements uid)
            then (
              (* something is wrong with the set up, this is an error during the
                  initialization of the program, we rather fail with a human
                  readable output *)
              let message =
                Printf.sprintf
                  "Type_generic %S requires %S for uid %S\n"
                  name
                  name'
                  (Typename.Uid.name uid)
              in
              prerr_endline message;
              raise (Broken_dependency message)))
          required
  ;;
end

(* Extending an existing generic *)
module type Extending = sig
  type 'a t
  type 'a computation = 'a t

  val ident : Ident.t

  (* generic_ident * typename or info *)
  exception Not_implemented of string * string

  module type S = sig
    type t

    include Typerepable.S with type t := t

    val compute : t computation
  end

  module type S1 = sig
    type 'a t

    include Typerepable.S1 with type 'a t := 'a t

    val compute : 'a computation -> 'a t computation
  end

  module type S2 = sig
    type ('a, 'b) t

    include Typerepable.S2 with type ('a, 'b) t := ('a, 'b) t

    val compute : 'a computation -> 'b computation -> ('a, 'b) t computation
  end

  module type S3 = sig
    type ('a, 'b, 'c) t

    include Typerepable.S3 with type ('a, 'b, 'c) t := ('a, 'b, 'c) t

    val compute
      :  'a computation
      -> 'b computation
      -> 'c computation
      -> ('a, 'b, 'c) t computation
  end

  module type S4 = sig
    type ('a, 'b, 'c, 'd) t

    include Typerepable.S4 with type ('a, 'b, 'c, 'd) t := ('a, 'b, 'c, 'd) t

    val compute
      :  'a computation
      -> 'b computation
      -> 'c computation
      -> 'd computation
      -> ('a, 'b, 'c, 'd) t computation
  end

  module type S5 = sig
    type ('a, 'b, 'c, 'd, 'e) t

    include Typerepable.S5 with type ('a, 'b, 'c, 'd, 'e) t := ('a, 'b, 'c, 'd, 'e) t

    val compute
      :  'a computation
      -> 'b computation
      -> 'c computation
      -> 'd computation
      -> 'e computation
      -> ('a, 'b, 'c, 'd, 'e) t computation
  end

  val register0 : (module S) -> unit
  val register1 : (module S1) -> unit
  val register2 : (module S2) -> unit
  val register3 : (module S3) -> unit
  val register4 : (module S4) -> unit
  val register5 : (module S5) -> unit

  (* special less scary type when the type has no parameters *)
  val register : 'a Typerep.t -> 'a computation -> unit

  (*
     Essentially because we cannot talk about a variable of kind * -> k
     val register1 : 'a 't Typerep.t -> ('a computation -> 'a 't computation) -> unit
     ...
  *)
end

(* Implementing a new generic *)
module type S_implementation = sig
  include Extending

  (* raise using the current ident *)
  val raise_not_implemented : string -> 'a

  type implementation = { generic : 'a. 'a Typerep.t -> 'a computation }

  (*
     Standard case, find a extended_implementation, or look in the content
  *)
  val _using_extended_implementation
    :  implementation
    -> 'a Typerep.Named.t
    -> 'a Typerep.t lazy_t option
    -> 'a computation

  (*
     This function allows you more control on what you want to do
  *)
  val find_extended_implementation
    :  implementation
    -> 'a Typerep.Named.t
    -> 'a computation option
end

module type S = sig
  include Extending

  val of_typerep : 'a Typerep.t -> [ `generic of 'a computation ]

  module Computation : Computation with type 'a t = 'a t
end

module Make_S_implementation (X : sig
  type 'a t

  val name : string
  val required : Ident.t list
end) : S_implementation with type 'a t = 'a X.t = struct
  type 'a t = 'a X.t
  type 'a computation = 'a t

  include Type_generic_intf.M (struct
    type 'a t = 'a computation
  end)

  (* we do not use core since we are earlier in the dependencies graph *)
  module Uid_table = struct
    include Hashtbl.Make (Typename.Uid)

    let find table key =
      if Lazy.is_val table
      then (
        let table = Lazy.force table in
        try Some (find table key) with
        | Base.Not_found_s _ | Stdlib.Not_found -> None)
      else None
    ;;

    let check_dependencies = Ident.check_dependencies X.name X.required

    let replace table key value =
      check_dependencies key;
      replace (Lazy.force table) key value
    ;;

    let mem table key =
      if Lazy.is_val table
      then (
        let table = Lazy.force table in
        mem table key)
      else false
    ;;
  end

  let size = 256
  let table0 = lazy (Uid_table.create size)
  let table1 = lazy (Uid_table.create size)
  let table2 = lazy (Uid_table.create size)
  let table3 = lazy (Uid_table.create size)
  let table4 = lazy (Uid_table.create size)
  let table5 = lazy (Uid_table.create size)

  let is_registered uid =
    Uid_table.mem table0 uid
    || Uid_table.mem table1 uid
    || Uid_table.mem table2 uid
    || Uid_table.mem table3 uid
    || Uid_table.mem table4 uid
    || Uid_table.mem table5 uid
  ;;

  let ident = { Ident.name = X.name; implements = is_registered }

  module Find0 (T : Typerep.Named.T0) : sig
    val compute : unit -> T.named computation option
  end = struct
    let compute () =
      match Uid_table.find table0 (Typename.uid T.typename_of_t) with
      | None -> None
      | Some rep ->
        let module S = (val rep : S) in
        let witness = Typename.same_witness_exn S.typename_of_t T.typename_of_named in
        let module L =
          Type_equal.Lift (struct
            type 'a t = 'a computation
          end)
        in
        Some (Type_equal.conv (L.lift witness) S.compute)
    ;;
  end

  module Find1 (T : Typerep.Named.T1) : sig
    val compute : unit -> (T.a computation -> T.a T.named computation) option
  end = struct
    let compute () =
      match Uid_table.find table1 (Typename.uid T.typename_of_t) with
      | None -> None
      | Some rep ->
        let module S1 = (val rep : S1) in
        let module Conv =
          Typename.Same_witness_exn_1
            (S1)
            (struct
              type 'a t = 'a T.named

              let typename_of_t = T.typename_of_named
            end)
        in
        let module L =
          Type_equal.Lift (struct
            type 'a t = T.a computation -> 'a computation
          end)
        in
        Some (Type_equal.conv (L.lift Conv.(witness.eq)) S1.compute)
    ;;
  end

  module Find2 (T : Typerep.Named.T2) : sig
    val compute
      :  unit
      -> (T.a computation -> T.b computation -> (T.a, T.b) T.named computation) option
  end = struct
    let compute () =
      match Uid_table.find table2 (Typename.uid T.typename_of_t) with
      | None -> None
      | Some rep ->
        let module S2 = (val rep : S2) in
        let module Conv =
          Typename.Same_witness_exn_2
            (S2)
            (struct
              type ('a, 'b) t = ('a, 'b) T.named

              let typename_of_t = T.typename_of_named
            end)
        in
        let module L =
          Type_equal.Lift (struct
            type 'a t = T.a computation -> T.b computation -> 'a computation
          end)
        in
        Some (Type_equal.conv (L.lift Conv.(witness.eq)) S2.compute)
    ;;
  end

  module Find3 (T : Typerep.Named.T3) : sig
    val compute
      :  unit
      -> (T.a computation
          -> T.b computation
          -> T.c computation
          -> (T.a, T.b, T.c) T.named computation)
         option
  end = struct
    let compute () =
      match Uid_table.find table3 (Typename.uid T.typename_of_t) with
      | None -> None
      | Some rep ->
        let module S3 = (val rep : S3) in
        let module Conv =
          Typename.Same_witness_exn_3
            (S3)
            (struct
              type ('a, 'b, 'c) t = ('a, 'b, 'c) T.named

              let typename_of_t = T.typename_of_named
            end)
        in
        let module L =
          Type_equal.Lift (struct
            type 'a t =
              T.a computation -> T.b computation -> T.c computation -> 'a computation
          end)
        in
        Some (Type_equal.conv (L.lift Conv.(witness.eq)) S3.compute)
    ;;
  end

  module Find4 (T : Typerep.Named.T4) : sig
    val compute
      :  unit
      -> (T.a computation
          -> T.b computation
          -> T.c computation
          -> T.d computation
          -> (T.a, T.b, T.c, T.d) T.named computation)
         option
  end = struct
    let compute () =
      match Uid_table.find table4 (Typename.uid T.typename_of_t) with
      | None -> None
      | Some rep ->
        let module S4 = (val rep : S4) in
        let module Conv =
          Typename.Same_witness_exn_4
            (S4)
            (struct
              type ('a, 'b, 'c, 'd) t = ('a, 'b, 'c, 'd) T.named

              let typename_of_t = T.typename_of_named
            end)
        in
        let module L =
          Type_equal.Lift (struct
            type 'a t =
              T.a computation
              -> T.b computation
              -> T.c computation
              -> T.d computation
              -> 'a computation
          end)
        in
        Some (Type_equal.conv (L.lift Conv.(witness.eq)) S4.compute)
    ;;
  end

  module Find5 (T : Typerep.Named.T5) : sig
    val compute
      :  unit
      -> (T.a computation
          -> T.b computation
          -> T.c computation
          -> T.d computation
          -> T.e computation
          -> (T.a, T.b, T.c, T.d, T.e) T.named computation)
         option
  end = struct
    let compute () =
      match Uid_table.find table5 (Typename.uid T.typename_of_t) with
      | None -> None
      | Some rep ->
        let module S5 = (val rep : S5) in
        let module Conv =
          Typename.Same_witness_exn_5
            (S5)
            (struct
              type ('a, 'b, 'c, 'd, 'e) t = ('a, 'b, 'c, 'd, 'e) T.named

              let typename_of_t = T.typename_of_named
            end)
        in
        let module L =
          Type_equal.Lift (struct
            type 'a t =
              T.a computation
              -> T.b computation
              -> T.c computation
              -> T.d computation
              -> T.e computation
              -> 'a computation
          end)
        in
        Some (Type_equal.conv (L.lift Conv.(witness.eq)) S5.compute)
    ;;
  end

  let unit = Typename.static

  let register0 compute =
    let module S = (val compute : S) in
    let uid = Typename.uid S.typename_of_t in
    Uid_table.replace table0 uid compute
  ;;

  let register1 compute =
    let module S1 = (val compute : S1) in
    let uid = Typename.uid (S1.typename_of_t unit) in
    Uid_table.replace table1 uid compute
  ;;

  let register2 compute =
    let module S2 = (val compute : S2) in
    let uid = Typename.uid (S2.typename_of_t unit unit) in
    Uid_table.replace table2 uid compute
  ;;

  let register3 compute =
    let module S3 = (val compute : S3) in
    let uid = Typename.uid (S3.typename_of_t unit unit unit) in
    Uid_table.replace table3 uid compute
  ;;

  let register4 compute =
    let module S4 = (val compute : S4) in
    let uid = Typename.uid (S4.typename_of_t unit unit unit unit) in
    Uid_table.replace table4 uid compute
  ;;

  let register5 compute =
    let module S5 = (val compute : S5) in
    let uid = Typename.uid (S5.typename_of_t unit unit unit unit unit) in
    Uid_table.replace table5 uid compute
  ;;

  let register (type a) typerep_of_a compute =
    let module S = struct
      type t = a

      let typename_of_t = Typerep.typename_of_t typerep_of_a
      let typerep_of_t = typerep_of_a
      let compute = compute
    end
    in
    register0 (module S : S)
  ;;

  (* IMPLEMENTATION *)

  type implementation = { generic : 'a. 'a Typerep.t -> 'a computation }

  let find_extended_implementation (type a) aux = function
    | Typerep.Named.T0 rep ->
      let module T = (val rep : Typerep.Named.T0 with type t = a) in
      let module Custom = Find0 (T) in
      (match Custom.compute () with
       | Some custom ->
         let Type_equal.T = T.witness in
         Some (custom : a computation)
       | None -> None)
    | Typerep.Named.T1 rep ->
      let module T = (val rep : Typerep.Named.T1 with type t = a) in
      let module Custom = Find1 (T) in
      (match Custom.compute () with
       | Some custom ->
         let custom = (custom (aux.generic T.a) : T.a T.named computation) in
         let Type_equal.T = T.witness in
         Some (custom : a computation)
       | None -> None)
    | Typerep.Named.T2 rep ->
      let module T = (val rep : Typerep.Named.T2 with type t = a) in
      let module Custom = Find2 (T) in
      (match Custom.compute () with
       | Some custom ->
         let custom =
           (custom (aux.generic T.a) (aux.generic T.b) : (T.a, T.b) T.named computation)
         in
         let Type_equal.T = T.witness in
         Some (custom : a computation)
       | None -> None)
    | Typerep.Named.T3 rep ->
      let module T = (val rep : Typerep.Named.T3 with type t = a) in
      let module Custom = Find3 (T) in
      (match Custom.compute () with
       | Some custom ->
         let custom =
           (custom (aux.generic T.a) (aux.generic T.b) (aux.generic T.c)
             : (T.a, T.b, T.c) T.named computation)
         in
         let Type_equal.T = T.witness in
         Some (custom : a computation)
       | None -> None)
    | Typerep.Named.T4 rep ->
      let module T = (val rep : Typerep.Named.T4 with type t = a) in
      let module Custom = Find4 (T) in
      (match Custom.compute () with
       | Some custom ->
         let custom =
           (custom (aux.generic T.a) (aux.generic T.b) (aux.generic T.c) (aux.generic T.d)
             : (T.a, T.b, T.c, T.d) T.named computation)
         in
         let Type_equal.T = T.witness in
         Some (custom : a computation)
       | None -> None)
    | Typerep.Named.T5 rep ->
      let module T = (val rep : Typerep.Named.T5 with type t = a) in
      let module Custom = Find5 (T) in
      (match Custom.compute () with
       | Some custom ->
         let custom =
           (custom
              (aux.generic T.a)
              (aux.generic T.b)
              (aux.generic T.c)
              (aux.generic T.d)
              (aux.generic T.e)
             : (T.a, T.b, T.c, T.d, T.e) T.named computation)
         in
         let Type_equal.T = T.witness in
         Some (custom : a computation)
       | None -> None)
  ;;

  exception Not_implemented of string * string

  let raise_not_implemented string = raise (Not_implemented (X.name, string))

  let _using_extended_implementation aux rep content =
    match find_extended_implementation aux rep with
    | Some computation -> computation
    | None ->
      (match content with
       | Some (lazy content) -> aux.generic content
       | None ->
         let typename = Typerep.Named.typename_of_t rep in
         let name = Typename.Uid.name (Typename.uid typename) in
         raise_not_implemented name)
  ;;
end

module _ = Hashtbl.Make (Typename.Key)

module Make (X : sig
  type 'a t

  val name : string
  val required : Ident.t list

  include Computation with type 'a t := 'a t
end) =
struct
  module Computation = X
  include Make_S_implementation (X)

  module Memo = Typename.Table (struct
    type 'a t = 'a X.Named.t
  end)

  module Helper = Helper (Typerep) (Computation)

  let of_typerep rep =
    let context = X.Named.Context.create () in
    let memo_table = Memo.create 32 in
    let rec of_typerep : type a. a Typerep.t -> a t = function
      | Typerep.Int -> X.int
      | Typerep.Int32 -> X.int32
      | Typerep.Int64 -> X.int64
      | Typerep.Nativeint -> X.nativeint
      | Typerep.Char -> X.char
      | Typerep.Float -> X.float
      | Typerep.String -> X.string
      | Typerep.Bytes -> X.bytes
      | Typerep.Bool -> X.bool
      | Typerep.Unit -> X.unit
      | Typerep.Option rep -> X.option (of_typerep rep)
      | Typerep.List rep -> X.list (of_typerep rep)
      | Typerep.Array rep -> X.array (of_typerep rep)
      | Typerep.Lazy rep -> X.lazy_t (of_typerep rep)
      | Typerep.Ref rep -> X.ref_ (of_typerep rep)
      | Typerep.Function (dom, rng) -> X.function_ (of_typerep dom) (of_typerep rng)
      | Typerep.Tuple tuple ->
        (* do NOT write [X.tuple2 (of_typerep a) (of_typerep b)]
           because of_typerep can contain a side effect and [a] should be executed
           before [b] *)
        (match tuple with
         | Typerep.Tuple.T2 (a, b) ->
           let ra = of_typerep a in
           let rb = of_typerep b in
           X.tuple2 ra rb
         | Typerep.Tuple.T3 (a, b, c) ->
           let ra = of_typerep a in
           let rb = of_typerep b in
           let rc = of_typerep c in
           X.tuple3 ra rb rc
         | Typerep.Tuple.T4 (a, b, c, d) ->
           let ra = of_typerep a in
           let rb = of_typerep b in
           let rc = of_typerep c in
           let rd = of_typerep d in
           X.tuple4 ra rb rc rd
         | Typerep.Tuple.T5 (a, b, c, d, e) ->
           let ra = of_typerep a in
           let rb = of_typerep b in
           let rc = of_typerep c in
           let rd = of_typerep d in
           let re = of_typerep e in
           X.tuple5 ra rb rc rd re)
      | Typerep.Record record ->
        X.record (Helper.map_record { Helper.map = of_typerep } record)
      | Typerep.Variant variant ->
        X.variant (Helper.map_variant { Helper.map = of_typerep } variant)
      | Typerep.Named (named, content) ->
        let typename = Typerep.Named.typename_of_t named in
        (match Memo.find memo_table typename with
         | Some shared -> X.Named.get_wip_computation shared
         | None ->
           (match find_extended_implementation { generic = of_typerep } named with
            | Some computation -> computation
            | None ->
              (match content with
               | None ->
                 let name = Typename.Uid.name (Typename.uid typename) in
                 raise_not_implemented name
               | Some (lazy content) ->
                 if X.Named.share content
                 then (
                   let shared = X.Named.init context typename in
                   Memo.set memo_table typename shared;
                   let computation = of_typerep content in
                   X.Named.set_final_computation shared computation)
                 else of_typerep content)))
    in
    let computation = of_typerep rep in
    `generic computation
  ;;
end