exception Defect of string
exception Not_implemented of string

module Machine = Merge_queue.Machine
module Signal = Merge_queue.Signal

module type INTERPRETER = sig
  type t
  type action
  type event

  val perform : t -> action -> emit:(event -> unit) -> unit
end

(* A failed system call, as the C library words its error, after the path or
   the call it failed on. *)
let system_error code call arg =
  Fmt.str "%s: %s (%s)"
    (if arg = "" then call else arg)
    (Unix.error_message code) call

let backend_error = function
  | Eio_unix.Unix_error (code, call, arg) -> Some (system_error code call arg)
  | _ -> None

let describe = function
  | Failure m | Invalid_argument m | Sys_error m -> m
  | Unix.Unix_error (code, call, arg) -> system_error code call arg
  | Eio.Io (err, _) as e -> (
      let backend =
        match err with
        | Eio.Exn.X b
        | Eio.Fs.E (Already_exists b | Not_found b | Permission_denied b) ->
            backend_error b
        | _ -> None
      in
      match backend with
      | Some s -> s
      | None ->
          String.concat " "
            (List.map String.trim
               (String.split_on_char '\n' (Fmt.str "%a" Eio.Exn.pp e))))
  | e -> Printexc.to_string e

let is_failure = function
  | Failure _ | Invalid_argument _ | Sys_error _ | Unix.Unix_error _ | Eio.Io _
    ->
      true
  | _ -> false

let failure e = if is_failure e then Some (describe e) else None

(* A frame of a backtrace that names a source file, as [FILE:LINE (FUNCTION)]. *)
let frame slot =
  match Printexc.Slot.location slot with
  | None -> None
  | Some { filename; line_number; _ } -> (
      match Printexc.Slot.name slot with
      | Some name -> Fmt.kstr Option.some "%s:%d (%s)" filename line_number name
      | None -> Fmt.kstr Option.some "%s:%d" filename line_number)

let with_backtrace m bt =
  let slots = Option.value ~default:[||] (Printexc.backtrace_slots bt) in
  match List.filter_map frame (Array.to_list slots) with
  | [] -> m
  | frames -> m ^ "; backtrace: " ^ String.concat ", " frames

(* [f ()], or [on_failure m] when it raises one of the shell's failures;
   any other exception, a cancellation or a defect, propagates. *)
let catching f ~on_failure =
  match f () with
  | v -> v
  | exception e when is_failure e -> on_failure (describe e)

module Make
    (M : Merge_queue.Machine.TOP)
    (I : INTERPRETER with type action = M.action and type event = M.event) =
struct
  (* A task by its id: its serial among the tasks of the run, its
     cancellation context once its fiber has made one, and whether the
     machine abandoned it, after which nothing it sends reaches the machine. *)
  type task = {
    serial : int;
    mutable cc : Eio.Cancel.t option;
    mutable abandoned : bool;
  }

  type loop = {
    events : M.event Eio.Stream.t;
    tasks : (Machine.Task.id, task) Hashtbl.t;
        (** The tasks running and not cancelled, by id. *)
    mutable running : int;
        (** The task fibers still running, cancelled ones included: each still
            owes the machine its answer. An abandoned one owes nothing. *)
    mutable stopped : (Machine.Task.id * task) list;
        (** The tasks cancelled whose fibers still run, newest first: those the
            machine may abandon. *)
    mutable serial : int;
    records : ((unit -> unit) -> unit) option;
        (** How the records of consecutive [Record] actions are written in one
            append, if they are. *)
  }

  let fail_defect fmt = Fmt.kstr (fun m -> raise (Defect m)) fmt

  (* Whether [e] is a defect when an interpreter raises it: it is none of the
     shell's failures, a cancellation or a path not implemented. *)
  let is_defect = function
    | Eio.Cancel.Cancelled _ | Not_implemented _ -> false
    | e -> not (is_failure e)

  (* The defect [e], raised performing [a], named after [a]. *)
  let named a = function
    | Defect m -> Fmt.kstr (fun m -> Defect m) "%a: %s" M.pp_action a m
    | e -> Fmt.kstr (fun m -> Defect m) "%a: %s" M.pp_action a (describe e)

  (* [f ()], which performs [a]; a defect it raises goes on named after [a],
     with its own backtrace, and any other exception as itself. *)
  let performing a f =
    match f () with
    | v -> v
    | exception e when is_defect e ->
        let bt = Printexc.get_raw_backtrace () in
        Printexc.raise_with_backtrace (named a e) bt

  (* An id names one task among the tasks running: a second task under the id
     of one not cancelled is a defect, never a task the machine could no
     longer cancel. A cancelled task leaves its id at once, and its fiber,
     once it ends, removes no task that took the id since. *)
  let start l ~sw h id (a : M.action) =
    if Hashtbl.mem l.tasks id then
      fail_defect "two tasks named %a: %a" Machine.Task.pp id M.pp_action a;
    l.serial <- l.serial + 1;
    let task = { serial = l.serial; cc = None; abandoned = false } in
    Hashtbl.replace l.tasks id task;
    l.running <- l.running + 1;
    Eio.Fiber.fork ~sw (fun () ->
        let push e = if not task.abandoned then Eio.Stream.add l.events e in
        (* A defect leaves the task running, so the loop does not take it for a
           machine waiting for nothing while the switch fails with it. *)
        let finished () =
          if not task.abandoned then l.running <- l.running - 1;
          l.stopped <-
            List.filter
              (fun (_, (k : task)) -> k.serial <> task.serial)
              l.stopped;
          match Hashtbl.find_opt l.tasks id with
          | Some k when k.serial = task.serial -> Hashtbl.remove l.tasks id
          | Some _ | None -> ()
        in
        match
          catching
            (fun () ->
              Eio.Cancel.sub (fun cc ->
                  task.cc <- Some cc;
                  Stopwatch.task (fun () ->
                      performing a (fun () -> I.perform h a ~emit:push)));
              finished ())
            ~on_failure:(fun m ->
              finished ();
              push (M.unanswered a (Machine.Failed m)))
        with
        | () -> ()
        | exception Eio.Cancel.Cancelled _ ->
            finished ();
            push (M.unanswered a Machine.Cancelled))

  (* A task stopped with a reason is cancelled with {!Machine.Task.Stopped},
     for the task to say why in what it leaves. *)
  let cancel ?(reason = Exit) l id =
    match Hashtbl.find_opt l.tasks id with
    | Some k ->
        Hashtbl.remove l.tasks id;
        l.stopped <- (id, k) :: l.stopped;
        Option.iter (fun cc -> Eio.Cancel.cancel cc reason) k.cc
    | None -> ()

  (* The newest task stopped under [id] owes the machine nothing more: what
     its fiber sends from now on is dropped. *)
  let abandon l id =
    match
      List.find_opt (fun (id', _) -> Machine.Task.equal id id') l.stopped
    with
    | Some (_, k) when not k.abandoned ->
        k.abandoned <- true;
        l.running <- l.running - 1;
        l.stopped <-
          List.filter (fun (_, (k' : task)) -> k'.serial <> k.serial) l.stopped
    | Some _ | None -> ()

  let is_record a =
    match M.task a with Machine.Task.Record -> true | _ -> false

  (* The [Record] actions at the head of [actions], and what follows them. *)
  let rec records_ahead acc = function
    | a :: rest when is_record a -> records_ahead (a :: acc) rest
    | rest -> (List.rev acc, rest)

  (* An action that fails stops the rest of its list: each later action is
     answered cancelled, unperformed. A cancellation is still made. The
     records of consecutive [Record] actions are one append: when it fails,
     none was written, the first is answered with what failed and the others
     cancelled. *)
  let perform_all l ~sw h actions =
    let push = Eio.Stream.add l.events in
    let now a =
      catching
        (fun () ->
          performing a (fun () -> I.perform h a ~emit:push);
          false)
        ~on_failure:(fun m ->
          push (M.unanswered a (Machine.Failed m));
          true)
    in
    let rec go ~failed = function
      | [] -> ()
      | a :: rest -> (
          match M.task a with
          | Machine.Task.Cancel id ->
              cancel l id;
              go ~failed rest
          | Machine.Task.Stop (id, why) ->
              cancel ~reason:(Machine.Task.Stopped why) l id;
              go ~failed rest
          | Machine.Task.Abandon id ->
              abandon l id;
              go ~failed rest
          | (Machine.Task.Now | Machine.Task.Record | Machine.Task.Start _)
            when failed ->
              push (M.unanswered a Machine.Cancelled);
              go ~failed rest
          | Machine.Task.Now -> go ~failed:(now a) rest
          | Machine.Task.Record -> (
              match l.records with
              | None -> go ~failed:(now a) rest
              | Some records ->
                  let group, rest = records_ahead [ a ] rest in
                  let failed =
                    catching
                      (fun () ->
                        records (fun () ->
                            List.iter
                              (fun a ->
                                performing a (fun () ->
                                    I.perform h a ~emit:push))
                              group);
                        false)
                      ~on_failure:(fun m ->
                        push (M.unanswered a (Machine.Failed m));
                        List.iter
                          (fun a -> push (M.unanswered a Machine.Cancelled))
                          (List.tl group);
                        true)
                  in
                  go ~failed rest)
          | Machine.Task.Start id ->
              start l ~sw h id a;
              go ~failed rest)
    in
    go ~failed:false actions

  let step l ~sw h = function
    | Ok (t, actions) ->
        perform_all l ~sw h actions;
        t
    | Error (Machine.Not_implemented { what; _ }) ->
        raise (Not_implemented what)
    | Error err -> fail_defect "%a" Machine.pp_error err

  (* The next event, or [None] once [deadline] passes first. An event taken in
     the round the deadline also passes is kept, never dropped: Fiber.first
     runs both to their end then, and the timer, due, fires on the next turn. *)
  let next_event l ~clock deadline =
    Stopwatch.waiting @@ fun () ->
    match deadline with
    | None -> Some (Eio.Stream.take l.events)
    | Some d ->
        Eio.Fiber.first
          ~combine:(fun a b -> match a with Some _ -> a | None -> b)
          (fun () -> Some (Eio.Stream.take l.events))
          (fun () ->
            Clock.sleep_until clock d;
            None)

  let due clock = function
    | Some d -> Ptime.compare d (Clock.now clock) <= 0
    | None -> false

  let rec drive l ~sw ~clock h t =
    match M.ending t with
    | Some _ -> t
    | None -> (
        let deadline = M.next_timeout t in
        if due clock deadline then
          drive l ~sw ~clock h (step l ~sw h (M.timer t ~now:(Clock.now clock)))
        else
          match Eio.Stream.take_nonblocking l.events with
          | Some e ->
              drive l ~sw ~clock h
                (step l ~sw h (M.incoming t ~now:(Clock.now clock) e))
          | None when l.running = 0 && Option.is_none deadline ->
              fail_defect "%a waits for nothing: no event, no task, no deadline"
                M.pp t
          | None -> (
              match next_event l ~clock deadline with
              | Some e ->
                  drive l ~sw ~clock h
                    (step l ~sw h (M.incoming t ~now:(Clock.now clock) e))
              | None ->
                  drive l ~sw ~clock h
                    (step l ~sw h (M.timer t ~now:(Clock.now clock)))))

  let watch l ~sw signal latch =
    Eio.Fiber.fork_daemon ~sw (fun () ->
        Eio_signal.await latch;
        Eio.Stream.add l.events (M.signal (signal ()));
        `Stop_daemon)

  let run ?spent ?records ~clock h t actions =
    Stopwatch.run clock @@ fun stopwatch ->
    let l =
      {
        events = Eio.Stream.create max_int;
        tasks = Hashtbl.create 8;
        running = 0;
        stopped = [];
        serial = 0;
        records;
      }
    in
    Eio_signal.install ~signals:[ Sys.sigint ] @@ fun int ->
    Eio_signal.install ~signals:[ Sys.sigterm ] @@ fun term ->
    Eio_signal.install ~signals:[ Sys.sighup ] @@ fun hup ->
    Eio.Switch.run @@ fun tasks_sw ->
    watch l ~sw:tasks_sw
      (fun () ->
        Signal.Interrupt { at_terminal = Console_eio.is_tty Unix.stdin })
      int;
    watch l ~sw:tasks_sw (fun () -> Signal.Terminate) term;
    watch l ~sw:tasks_sw (fun () -> Signal.Hangup) hup;
    perform_all l ~sw:tasks_sw h actions;
    let t = drive l ~sw:tasks_sw ~clock h t in
    (* The machine has ended: a task still running is stopped, saying so. *)
    let ended =
      Machine.Task.Stopped "the mq call that started it ended first"
    in
    Hashtbl.iter
      (fun _ k -> Option.iter (fun cc -> Eio.Cancel.cancel cc ended) k.cc)
      l.tasks;
    Option.iter (fun f -> f (Stopwatch.read stopwatch)) spent;
    t
end