(** [flan dev]: one long-lived session, the program it belongs to running
    beside it, and a socket an editor talks to.

    This is the piece between an editor and everything else. What it adds over
    [flan reload] is that the session *persists*: a [defvar] added by one
    evaluation is part of the program the next one is checked against, and the
    set of names the running process was built with is the one from the build
    this daemon actually made. A CLI that rebuilds its session from source each
    time cannot have either.

    It owns the build, which is what makes its layout rules mean anything: a
    session's struct layouts and global types describe the memory of a process
    only if it is the session that compiled it. So the daemon launches the
    program rather than attaching to one. *)

type t = {
  session : Session.t;
  child   : int;                 (* the running program *)
  agent   : string;              (* where it listens for modules *)
  dir     : string;              (* modules are built here, one per eval *)
  stdout  : Unix.file_descr;     (* the program's output, on its way to here *)
  out     : Buffer.t;            (* ...buffered until an editor asks for it *)
  mutable n : int;               (* dlopen caches by path: never reuse one *)
}

(* The program's stdout is a pipe into this process, so that an editor can see
   it. That makes draining it a *liveness* requirement and not a nicety: a pipe
   nobody reads fills at 64K and the next write blocks the program forever. So
   it is read from the accept loop's select, not only when someone asks. *)
let capacity = 256 * 1024

let drain t =
  let b = Bytes.create 8192 in
  let rec go () =
    match Unix.select [ t.stdout ] [] [] 0. with
    | [], _, _ -> ()
    | _ ->
      (match Unix.read t.stdout b 0 8192 with
       | 0 -> ()
       | n ->
         Buffer.add_subbytes t.out b 0 n;
         (* Bounded: a program that prints every frame must not grow this
            process without limit. The newest text is the useful end. *)
         if Buffer.length t.out > capacity then begin
           let keep = Buffer.sub t.out (Buffer.length t.out - capacity) capacity in
           Buffer.clear t.out;
           Buffer.add_string t.out keep
         end;
         go ()
       | exception Unix.Unix_error (Unix.EAGAIN, _, _) -> ()
       | exception Unix.Unix_error (Unix.EWOULDBLOCK, _, _) -> ()
       | exception Unix.Unix_error _ -> ())
  in
  go ()

let take t =
  drain t;
  let s = Buffer.contents t.out in
  Buffer.clear t.out;
  s

let await ?(ms = 5000) f =
  let rec go ms =
    if f () then true
    else if ms <= 0 then false
    else begin ignore (Unix.select [] [] [] 0.005); go (ms - 5) end
  in
  go ms

(* ── Delivery ──────────────────────────────────────────────────────── *)

(* The agent answers "ok" when it has queued a module, and anything else is a
   refusal with a reason. Reporting that back rather than swallowing it is what
   keeps a failed delivery from looking like a successful evaluation — the
   whole class of bug this socket makes possible. *)
let deliver t path =
  let s = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
  Fun.protect
    ~finally:(fun () -> try Unix.close s with Unix.Unix_error _ -> ())
    (fun () ->
       Unix.connect s (Unix.ADDR_UNIX t.agent);
       let msg = path ^ "\n" in
       ignore (Unix.write_substring s msg 0 (String.length msg));
       let b = Bytes.create 1024 in
       let buf = Buffer.create 64 in
       let rec drain () =
         match Unix.read s b 0 1024 with
         | 0 -> ()
         | n -> Buffer.add_subbytes buf b 0 n; drain ()
         | exception Unix.Unix_error _ -> ()
       in
       drain ();
       String.trim (Buffer.contents buf))

(* Read back the value of the last expression evaluated, with the counter that
   says whether it is a new one. The thunk runs on the game thread whenever the
   program next reaches a frame boundary, which is not a moment the daemon gets
   to know about, so this waits for the counter to move rather than assuming it
   has. *)
let result t =
  let s = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
  Fun.protect
    ~finally:(fun () -> try Unix.close s with Unix.Unix_error _ -> ())
    (fun () ->
       Unix.connect s (Unix.ADDR_UNIX t.agent);
       ignore (Unix.write_substring s "result\n" 0 7);
       let b = Bytes.create 4096 in
       let buf = Buffer.create 256 in
       let rec drain () =
         match Unix.read s b 0 4096 with
         | 0 -> ()
         | n -> Buffer.add_subbytes buf b 0 n; drain ()
         | exception Unix.Unix_error _ -> ()
       in
       drain ();
       let text = Buffer.contents buf in
       match String.index_opt text '\n' with
       | None -> None
       | Some i ->
         let header = String.sub text 0 i in
         let body = String.sub text (i + 1) (String.length text - i - 1) in
         (match String.split_on_char ' ' header with
          | [ g; _ ] ->
            (match Int64.of_string_opt g with
             | Some g -> Some (g, body)
             | None -> None)
          | _ -> None))

let alive t =
  match Unix.waitpid [ Unix.WNOHANG ] t.child with
  | 0, _ -> true
  | _ -> false
  | exception Unix.Unix_error _ -> false

(* ── Ops ───────────────────────────────────────────────────────────── *)

(* Every reply is a plist with a :status, so an editor can dispatch on one key
   and never has to guess whether a missing field means failure. *)
let ok fields =
  "(:status \"ok\"" ^ String.concat "" (List.map (fun f -> " " ^ f) fields) ^ ")"

(* Anything the program printed since the last reply rides along with this one.
   An editor that had to ask separately would miss the output an evaluation
   itself caused, which is the output anyone actually wants to see. *)
let with_output t reply =
  match take t with
  | "" -> reply
  | text ->
    let i = String.length reply - 1 in
    String.sub reply 0 i ^ " :output " ^ Wire.quote text ^ ")"

let error ?loc msg =
  "(:status \"error\" :message " ^ Wire.quote msg
  ^ (match loc with None -> "" | Some l -> " :loc " ^ Wire.quote l)
  ^ ")"

let eval t ~code ~origin =
  if not (alive t) then error "the program exited; restart flan dev"
  else
    match Session.eval ~origin t.session code with
    | c when not c.Session.installs ->
      (* Accepted into the session and nothing to send: a declaration the
         program already has, with no body and no new storage. Saying "ok" and
         shipping an empty module would report success for a change that cannot
         have taken effect. *)
      ok
        [ ":names " ^ Wire.strings c.Session.names; ":fns ()";
          ":note " ^ Wire.quote "nothing to install" ]
    | c ->
      t.n <- t.n + 1;
      let out = Filename.concat t.dir (Printf.sprintf "m%d.so" t.n) in
      (match Build.shared ~opts:{ Build.default with Build.dev = true }
               ~ir:c.Session.ir ~out () with
       | timing ->
         (match deliver t out with
          | "ok" ->
            ok
              [ ":names " ^ Wire.strings c.Session.names;
                ":fns " ^ Wire.strings c.Session.fns;
                Printf.sprintf ":ms %.1f"
                  (timing.Build.llc_ms +. timing.Build.link_ms) ]
          | reply -> error ("the program refused the module: " ^ reply)
          | exception Unix.Unix_error (e, _, _) ->
            error
              ("cannot reach the program on " ^ t.agent ^ ": "
               ^ Unix.error_message e))
       | exception Failure m -> error m)
    | exception Loc.Error (l, msg) -> error ~loc:(Loc.to_string l) msg

(* Redefining a name installs a body; evaluating an expression has no name to
   install into, so the module carries a thunk the agent runs once. The value
   comes back through the runtime rather than through this reply, because the
   frame boundary it runs at is the program's to choose. *)
let eval_expr t ~code ~origin =
  if not (alive t) then error "the program exited; restart flan dev"
  else
    match Session.eval_expr ~origin t.session code with
    | c ->
      let before = match result t with Some (g, _) -> g | None -> 0L in
      t.n <- t.n + 1;
      let out = Filename.concat t.dir (Printf.sprintf "e%d.so" t.n) in
      (match Build.shared ~opts:{ Build.default with Build.dev = true }
               ~ir:c.Session.ir ~out () with
       | _ ->
         (match deliver t out with
          | "ok" ->
            let rec wait ms =
              match result t with
              | Some (g, v) when Int64.compare g before > 0 -> Some v
              | _ when ms <= 0 -> None
              | _ ->
                ignore (Unix.select [] [] [] 0.005);
                if alive t then wait (ms - 5) else None
            in
            (match wait 5000 with
             | Some v -> ok [ ":value " ^ Wire.quote v ]
             | None ->
               error
                 "the program did not reach a frame boundary; is it calling \
                  (agent/poll)?")
          | reply -> error ("the program refused the module: " ^ reply)
          | exception Unix.Unix_error (e, _, _) ->
            error ("cannot reach the program: " ^ Unix.error_message e))
       | exception Failure m -> error m)
    | exception Loc.Error (l, msg) -> error ~loc:(Loc.to_string l) msg

let describe t =
  ok
    [ ":fns "
      ^ Wire.strings
          (List.map (fun (f : Tast.fn) -> f.Tast.name)
             t.session.Session.program.Tast.fns);
      ":globals "
      ^ Wire.strings
          (List.map (fun (g : Tast.global) -> g.Tast.gname)
             t.session.Session.program.Tast.globals);
      ":alive " ^ (if alive t then "t" else "nil") ]

let handle t req =
  match Wire.string_field req "op" with
  | Some "eval" ->
    (match Wire.string_field req "code" with
     | Some code ->
       let origin =
         match Wire.string_field req "file" with Some f -> f | None -> "<editor>"
       in
       eval t ~code ~origin
     | None -> error "eval needs :code")
  | Some "eval-expr" ->
    (match Wire.string_field req "code" with
     | Some code ->
       let origin =
         match Wire.string_field req "file" with Some f -> f | None -> "<editor>"
       in
       eval_expr t ~code ~origin
     | None -> error "eval-expr needs :code")
  | Some "describe" -> describe t
  | Some "close" -> ok []
  | Some op -> error ("unknown op: " ^ op)
  | None -> error "no :op"

(* ── The loop ──────────────────────────────────────────────────────── *)

(* One connection at a time. An editor is one client, evaluations are
   sequential by nature — each one is checked against the program the last one
   left behind — and a second concurrent evaluation would be racing for the
   same session anyway. *)
(* Returns whether the client asked to end the session. One editor per daemon,
   so [close] shuts the whole thing down rather than waiting for another
   connection nobody is going to make. *)
let serve t fd =
  let rec go () =
    match Wire.recv fd with
    | src ->
      let op, reply =
        match Wire.parse src with
        | req -> (Wire.string_field req "op", handle t req)
        | exception Loc.Error (_, m) -> (None, error ("bad request: " ^ m))
      in
      Wire.send fd (with_output t reply);
      if op = Some "close" then true else go ()
    | exception Wire.Closed -> false
    | exception Unix.Unix_error _ -> false
  in
  go ()

let start ~file ~sock =
  let t0 = Unix.gettimeofday () in
  let session, l = Session.create ~file in
  let dir =
    Filename.concat (Filename.get_temp_dir_name ())
      (Printf.sprintf "flan-dev-%d" (Unix.getpid ()))
  in
  (try Unix.mkdir dir 0o700 with Unix.Unix_error (Unix.EEXIST, _, _) -> ());
  let exe = Filename.concat dir "program" in
  ignore
    (Build.executable ~opts:{ Build.default with Build.dev = true }
       ~csrcs:l.Load.csrcs ~lflags:l.Load.lflags session.Session.host ~out:exe);
  let agent = Filename.concat dir "agent.sock" in

  (* The program's source names some socket path; the daemon is the one that
     knows where it wants to talk to it, so it overrides through the
     environment. Guessing instead would fail silently — everything compiles,
     the module is built, and nothing ever receives it. *)
  Unix.putenv "FLAN_AGENT_SOCKET" agent;
  (* Through a pipe, so the program's own output can reach an editor instead of
     only the terminal the daemon was started in. *)
  let rd, wr = Unix.pipe ~cloexec:false () in
  let child = Unix.create_process exe [| exe |] Unix.stdin wr Unix.stderr in
  Unix.close wr;
  Unix.set_nonblock rd;

  (* Wait for it to bind before accepting an evaluation. One that arrives first
     would fail for a reason that reads like a compiler bug. *)
  if not (await (fun () -> Sys.file_exists agent)) then begin
    (try Unix.kill child Sys.sigterm with Unix.Unix_error _ -> ());
    failwith
      ("the program never listened on " ^ agent
       ^ " — does it call (agent/start ...)?")
  end;

  let t =
    { session; child; agent; dir; stdout = rd; out = Buffer.create 4096; n = 0 }
  in
  (try Unix.unlink sock with Unix.Unix_error _ -> ());
  let ls = Unix.socket Unix.PF_UNIX Unix.SOCK_STREAM 0 in
  Unix.bind ls (Unix.ADDR_UNIX sock);
  Unix.listen ls 4;
  Printf.eprintf "flan dev: %s ready on %s (%.0fms)\n%!" file sock
    ((Unix.gettimeofday () -. t0) *. 1000.);
  (* [accept] would block past the program's own exit, so it is waited on with
     a timeout and the child checked each time round: a daemon whose program
     has finished has nothing left to do, and an editor waiting on it would
     wait forever. *)
  let rec accept_loop () =
    if alive t then
      (* The program's pipe is in the same select as the listening socket: it
         has to be drained whether or not an editor is asking for anything. *)
      match Unix.select [ ls; t.stdout ] [] [] 0.2 with
      | [], _, _ -> accept_loop ()
      | ready, _, _ when not (List.mem ls ready) -> drain t; accept_loop ()
      | _ ->
        (match Unix.accept ls with
         | fd, _ ->
           let closed = serve t fd in
           (try Unix.close fd with Unix.Unix_error _ -> ());
           if not closed then accept_loop ()
         | exception Unix.Unix_error (Unix.EINTR, _, _) -> accept_loop ())
      | exception Unix.Unix_error (Unix.EINTR, _, _) -> accept_loop ()
  in
  Fun.protect
    ~finally:(fun () ->
      (try Unix.kill child Sys.sigterm with Unix.Unix_error _ -> ());
      (try Unix.close ls with Unix.Unix_error _ -> ());
      (try Unix.close rd with Unix.Unix_error _ -> ());
      (try Unix.unlink sock with Unix.Unix_error _ -> ()))
    accept_loop