open Rnt

module Error = struct
  open Concepts.Condition

  let closed () = condition "beam-closed" "The BEAM node is gone" empty
  let taken () = condition "beam-taken" "An evaluator is already registered as karuta:beam" empty

  let malformed () =
    condition "malformed-beam-program"
      "A beam program is not {Module, Function, Beam, #{Name => Domain}} of binaries" empty

  let not_a_program () =
    condition "not-a-beam-program" "The evaluator was given a program of another kind" empty

  let inputs () = condition "beam-program-inputs" "A beam program takes no inputs" empty

  let failed message =
    condition "beam-failed" "The BEAM node failed to answer a call"
      ("message" |=| Concepts.Value.String message)

  let unexpected () =
    condition "unexpected-beam-message" "The BEAM node said something out of turn" empty

  let unnamed () = condition "unnamed-beam-program" "A program was defined at an empty path" empty

  let program_taken path =
    condition "beam-program-taken" "A program is already defined at the path"
      ("path" |=| Concepts.Value.String (Kernel.Path.to_string path))

  let unknown_cursor id =
    condition "unknown-beam-cursor" "The BEAM node named a cursor it does not hold"
      ("cursor" |=| Concepts.Value.Integer id)
end

module type TRANSPORT = sig
  type t

  val send : t -> string -> (unit, Concepts.Condition.condition) result
  val receive : t -> (string, Concepts.Condition.condition) result
  val stop : t -> unit
end

module Port = struct
  type t = {input: in_channel; output: out_channel}

  let send {output; _} message =
    let header = Bytes.create 4 in
    Bytes.set_int32_be header 0 (Int32.of_int (String.length message));
    try
      output_bytes output header;
      output_string output message;
      Ok (flush output)
    with Sys_error _ -> Error (Error.closed ())

  let receive {input; _} =
    try
      let header = really_input_string input 4 in
      Ok (really_input_string input (Int32.to_int (String.get_int32_be header 0) land 0xffff_ffff))
    with End_of_file | Sys_error _ -> Error (Error.closed ())

  let start ?(erl = "erl") directories =
    let boot =
      "try [{module, _} = code:load_abs(filename:rootname(F)) || D <- init:get_plain_arguments(), \
       F <- filelib:wildcard(D ++ \"/*.beam\")], rnt:start() catch C:R:S -> \
       io:format(standard_error, \"~p:~p~n~p~n\", [C, R, S]), halt(1) end."
    in
    match
      Unix.open_process_args erl
        (Array.of_list ([erl; "-noshell"; "-noinput"; "-eval"; boot; "-extra"] @ directories))
    with
    | input, output -> Ok {input; output}
    | exception Unix.Unix_error _ -> Error (Error.closed ())

  let stop {input; output} = ignore (Unix.close_process (input, output))
end

let name = "karuta"

type program =
  { module_: string;
    function_: string;
    beam: string;
    heading: Protocols.Schematics.relation_description }

type Kernel.Program.kind += Beam of program

let decode source =
  let domain _ = function
    | Concepts.Value.String domain -> Some {Protocols.Schematics.domain; provenance= []}
    | Concepts.Value.Integer _ -> None
  in
  match Translator.read 4 source with
  | Some ([module_; function_; beam], attributes)
    when String.length module_ <= 0xffff && String.length function_ <= 0xffff ->
      let heading = BatMap.String.filter_map domain attributes in
      if BatMap.String.cardinal heading = BatMap.String.cardinal attributes then
        Ok (Beam {module_; function_; beam; heading})
      else Error (Error.malformed ())
  | _ -> Error (Error.malformed ())

(* Both ends speak over one pipe, and each serves the other's requests
   while it waits for its own answer, so a BEAM program RNT is reading
   can call back into RNT, and so on, as deep as the calls nest.

   RNT asks the node:
   - h                                 run every loaded rnt_init/0  -> o | e
   - i Id:64 M:16 Module F:16 Function B:32 Beam
                                       start Module:Function() from Beam
   - f Id:64 Limit:32                  the next batch of it         -> b | e
   - c Id:64                           forget it
   The node asks RNT:
   - D Path Evaluator:16 Source        define a program at Path     -> o | e
   - X Path                            invoke the program at Path   -> r Id:64 | e
   - E Evaluator:16 Source             invoke a program of no path  -> r Id:64 | e
   - N Id:64 Limit:32                  the next batch of its result -> b | e
   - C Id:64                           release that result
   A Path is Count:16 then each segment as Length:16 Segment, and a
   batch b is Exhausted:8 then each tuple as Size:32 term_to_binary. *)
module Make (T : TRANSPORT) = struct
  (* [root] is borrowed, not held: the evaluator lives under it, so
     holding it would keep both alive forever. [cursors] are the
     results the node invoked and has not finished reading. *)
  type node =
    { transport: T.t;
      root: Protocols.Handle.t;
      cursors: (int, Protocols.Handle.t * Protocols.Handle.t) Hashtbl.t;
      mutable next: int }

  let node transport root = {transport; root; cursors= Hashtbl.create 8; next= 0}

  module Request = struct
    type t =
      | Define of string list * string * string
      | Invoke of string list
      | Evaluate of string * string
      | Next of int * int
      | Close of int

    exception Short

    let of_message message =
      let length = String.length message in
      let at = ref 1 in
      let take n =
        if !at + n > length then raise Short;
        let p = !at in
        at := p + n;
        p
      in
      let u16 () = String.get_uint16_be message (take 2) in
      let u32 () = Int32.to_int (String.get_int32_be message (take 4)) land 0xffff_ffff in
      let u64 () = Int64.to_int (String.get_int64_be message (take 8)) in
      let sized () =
        let n = u16 () in
        String.sub message (take n) n
      in
      let path () = List.init (u16 ()) (fun _ -> sized ()) in
      let finished request = if !at = length then Some request else None in
      match message.[0] with
      | 'D' ->
          let path = path () in
          let evaluator = sized () in
          Some (Define (path, evaluator, String.sub message !at (length - !at)))
      | 'E' ->
          let evaluator = sized () in
          Some (Evaluate (evaluator, String.sub message !at (length - !at)))
      | 'X' -> finished (Invoke (path ()))
      | 'N' ->
          let id = u64 () in
          finished (Next (id, u32 ()))
      | 'C' -> finished (Close (u64 ()))
      | _ -> None
      | exception Short -> None
  end

  let batch_message {Protocols.Cursor.tuples; exhausted} =
    let buffer = Buffer.create 64 in
    Buffer.add_char buffer 'b';
    Buffer.add_char buffer (if exhausted then '\001' else '\000');
    BatFingerTree.iter
      (fun tuple ->
        let term = Translator.encode tuple in
        Buffer.add_int32_be buffer (Int32.of_int (String.length term));
        Buffer.add_string buffer term )
      tuples;
    Buffer.contents buffer

  let cursor_message id =
    let buffer = Buffer.create 9 in
    Buffer.add_char buffer 'r';
    Buffer.add_int64_be buffer (Int64.of_int id);
    Buffer.contents buffer

  let forget node id =
    Hashtbl.find_opt node.cursors id
    |> Option.iter (fun (relation, cursor) ->
        Hashtbl.remove node.cursors id;
        Protocols.Handle.release cursor;
        Protocols.Handle.release relation )

  let define node path evaluator source =
    let open Utilities.Result in
    match List.rev path with
    | [] -> Error (Error.unnamed ())
    | key :: directory -> (
        let* program = Kernel.Program.load node.root ~evaluator ~source in
        match
          Kernel.Path.update node.root
            (Kernel.Path.of_list (List.rev directory))
            key None (Some program)
        with
        | Ok true -> Ok "o"
        | Ok false ->
            Protocols.Handle.release program;
            Error (Error.program_taken (Kernel.Path.of_list path))
        | Error c -> Protocols.Handle.release program; Error c )

  (* Runs [program], which it releases, keeping a cursor on the result
     for the node to read. *)
  let invoke node program =
    let open Utilities.Result in
    let* relation =
      Protocols.Handle.releasing program (fun () ->
          let* executable = Protocols.Executable.require program in
          Protocols.Executable.invoke executable [] ~within:node.root )
    in
    match Protocols.Enumerable.cursor relation with
    | Error c -> Protocols.Handle.release relation; Error c
    | Ok cursor ->
        let id = node.next in
        node.next <- id + 1;
        Hashtbl.replace node.cursors id (relation, cursor);
        Ok (cursor_message id)

  let fetch node id limit =
    let open Utilities.Result in
    let* _, cursor =
      Hashtbl.find_opt node.cursors id |> Option.to_result ~none:(Error.unknown_cursor id)
    in
    let* scan = Protocols.Cursor.require cursor in
    let* batch = Protocols.Cursor.fetch scan limit in
    if batch.exhausted then forget node id;
    Ok (batch_message batch)

  let serve node request =
    let answer = function
      | Ok message -> T.send node.transport message
      | Error c -> T.send node.transport ("e" ^ Concepts.Condition.to_string_hum c)
    in
    match request with
    | Request.Define (path, evaluator, source) -> answer (define node path evaluator source)
    | Request.Invoke path ->
        answer (Result.bind (Kernel.Path.lookup node.root (Kernel.Path.of_list path)) (invoke node))
    | Request.Evaluate (evaluator, source) ->
        answer (Result.bind (Kernel.Program.load node.root ~evaluator ~source) (invoke node))
    | Request.Next (id, limit) -> answer (fetch node id limit)
    | Request.Close id -> Ok (forget node id)

  let batch_of message =
    let open Utilities.Result in
    let length = String.length message in
    let rec tuples at read =
      if at = length then Ok (List.rev read)
      else if at + 4 > length then Error (Error.unexpected ())
      else
        let size = Int32.to_int (String.get_int32_be message at) land 0xffff_ffff in
        if at + 4 + size > length then Error (Error.unexpected ())
        else
          let* tuple = Translator.decode (String.sub message (at + 4) size) in
          tuples (at + 4 + size) (tuple :: read)
    in
    if length < 2 then Error (Error.unexpected ())
    else
      let* read = tuples 2 [] in
      Ok {Protocols.Cursor.tuples= BatFingerTree.of_list read; exhausted= message.[1] <> '\000'}

  (* The node's answer to what RNT last asked, serving whatever the
     node asks in between. *)
  let rec await node =
    let open Utilities.Result in
    let* message = T.receive node.transport in
    match if message = "" then ' ' else message.[0] with
    | 'b' -> Result.map (fun batch -> `Batch batch) (batch_of message)
    | 'o' when message = "o" -> Ok `Done
    | 'e' -> Error (Error.failed (String.sub message 1 (String.length message - 1)))
    | _ -> (
      match Request.of_message message with
      | None -> Error (Error.unexpected ())
      | Some request ->
          let* () = serve node request in
          await node )

  (* Lets the node's rnt_init/0 functions define their programs. *)
  let handshake node =
    let open Utilities.Result in
    let* () = T.send node.transport "h" in
    let* reply = await node in
    match reply with `Done -> Ok () | `Batch _ -> Error (Error.unexpected ())

  module Poll = struct
    let invoke_message id {module_; function_; beam; _} =
      let buffer = Buffer.create (64 + String.length beam) in
      Buffer.add_char buffer 'i';
      Buffer.add_int64_be buffer (Int64.of_int id);
      List.iter
        (fun name ->
          Buffer.add_uint16_be buffer (String.length name);
          Buffer.add_string buffer name )
        [module_; function_];
      Buffer.add_int32_be buffer (Int32.of_int (String.length beam));
      Buffer.add_string buffer beam;
      Buffer.contents buffer

    let fetch_message id limit =
      let buffer = Buffer.create 13 in
      Buffer.add_char buffer 'f';
      Buffer.add_int64_be buffer (Int64.of_int id);
      Buffer.add_int32_be buffer (Int32.of_int (max 0 (min limit 0xffff_ffff)));
      Buffer.contents buffer

    let close_message id =
      let buffer = Buffer.create 9 in
      Buffer.add_char buffer 'c';
      Buffer.add_int64_be buffer (Int64.of_int id);
      Buffer.contents buffer

    class t node =
      object (self)
        inherit Kernel.Lifecycle.counted
        inherit Kernel.Identity.of_id
        method to_string = name

        method fetch limit =
          let open Utilities.Result in
          let* () = T.send node.transport (fetch_message (Oo.id self) limit) in
          let* reply = await node in
          match reply with `Batch batch -> Ok batch | `Done -> Error (Error.unexpected ())

        method destroy = ignore (T.send node.transport (close_message (Oo.id self)))
        method protocols : Protocols.Handle.protocol list = Protocols.[Cursor.make self]
      end

    let start node program =
      let poll = new t node in
      T.send node.transport (invoke_message (Oo.id poll) program)
      |> Result.map (fun () -> Protocols.Handle.make poll)
  end

  class evaluator node =
    object (self)
      inherit Kernel.Lifecycle.counted
      inherit Kernel.Identity.of_id
      method to_string = name

      method destroy =
        Hashtbl.to_seq_keys node.cursors |> List.of_seq |> List.iter (forget node);
        T.stop node.transport

      method parse source = decode source

      method eval :
          Kernel.Program.kind ->
          Protocols.Handle.t list ->
          within:Protocols.Handle.t ->
          (Protocols.Handle.t, Concepts.Condition.condition) result =
        fun kind inputs ~within:_ ->
          match kind, inputs with
          | Beam program, [] ->
              Ok
                (Kernel.EphemeralRelation.instantiate program.heading
                   ~modes:(fun () -> Ok [Concepts.Mode.mode [] Concepts.Cardinality.Countable])
                   (Kernel.EphemeralRelation.temporary program.heading)
                   (fun () -> Poll.start node program) )
          | Beam _, _ -> Error (Error.inputs ())
          | _ -> Error (Error.not_a_program ())

      method protocols : Protocols.Handle.protocol list = [Protocols.Evaluator.make self]
    end
end

module Node = Make (Port)

(* \karuta\program, where the node can define programs of its own. *)
let programs root =
  let open Utilities.Result in
  let* found = Kernel.Path.find root Kernel.Path.(name @/ "program" @/ this) in
  match found with
  | Some programs -> Ok (Protocols.Handle.release programs)
  | None -> (
      let programs = Kernel.Namespace.make () in
      match Kernel.Path.update root Kernel.Path.(name @/ this) "program" None (Some programs) with
      | Ok true -> Ok ()
      | Ok false -> Ok (Protocols.Handle.release programs)
      | Error c -> Protocols.Handle.release programs; Error c )

let register root directories =
  let open Utilities.Result in
  let* port = Port.start directories in
  let node = Node.node port root in
  let evaluator = Protocols.Handle.make (new Node.evaluator node) in
  let bound =
    let* () = Kernel.Initialization.mount root name in
    let* () = programs root in
    Kernel.Path.update root Kernel.Path.(name @/ "evaluator" @/ this) "beam" None (Some evaluator)
  in
  match bound with
  | Ok true -> Node.handshake node
  | Ok false ->
      Protocols.Handle.release evaluator;
      Error (Error.taken ())
  | Error c ->
      Protocols.Handle.release evaluator;
      Error c