(* The agent, end to end (NEXT.md, dev loop step 3).

   test_reload.ml proves the primitive with a C harness driving it. This one
   proves the thing the dev loop actually is: a program that is running its own
   loop, a redefinition arriving over a socket while it runs, and the swap
   becoming visible at a point the program chose.

   The split the agent exists for is between two threads. [dlopen] relocates a
   module and takes the loader lock — milliseconds, unbounded — so it happens
   on the listener thread. [flan_reload_install] is one store per function, and
   it must not land while a redefined function is on the stack, so it happens
   on the game thread when it asks. Everything here is arranged to make that
   observable rather than to make it fast. *)

open Flan

let failures = ref 0
let fail fmt = Printf.ksprintf (fun s -> incr failures; print_endline ("FAIL  " ^ s)) fmt

let scratch = Filename.get_temp_dir_name ()
let tmp name = Filename.concat scratch ("flan-agent-" ^ name)

(* Poll for a condition rather than sleeping a fixed time: the program has to
   bind its socket before there is anything to connect to, and how long that
   takes is not ours to predict. *)
let rec await ?(ms = 3000) f =
  if f () then true
  else if ms <= 0 then false
  else begin
    ignore (Unix.select [] [] [] 0.005);
    await ~ms:(ms - 5) f
  end

(* The socket file appears at [bind], which is a moment before [listen], so a
   connect can lose that race and get ECONNREFUSED. Retry rather than sleep. *)
let rec connect ?(ms = 2000) path =
  let s = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
  match Unix.connect s (Unix.ADDR_UNIX path) with
  | () -> s
  | exception Unix.Unix_error (Unix.ECONNREFUSED, _, _) when ms > 0 ->
    Unix.close s;
    ignore (Unix.select [] [] [] 0.005);
    connect ~ms:(ms - 5) path

let send path line =
  let s = connect path in
  let msg = line ^ "\n" in
  ignore (Unix.write_substring s msg 0 (String.length msg));
  (* Read to EOF, not once: a reply arrives in several pieces, and closing
     after the first one is what made the agent take SIGPIPE. *)
  let buf = Bytes.create 512 in
  let b = Buffer.create 512 in
  let rec drain () =
    match Unix.read s buf 0 512 with
    | 0 -> ()
    | n -> Buffer.add_subbytes b buf 0 n; drain ()
    | exception Unix.Unix_error _ -> ()
  in
  drain ();
  Unix.close s;
  Buffer.contents b

let () =
  match Sys.command "command -v clang > /dev/null 2>&1 && command -v llc > /dev/null 2>&1" with
  | 0 ->
    (* The session is the program the process is about to be built from. Going
       through it rather than calling Emit directly is the point: it is what
       knows [tick] is a name the host has, so the module binds to its cell as
       a symbol instead of inventing a registry entry nobody publishes. *)
    let t, l = Session.create ~file:"programs/agent.flan" in

    (* A dev build, because that is what has cells to install into and exports
       them. The agent's own C and its -lpthread come from the package. *)
    let dev = { Build.default with Build.dev = true } in
    let exe = tmp "prog" in
    ignore
      (Build.executable ~opts:dev ~csrcs:l.Load.csrcs ~lflags:l.Load.lflags
         t.Session.host ~out:exe);

    (* Two evaluations from the one session, which is the daemon's loop and
       the thing no earlier test does. The first introduces a global the
       process was never built with; the second only reads it, and can only
       come back with 1007 if it found the storage the first one allocated
       rather than a fresh zeroed copy of it. *)
    let build_module src name =
      let c = Session.eval t src in
      let out = tmp name in
      ignore (Build.shared ~opts:dev ~ir:c.Session.ir ~out ());
      out
    in
    let so1 =
      build_module
        "(defvar acc i64) (defn tick [] i64 (set acc (+ acc 1000)) acc)"
        "tick1.so"
    in
    let so2 =
      build_module "(defn tick [] i64 (set acc (+ acc 7)) acc)" "tick2.so"
    in

    let sock = tmp "sock" in
    let out = tmp "out" in
    (try Sys.remove sock with Sys_error _ -> ());
    let fd = Unix.openfile out [ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC ] 0o600 in
    let pid =
      Unix.create_process exe [| exe; sock |] Unix.stdin fd Unix.stderr
    in
    Unix.close fd;

    if not (await (fun () -> Sys.file_exists sock)) then begin
      fail "the program never bound its socket";
      (try Unix.kill pid Sys.sigkill with Unix.Unix_error _ -> ())
    end else begin
      (* Refusing junk comes first, while the program is still running: the
         daemon is a separate process and can send anything, and a bad path
         must not take down the program it was sent to. Nothing is queued by
         it, so the program is still waiting afterwards. *)
      let reply = send sock "/nonexistent/nope.so" in
      if String.length reply < 4 || String.sub reply 0 4 <> "err " then
        fail "a bad path was not refused: %S" reply;

      (* "ok" means queued, not installed — the store happens on the other
         thread, at a time this one does not choose. *)
      let lines () =
        let text = In_channel.with_open_bin out In_channel.input_all in
        List.length (String.split_on_char '\n' text) - 1
      in
      let reply = send sock so1 in
      if reply <> "ok\n" then fail "agent replied %S, wanted \"ok\\n\"" reply;
      (* Wait for the program to have consumed the first module before sending
         the second. Both at once is a legitimate thing for the agent to do —
         one poll installs everything queued — but then only the last one is
         ever observed and the sequencing is not what was tested. *)
      if not (await (fun () -> lines () >= 2)) then
        fail "the first reload was never installed";
      let reply = send sock so2 in
      if reply <> "ok\n" then fail "agent replied %S, wanted \"ok\\n\"" reply;
      let _, status = Unix.waitpid [] pid in
      let text = In_channel.with_open_bin out In_channel.input_all in
      (* 1 from the original [tick]; 1000 from a body that did not exist when
         the program started, over a global that did not either; 1007 from a
         second body that only reads it. That last number is the whole point of
         doing this twice — a registry that handed out fresh storage per module
         would say 7. *)
      if status <> Unix.WEXITED 0 || text <> "1\n1000\n1007\n" then
        fail "agent reload\n        got: %S (%s)\n     wanted: %S" text
          (match status with
           | Unix.WEXITED c -> Printf.sprintf "exit %d" c
           | Unix.WSIGNALED c -> Printf.sprintf "signal %d" c
           | Unix.WSTOPPED c -> Printf.sprintf "stopped %d" c)
          "1\n1000\n1007\n"
    end;

    List.iter (fun f -> try Sys.remove f with Sys_error _ -> ())
      [ exe; so1; so2; sock; out ];
    if !failures = 0 then print_endline "agent: all tests passed"
    else begin
      Printf.printf "\n%d failure(s)\n" !failures;
      exit 1
    end
  | _ -> print_endline "agent: skipped (no clang or llc on PATH)"