(* 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);

    (* And the same program without [--dev], which has to *link*. Its [main]
       calls [agent/start], so the package is reached and its C comes with it
       even through [Reach.link] — and that C refers to the dev runtime, so
       leaving it out made this an undefined symbol at the link rather than a
       missing flag. Nothing is run: with no cells the agent refuses every
       module, and linking is the whole claim. *)
    (match
       let p, csrcs, lflags = Reach.link l t.Session.host in
       Build.executable ~opts:Build.default ~csrcs ~lflags p
         ~out:(tmp "prog-release")
     with
     | _ -> ()
     | exception Failure m -> fail "a release build of the agent: %s" m);

    (* 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;

    (* ── The break loop, spec-conditions.md §2 ──────────────────────── *)

    (* The claim is that an unhandled [error] stops rather than dying, and can
       be resumed into a restart chosen from outside. A program that died would
       exit 134 with no output; one that stopped and was never resumed would
       hang and be killed by the timeout. Only a resume produces both numbers,
       and they differ, so a loop that always took the same restart fails. *)
    let bsock = tmp "break.sock" and bout = tmp "break.out" in
    (try Sys.remove bsock with Sys_error _ -> ());
    let bt, bl = Session.create ~file:"programs/break.flan" in
    let bexe = tmp "break" in
    ignore
      (Build.executable ~opts:dev ~csrcs:bl.Load.csrcs ~lflags:bl.Load.lflags
         bt.Session.host ~out:bexe);
    let bfd = Unix.openfile bout [ Unix.O_WRONLY; Unix.O_CREAT; Unix.O_TRUNC ] 0o600 in
    let env =
      Array.append (Unix.environment ()) [| "FLAN_AGENT_SOCKET=" ^ bsock |]
    in
    let bpid =
      Unix.create_process_env bexe [| bexe |] env Unix.stdin bfd bfd
    in
    Unix.close bfd;
    if not (await (fun () -> Sys.file_exists bsock)) then
      fail "the broken program never listened"
    else begin
      (* It has to be *stopped* before the restarts are knowable: the listener
         refuses to walk a stack the game thread is still running on. *)
      let listed = ref "" in
      if not
           (await (fun () ->
              listed := send bsock "restarts";
              !listed <> "" && not (String.length !listed >= 3
                                    && String.sub !listed 0 3 = "err")))
      then fail "the program never reached the break loop: %S" !listed
      else begin
        (* Innermost first, and both on offer. *)
        if !listed <> "retry\nuse-placeholder\n.\n" then
          fail "restarts on offer\n        got: %S\n     wanted: %S" !listed
            "retry\nuse-placeholder\n.\n";
        (* A name nothing offers is refused *here*, before the reply. Answering
           ok and discovering it on the game thread would report success for
           something that cannot happen. *)
        let bad = send bsock "restart nonesuch" in
        if not (String.length bad >= 3 && String.sub bad 0 3 = "err") then
          fail "a restart nobody offers was accepted: %S" bad;
        ignore (send bsock "restart retry");
        (* Wait for the *result* of that choice before making the next one.
           Asking whether it is stopped is not enough: it is still stopped in
           the first break until the resume lands, and a second choice sent
           then would be taken by the first one — which passes the listing
           check and then hangs, because the second break never gets an
           answer. The printed 7 is the only proof the first resume happened. *)
        let printed () =
          let t = In_channel.with_open_bin bout In_channel.input_all in
          List.exists (String.equal "7") (String.split_on_char '\n' t)
        in
        if not (await printed) then
          fail "the first restart never produced its value"
        else if not (await (fun () -> send bsock "restarts"
                                      = "retry\nuse-placeholder\n.\n"))
        then fail "the program never stopped a second time"
        else ignore (send bsock "restart use-placeholder")
      end
    end;
    let bstatus = ref (Unix.WEXITED 0) in
    let reaped =
      await ~ms:5000 (fun () ->
        match Unix.waitpid [ Unix.WNOHANG ] bpid with
        | 0, _ -> false
        | _, s -> bstatus := s; true)
    in
    if not reaped then begin
      (try Unix.kill bpid Sys.sigkill with Unix.Unix_error _ -> ());
      fail "the program never resumed out of the break loop"
    end
    else begin
      let text = In_channel.with_open_bin bout In_channel.input_all in
      let want = "7\n-1\n" in
      let got =
        String.concat "\n"
          (List.filter
             (fun l -> l <> "" && not (String.length l >= 5 && String.sub l 0 5 = "flan:")
                       && not (String.length l >= 2 && String.sub l 0 2 = "  "))
             (String.split_on_char '\n' text))
      in
      if !bstatus <> Unix.WEXITED 0 || got ^ "\n" <> want then
        fail "break loop transcript\n        got: %S\n     wanted: %S" got want
    end;

    List.iter (fun f -> try Sys.remove f with Sys_error _ -> ())
      [ exe; so1; so2; sock; out; bsock; bout; bexe ];
    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)"