(* The beam evaluator. Its programs are functions on a BEAM node, each
   taking one relation per input and giving back a relation of the
   heading the program declares. rnt-node.lfe serves them. Tuples cross
   as Translator writes them, and a message is a byte for its kind
   followed by big-endian integers and binaries:

     RNT   i Cursor:32 M:16 Module F:16 Function Reader:32...  invoke
     RNT   f Cursor:32 Limit:32                                 fetch
     RNT   c Cursor:32                                          close
     node  n Reader:32 Limit:32                                 next
     both  b Exhausted:8 (Size:32 Tuple)...                     a batch
     both  e Message                                            a failure

   fetch and next are answered with a batch or a failure, invoke and
   close not at all, so a function that cannot be applied fails the
   first fetch. Calls nest: until its answer comes, whoever waits serves
   the calls that come instead. *)

open Rnt

module Error = struct
  open Concepts.Condition

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

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

  let closed () = condition "beam-closed" "The BEAM node is gone" empty

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

  let unheld_input () =
    condition "unheld-input" "An input could not be held for as long as the program needs it" empty
end

module type TRANSPORT = sig
  type t

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

(* A node RNT starts as its child and speaks to over the child's stdin
   and stdout, each message behind its length in four bytes, as a port
   opened with {packet, 4} frames it. *)
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)))
    with End_of_file | Sys_error _ -> Error (Error.closed ())

  (* The node compiles every LFE module in [directories], which must
     hold rnt.lfe and rnt-node.lfe, and serves until [stop] closes its
     stdin. If it cannot, it says why on stderr and halts. It is started
     with erl and not lfe, whose shell would read stdin, so LFE has to be
     where erl finds it, as ERL_LIBS says. *)
  let start ?(erl = "erl") directories =
    let boot =
      "try [begin {ok, Ms, _} = lfe_comp:file(F, [binary, return]), [{module, M} = \
       code:load_binary(M, F, B) || {ok, M, B, _} <- Ms] end || D <- init:get_plain_arguments(), F \
       <- filelib:wildcard(D ++ \"/*.lfe\")], 'rnt-node':start() catch C:R:S -> \
       io:format(standard_error, \"~p:~p~n~p~n\", [C, R, S]), halt(1) end."
    in
    let input, output =
      Unix.open_process_args erl
        (Array.of_list ([erl; "-noshell"; "-noinput"; "-eval"; boot; "-extra"] @ directories))
    in
    {input; output}

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

let subsystem = "karuta"
let language = "beam"
let name = subsystem ^ ":" ^ language

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

type Kernel.Program.kind += Beam of program

module Bencode = Concepts.Encoding.Bencode

let encode {module_; function_; heading} =
  Bencode.Dict
    [ "module", Bencode.String module_;
      "function", Bencode.String function_;
      ( "heading",
        Bencode.Dict
          ( BatMap.String.bindings heading
          |> List.map (fun (attribute, {Protocols.Schematics.domain; _}) ->
              attribute, Bencode.String domain ) ) ) ]
  |> Bencode.to_bytes
  |> Bytes.to_string

let decode source =
  let open Utilities.Result in
  let* program = Bencode.of_bytes (Bytes.of_string source) in
  let* module_ = Concepts.Encoding.Field.require_string "module" program in
  let* function_ = Concepts.Encoding.Field.require_string "function" program in
  let* heading = Bencode.field "heading" program |>> Bencode.as_dict in
  let* heading =
    heading
    |> List.map (fun (attribute, domain) ->
        Bencode.as_string domain
        |> Result.map (fun domain -> attribute, {Protocols.Schematics.domain; provenance= []}) )
    |> Utilities.List.sequence
  in
  Ok {module_; function_; heading= BatMap.String.of_list heading}

let int32 buffer n = Buffer.add_int32_be buffer (Int32.of_int n)

let invoke_message cursor {module_; function_; _} readers =
  let buffer = Buffer.create 64 in
  Buffer.add_char buffer 'i';
  int32 buffer cursor;
  List.iter
    (fun name ->
      Buffer.add_uint16_be buffer (String.length name);
      Buffer.add_string buffer name )
    [module_; function_];
  List.iter (int32 buffer) readers;
  Buffer.contents buffer

let fetch_message cursor limit =
  let buffer = Buffer.create 9 in
  Buffer.add_char buffer 'f'; int32 buffer cursor; int32 buffer limit; Buffer.contents buffer

let close_message cursor =
  let buffer = Buffer.create 5 in
  Buffer.add_char buffer 'c'; int32 buffer cursor; Buffer.contents buffer

let batch_message {Protocols.Cursor.tuples; exhausted} =
  let buffer = Buffer.create 256 in
  Buffer.add_char buffer 'b';
  Buffer.add_uint8 buffer (Bool.to_int exhausted);
  BatFingerTree.iter
    (fun tuple ->
      let term = Translator.encode tuple in
      int32 buffer (String.length term);
      Buffer.add_string buffer term )
    tuples;
  Buffer.contents buffer

let failure_message condition = "e" ^ Concepts.Condition.to_string_hum condition

let batch_of message =
  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
        Result.bind
          (Translator.decode (String.sub message (at + 4) size))
          (fun tuple -> tuples (at + 4 + size) (tuple :: read))
  in
  Result.map
    (fun read ->
      {Protocols.Cursor.tuples= BatFingerTree.of_list read; exhausted= message.[1] <> '\000'} )
    (tuples 2 [])

module Make (T : TRANSPORT) = struct
  (* The cursors over its inputs that the node reads through us, by the
     reader id it knows each by, and the last id given to a reader or a
     cursor. *)
  type node = {transport: T.t; readers: (int, Protocols.Handle.t) Hashtbl.t; mutable last: int}

  let fresh node =
    node.last <- node.last + 1;
    node.last

  let next node message =
    let open Utilities.Result in
    let reader = Int32.to_int (String.get_int32_be message 1) in
    let limit = Int32.to_int (String.get_int32_be message 5) in
    let* cursor =
      Hashtbl.find_opt node.readers reader |> Option.to_result ~none:(Error.unexpected ())
    in
    let* scan = Protocols.Cursor.require cursor in
    Protocols.Cursor.fetch scan limit

  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' when String.length message >= 2 -> batch_of message
    | 'e' -> Error (Error.failed (String.sub message 1 (String.length message - 1)))
    | 'n' when String.length message = 9 ->
        let answer =
          match next node message with
          | Ok batch -> batch_message batch
          | Error c -> failure_message c
        in
        let* () = T.send node.transport answer in
        await node
    | _ -> Error (Error.unexpected ())

  let forget node readers =
    List.iter
      (fun reader ->
        Hashtbl.find_opt node.readers reader |> Option.iter Protocols.Handle.release;
        Hashtbl.remove node.readers reader )
      readers

  class cursor node id readers =
    object (self)
      inherit Kernel.Lifecycle.counted
      inherit Kernel.Identity.of_id
      method to_string = "beam-cursor"

      method fetch limit =
        Result.bind (T.send node.transport (fetch_message id limit)) (fun () -> await node)

      method destroy =
        ignore (T.send node.transport (close_message id));
        forget node readers

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

  let open_readers node inputs =
    Utilities.List.foldm_left
      (fun readers input ->
        match Protocols.Enumerable.cursor input with
        | Ok cursor ->
            let reader = fresh node in
            Hashtbl.replace node.readers reader cursor;
            Ok (reader :: readers)
        | Error c -> forget node readers; Error c )
      [] inputs
    |> Result.map List.rev

  let invoke node program inputs =
    let open Utilities.Result in
    let* readers = open_readers node inputs in
    let id = fresh node in
    match T.send node.transport (invoke_message id program readers) with
    | Ok () -> Ok (Protocols.Handle.make (new cursor node id readers))
    | Error c -> forget node readers; Error c

  (* The relation reads its inputs each time it is enumerated, so it holds them. *)
  let hold inputs =
    let held = List.filter_map Protocols.Handle.copy inputs in
    if List.length held = List.length inputs then Ok held
    else begin
      List.iter Protocols.Handle.release held;
      Error (Error.unheld_input ())
    end

  (* TODO: A program should declare its modes as it declares its heading.
     Until then every program is taken to end. *)
  let generate node program inputs =
    let open Utilities.Result in
    let* inputs = hold inputs in
    Ok
      (Kernel.EphemeralRelation.instantiate ~inputs program.heading
         ~modes:(fun () -> Ok [Concepts.Mode.mode [] Concepts.Cardinality.Finite])
         (Kernel.EphemeralRelation.temporary program.heading)
         (fun () -> invoke node program inputs) )

  class evaluator node =
    object (self)
      inherit Kernel.Lifecycle.null
      inherit Kernel.Identity.of_id
      method to_string = "beam-evaluator"
      method parse source = decode source |> Result.map (fun program -> Beam program)

      method eval kind inputs ~within:_ =
        match kind with
        | Beam program -> generate node program inputs
        | _ -> Error (Error.not_a_program ())

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

  (* \karuta is the subsystem's root, and its programs name the evaluator karuta:beam. *)
  let register root transport =
    let open Utilities.Result in
    let node = {transport; readers= Hashtbl.create 8; last= 0} in
    let* _ = Kernel.Initialization.mount root subsystem in
    let* _ =
      Kernel.Path.update root
        Kernel.Path.(subsystem @/ "evaluator" @/ this)
        language None
        (Some (Protocols.Handle.make (new evaluator node)))
    in
    Ok ()
end

let program ?root ~module_ ~function_ heading =
  let program = {module_; function_; heading} in
  Kernel.Program.make ?root ~evaluator:name ~source:(encode program) (Beam program)