(* Much of this code is from Eio_posix.

   Copyright (C) 2021 Anil Madhavapeddy Copyright (C) 2022 Thomas Leonard

   Permission to use, copy, modify, and distribute this software for any purpose
   with or without fee is hereby granted, provided that the above copyright notice
   and this permission notice appear in all copies.

   THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
   REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
   FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
   INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
   LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
   OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
   PERFORMANCE OF THIS SOFTWARE. *)

open Eio.Std

module Process = struct
  type t = {
    pid : int;
    exit_status : Unix.process_status Promise.t;
    lock : Mutex.t;
  }
  (* When [lock] is unlocked, [exit_status] is resolved iff the process has been reaped. *)

  let exit_status t = t.exit_status
  let pid t = t.pid

  module Fork_action = Eio_unix.Private.Fork_action

  (* Read a (typically short) error message from a child process. *)
  let rec read_response fd =
    let buf = Bytes.create 256 in
    match Eio_posix.Low_level.read fd buf 0 (Bytes.length buf) with
    | 0 -> ""
    | n -> Bytes.sub_string buf 0 n ^ read_response fd

  let with_pipe fn =
    Switch.run @@ fun sw ->
    let r, w = Eio_posix.Low_level.pipe ~sw in
    fn r w

  let signal t signal =
    (* We need the lock here so that one domain can't signal the process exactly as another is reaping it. *)
    Mutex.lock t.lock;
    Fun.protect ~finally:(fun () -> Mutex.unlock t.lock) @@ fun () ->
    if not (Promise.is_resolved t.exit_status) then Unix.kill t.pid signal
  (* else process has been reaped and t.pid is invalid *)

  external eio_spawn :
    Unix.file_descr -> Eio_unix.Private.Fork_action.c_action list -> int
    = "caml_eio_posix_spawn"

  (* Wait for [pid] to exit and then resolve [exit_status] to its status. *)
  let reap t exit_status =
    Eio.Condition.loop_no_mutex Eio_unix.Process.sigchld (fun () ->
        Mutex.lock t.lock;
        match Unix.waitpid [ WNOHANG ] t.pid with
        | 0, _ ->
            Mutex.unlock t.lock;
            None (* Not ready; wait for next SIGCHLD *)
        | p, status ->
            assert (p = t.pid);
            Promise.resolve exit_status status;
            Mutex.unlock t.lock;
            Some ())

  let iter_switch ~f = function
    | Merry.Types.Async -> ()
    | Merry.Types.Switched sw -> f sw

  let spawn ?delay_reap ~mode actions =
    with_pipe @@ fun errors_r errors_w ->
    Eio_unix.Private.Fork_action.with_actions actions @@ fun c_actions ->
    iter_switch ~f:Switch.check mode;
    let exit_status, set_exit_status = Promise.create () in
    let t =
      let pid =
        Eio_unix.Fd.use_exn "errors-w" errors_w @@ fun errors_w ->
        Eio.Private.Trace.with_span "spawn" @@ fun () ->
        eio_spawn errors_w c_actions
      in
      Eio_unix.Fd.close errors_w;
      { pid; exit_status; lock = Mutex.create () }
    in
    let () =
      iter_switch
        ~f:(fun sw ->
          let hook =
            Switch.on_release_cancellable sw (fun () ->
                (* Kill process (if still running) *)
                signal t Sys.sigkill;
                (* The switch is being released, so either the daemon fiber got
             cancelled or it hasn't started yet (and never will start). *)
                if not (Promise.is_resolved t.exit_status) then
                  (* Do a (non-cancellable) waitpid here to reap the child. *)
                  reap t set_exit_status)
          in
          Fiber.fork_daemon ~sw (fun () ->
              Option.iter Eio.Promise.await delay_reap;
              reap t set_exit_status;
              Switch.remove_hook hook;
              `Stop_daemon))
        mode
    in
    (* Check for errors starting the process. *)
    match read_response errors_r with
    | "" -> t (* Success! Execing the child closed [errors_w] and we got EOF. *)
    | err -> failwith err
end

module Process_impl = struct
  type t = Process.t
  type tag = [ `Generic | `Unix ]

  let pid = Process.pid

  let await t =
    match Eio.Promise.await @@ Process.exit_status t with
    | Unix.WEXITED i -> `Exited i
    | Unix.WSIGNALED i -> `Signaled i
    | Unix.WSTOPPED _ -> assert false

  let signal = Process.signal
end

let process =
  let handler = Eio.Process.Pi.process (module Process_impl) in
  fun proc -> Eio.Resource.T (proc, handler)

let read_of_fd ~mode ~default ~to_close v =
  match (mode, v) with
  | Merry.Types.Async, _ | _, None -> default
  | Merry.Types.Switched sw, Some f -> (
      match Eio_unix.Resource.fd_opt f with
      | Some fd -> fd
      | None ->
          let r, w = Eio_unix.pipe sw in
          Fiber.fork ~sw (fun () ->
              Eio.Flow.copy f w;
              Eio.Flow.close w);
          let r = Eio_unix.Resource.fd r in
          to_close := r :: !to_close;
          r)

let write_of_fd ~mode ~default ~to_close v =
  match (mode, v) with
  | Merry.Types.Async, _ | _, None -> default
  | Merry.Types.Switched sw, Some f -> (
      match Eio_unix.Resource.fd_opt f with
      | Some fd -> fd
      | None ->
          let r, w = Eio_unix.pipe sw in
          Fiber.fork ~sw (fun () ->
              Eio.Flow.copy r f;
              Eio.Flow.close r);
          let w = Eio_unix.Resource.fd w in
          to_close := w :: !to_close;
          w)

let with_close_list fn =
  let to_close = ref [] in
  let close () = List.iter Eio_unix.Fd.close !to_close in
  match fn to_close with
  | x ->
      close ();
      x
  | exception ex ->
      let bt = Printexc.get_raw_backtrace () in
      close ();
      Printexc.raise_with_backtrace ex bt

let get_executable ~args = function
  | Some exe -> exe
  | None -> (
      match args with
      | [] -> invalid_arg "Arguments list is empty and no executable given!"
      | x :: _ -> x)

let get_env = function Some e -> e | None -> Unix.environment ()

external action_dups : unit -> Eio_unix.Private.Fork_action.fork_fn
  = "eio_unix_fork_dups"

let action_dups = action_dups ()

let rec with_fds mapping k =
  match mapping with
  | [] -> k []
  | (dst, src, _) :: xs ->
      Eio_unix.Fd.use_exn "inherit_fds" src @@ fun src ->
      with_fds xs @@ fun xs -> k ((dst, (Obj.magic src : int)) :: xs)

let inherit_fds m =
  let blocking =
    m
    |> List.filter_map (fun (dst, _, flags) ->
        match flags with
        | `Blocking -> Some (dst, true)
        | `Nonblocking -> Some (dst, false)
        | `Preserve_blocking -> None)
  in
  with_fds m @@ fun m ->
  (* TODO: investigate -- the plan from Eio seems to also invert the list of redirections.
     This is problematic for redirections, so we have copied the entire action here. *)
  let plan = Eio_unix__.Inherit_fds.plan m |> List.rev in
  Eio_unix.Private.Fork_action.
    { run = (fun k -> k (Obj.repr (action_dups, plan, blocking))) }

let spawn_unix () ?delay_reap ~mode ~fork_actions ?pgid ?uid ?gid ~env ~fds
    ~executable ~cwd args =
  let open Eio_posix in
  let actions =
    [
      inherit_fds fds;
      Low_level.Process.Fork_action.execve executable ~argv:(Array.of_list args)
        ~env;
    ]
  in
  let actions =
    match pgid with
    | None -> actions
    | Some pgid -> Low_level.Process.Fork_action.setpgid pgid :: actions
  in
  let actions =
    match uid with
    | None -> actions
    | Some uid -> Eio_unix.Private.Fork_action.setuid uid :: actions
  in
  let actions =
    match gid with
    | None -> actions
    | Some gid -> Eio_unix.Private.Fork_action.setgid gid :: actions
  in
  let actions = actions @ fork_actions in
  let with_actions cwd fn =
    let ((dir, path) : Eio.Fs.dir_ty Eio.Path.t) = cwd in
    match Eio_posix__.Fs.as_posix_dir dir with
    | None -> Fmt.invalid_arg "cwd is not an OS directory!"
    | Some dirfd ->
        Switch.run ~name:"spawn_unix" @@ fun launch_sw ->
        let cwd =
          Eio_posix__.Err.run
            (fun () ->
              let flags = Low_level.Open_flags.(rdonly + directory) in
              Low_level.openat ~sw:launch_sw ~mode:0 dirfd path flags)
            ()
        in
        fn (Low_level.Process.Fork_action.fchdir cwd :: actions)
  in
  with_actions cwd @@ fun actions ->
  process (Process.spawn ?delay_reap ~mode actions)

let fd_equal_int fd i =
  Eio_unix.Fd.use_exn "fd_equal_int" fd @@ fun ufd ->
  let ufd_int = (Obj.magic ufd : int) in
  Int.equal i ufd_int

let pp_redirections ppf (i, fd, _) = Fmt.pf ppf "(%i,%a)" i Eio_unix.Fd.pp fd

let run ~mode ?delay_reap _ ?stdin ?stdout ?stderr ?(fds = [])
    ?(fork_actions = []) ~pgid ~cwd ?env ?executable args =
  with_close_list @@ fun to_close ->
  let check_fd n = function
    | Merry.Types.Redirect (m, _, _) -> Int.equal n m
    | Merry.Types.Close fd -> fd_equal_int fd n
  in
  let fd_exists n = List.exists (check_fd n) fds in
  let std_fds =
    (if fd_exists 0 then []
     else
       [
         ( 0,
           read_of_fd ~mode stdin ~default:Eio_unix.Fd.stdin ~to_close,
           `Blocking );
       ])
    @ (if fd_exists 1 then []
       else begin
         [
           ( 1,
             write_of_fd ~mode stdout ~default:Eio_unix.Fd.stdout ~to_close,
             `Blocking );
         ]
       end)
    @
    if fd_exists 2 then []
    else
      [
        ( 2,
          write_of_fd ~mode stderr ~default:Eio_unix.Fd.stderr ~to_close,
          `Blocking );
      ]
  in
  let need_close, fds =
    List.fold_left
      (fun (cs, fs) -> function
        | Merry.Types.Redirect (a, b, c) -> (cs, (a, b, c) :: fs)
        | Close fd -> (fd :: cs, fs))
      ([], []) fds
    |> fun (cs, fs) -> (List.rev cs, List.rev fs)
  in
  List.iter Eio_unix.Fd.close need_close;
  let fds = std_fds @ fds in
  let executable = get_executable executable ~args in
  let env = get_env env in
  spawn_unix ?delay_reap ~mode ~fork_actions ~cwd ~pgid ~fds ~env ~executable ()
    args