(** [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 *)
  mutable n : int;               (* dlopen caches by path: never reuse one *)
}

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))

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) ^ ")"

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

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 "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 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;
  let child = Unix.create_process exe [| exe |] Unix.stdin Unix.stdout Unix.stderr in

  (* 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; 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
      match Unix.select [ ls ] [] [] 0.2 with
      | [], _, _ -> 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.unlink sock with Unix.Unix_error _ -> ()))
    accept_loop