module M = Merge_queue
module T = M.Machine.Task
open World

type env = {
  world : World.t;
  pick : int -> int;
  mutable failures : int;
  mutable doubles : int;
  mutable next_pid : int;
  mutable widest : int;
  mutable killed_runs : (string * Ptime.t) list;
  mutable defected_runs : (string * Ptime.t) list;
}

type step = Timer | Event | Unblock | Finish of T.id

type stage = {
  run : string option;
  own : unit -> step option;
  due : unit -> step list;
  wakes : unit -> Ptime.t option list;
  enter : step -> (unit -> unit) option;
  check_end : unit -> unit;
  signal : M.Signal.t -> unit;
}

and t = {
  call : Call.t;
  mutable stage : stage;
  mutable ended : bool;
  mutable summary : Call.summary option;
  at_end : unit -> unit;
  after_kill : unit -> unit;
}

type next = {
  summary : Call.summary option;
  next : (t -> stage * (unit -> unit)) option;
}

let later t s =
  Option.value ~default:t
    (Ptime.add_span t (Ptime.Span.of_float_s s |> Option.get))

let is_due w = function Some d -> Ptime.compare d (now w) <= 0 | None -> false

(* A process that ends with its run's name in the lock's file is a run the
   next one finds interrupted: killed, or with the release of the file
   cancelled by a write that failed before it, its end among them. One that
   ends in a defect so, with no write failed, has not ended its run as a run
   must (lib/run.mli). *)
let finish env p (e : Call.ending) =
  let w = env.world in
  (match p.stage.run with
  | Some run when w.run_lock = Some p.call.pid && List.mem run w.run_file -> (
      match e with
      | Defect _ when not p.call.faulted ->
          env.defected_runs <- (run, now w) :: env.defected_runs
      | Defect _ | Killed _ | Ended _ | Unfinished ->
          env.killed_runs <- (run, now w) :: env.killed_runs)
  | Some _ | None -> ());
  p.call.ending <- e;
  p.call.ended_at <- Some (now w);
  p.ended <- true;
  exited w p.call.pid;
  p.at_end ();
  log w "p%d ends: %s" p.call.pid
    (match e with
    | Ended s -> s.line
    | Killed what -> "killed after " ^ what
    | Defect m -> "defect: " ^ m
    | Unfinished -> "unfinished");
  log w "  %a" pp_brief w

module Make
    (Mc : M.Machine.TOP)
    (I : Interpreter.S with type action = Mc.action and type event = Mc.event) =
struct
  (* [Seen f] follows the event an answer gave, and [f] runs as the machine
     takes that event. *)
  type item =
    | Plain of Mc.event
    | Sig of M.Signal.t * Mc.event
    | Seen of (unit -> unit)

  type task = {
    id : T.id;
    action : Mc.action;
    due : Ptime.t option;
    ready : (unit -> bool) option;
    answer : unit -> unit;
    seen : (unit -> unit) option;
  }

  type blocked = {
    until : Ptime.t option;
    answer : unit -> unit;
    seen : (unit -> unit) option;
    rest : Mc.action list;
  }

  type st = {
    mutable machine : Mc.t;
    queue : item Queue.t;
    mutable tasks : task list;
    mutable blocked : blocked option;
  }

  let stage ?run ?last env p (machine, actions) ~ended =
    let w = env.world and pid = p.call.pid in
    let st = { machine; queue = Queue.create (); tasks = []; blocked = None } in
    let push e = Queue.add (Plain e) st.queue in
    let answer f seen =
      f ();
      Option.iter (fun f -> Queue.add (Seen f) st.queue) seen
    in
    let due_of (d : Interpreter.deferred) =
      Option.map (later (now w)) d.after
    in
    let cancel id =
      match List.partition (fun (k : task) -> T.equal k.id id) st.tasks with
      | [ k ], rest ->
          st.tasks <- rest;
          log w "p%d cancels %a" pid T.pp id;
          push (Mc.unanswered k.action M.Machine.Cancelled)
      | _ -> ()
    in
    (* An action the fault budget lets fail, before its first write: what a
       full disk does. *)
    let inject a =
      if env.failures > 0 && I.failable a && env.pick 40 = 1 then (
        env.failures <- env.failures - 1;
        I.failed p.call a;
        p.call.faulted <- true;
        Some "no space left on device")
      else None
    in
    let perform_one a =
      log w "p%d -> %a" pid Mc.pp_action a;
      match inject a with
      | Some why ->
          log w "  p%d: the action fails: %s" pid why;
          `Failed why
      | None -> (
          match I.perform w p.call a ~emit:push with
          | None -> `Done
          | Some d -> `Deferred d
          | exception Failure m -> `Failed m)
    in
    let rec perform_all ~failed = function
      | [] -> ()
      | a :: rest -> (
          match Mc.task a with
          | T.Cancel id | T.Stop (id, _) ->
              cancel id;
              perform_all ~failed rest
          | (T.Now | T.Record | T.Start _) when failed ->
              (* Cancelled by the write that failed before it: not
                 performed, as a failed action is not. *)
              I.failed p.call a;
              push (Mc.unanswered a M.Machine.Cancelled);
              perform_all ~failed rest
          (* A record performed alone: a fault injected into any record of a
             run of them covers the shell's, where the append of the run
             fails at its first. *)
          | T.Now | T.Record -> (
              match perform_one a with
              | `Done -> perform_all ~failed:false rest
              | `Failed why ->
                  push (Mc.unanswered a (M.Machine.Failed why));
                  perform_all ~failed:true rest
              | `Deferred d ->
                  st.blocked <-
                    Some
                      {
                        until = due_of d;
                        answer = d.answer;
                        seen = d.seen;
                        rest;
                      })
          | T.Start id -> (
              match perform_one a with
              | `Done -> perform_all ~failed rest
              | `Failed why ->
                  push (Mc.unanswered a (M.Machine.Failed why));
                  perform_all ~failed rest
              | `Deferred d ->
                  st.tasks <-
                    {
                      id;
                      action = a;
                      due = due_of d;
                      ready = d.ready;
                      answer = d.answer;
                      seen = d.seen;
                    }
                    :: st.tasks;
                  perform_all ~failed rest))
    in
    let check_end () =
      if not p.ended then
        match Mc.ending st.machine with
        | Some e -> (
            Option.iter (fun f -> f st.machine) last;
            let n = ended e in
            Option.iter (fun s -> p.summary <- Some s) n.summary;
            match n.next with
            | None ->
                finish env p
                  (match p.summary with
                  | Some s -> Ended s
                  | None -> Defect "the call ended with no outcome")
            | Some mk ->
                let stage, start = mk p in
                p.stage <- stage;
                start ();
                p.stage.check_end ())
        | None ->
            if
              Queue.is_empty st.queue && st.tasks = []
              && Option.is_none st.blocked
              && Option.is_none (Mc.next_timeout st.machine)
            then
              Fmt.kstr
                (fun m -> finish env p (Defect m))
                "%a waits for nothing" Mc.pp st.machine
    in
    let apply = function
      | Ok (m, actions) ->
          env.widest <- max env.widest (List.length actions);
          st.machine <- m;
          perform_all ~failed:false actions
      | Error e ->
          Fmt.kstr (fun m -> finish env p (Defect m)) "%a" M.Machine.pp_error e
    in
    let own () =
      match st.blocked with
      | Some b -> if is_due w b.until then Some Unblock else None
      | None ->
          if is_due w (Mc.next_timeout st.machine) then Some Timer
          else if not (Queue.is_empty st.queue) then Some Event
          else None
    in
    let due () =
      List.filter_map
        (fun (k : task) ->
          let ready = match k.ready with Some f -> f () | None -> false in
          if ready || is_due w k.due then Some (Finish k.id) else None)
        st.tasks
    in
    let wakes () =
      (match st.blocked with
        | Some b -> [ b.until ]
        | None -> [ Mc.next_timeout st.machine ])
      @ List.map (fun (k : task) -> k.due) st.tasks
    in
    let enter = function
      | Timer ->
          Some
            (fun () ->
              log w "p%d: timer" pid;
              apply (Mc.timer st.machine ~now:(now w)))
      | Event ->
          let rec take () =
            match Queue.take st.queue with
            | Plain e -> e
            | Sig (sg, e) ->
                p.call.signals <- sg :: p.call.signals;
                e
            | Seen f ->
                f ();
                take ()
          in
          let e = take () in
          let rec seen acc =
            match Queue.peek_opt st.queue with
            | Some (Seen f) ->
                ignore (Queue.take st.queue);
                seen (f :: acc)
            | Some (Plain _ | Sig _) | None -> List.rev acc
          in
          let seen = seen [] in
          (* An answer given twice to a run that holds the lock, as a defect
             of the shell would give it: what the run does with an answer it
             cannot take. *)
          if env.doubles > 0 && w.run_lock = Some pid && env.pick 60 = 1 then (
            env.doubles <- env.doubles - 1;
            p.call.doubled <- true;
            log w "  p%d: the answer is given twice" pid;
            push e);
          Some
            (fun () ->
              List.iter (fun f -> f ()) seen;
              log w "p%d <- %a" pid Mc.pp_event e;
              apply (Mc.incoming st.machine ~now:(now w) e))
      | Unblock -> (
          match st.blocked with
          | None -> None
          | Some b ->
              st.blocked <- None;
              Some
                (fun () ->
                  answer b.answer b.seen;
                  perform_all ~failed:false b.rest))
      | Finish id -> (
          match List.partition (fun (k : task) -> T.equal k.id id) st.tasks with
          | k :: _, rest ->
              st.tasks <- rest;
              Some (fun () -> answer k.answer k.seen)
          | [], _ -> None)
    in
    let signal sg = Queue.add (Sig (sg, Mc.signal sg)) st.queue in
    ( { run; own; due; wakes; enter; check_end; signal },
      fun () -> perform_all ~failed:false actions )
end

let idle =
  {
    run = None;
    own = (fun () -> None);
    due = (fun () -> []);
    wakes = (fun () -> []);
    enter = (fun _ -> None);
    check_end = ignore;
    signal = ignore;
  }

let v call ~at_end ~after_kill first =
  let p =
    { call; stage = idle; ended = false; summary = None; at_end; after_kill }
  in
  let stage, start = first p in
  p.stage <- stage;
  (p, start)

let call p = p.call
let running p = not p.ended
let own p = p.stage.own ()
let due p = p.stage.due ()
let wakes p = p.stage.wakes ()
let enter p s = p.stage.enter s
let check_end p = p.stage.check_end ()
let signal p s = p.stage.signal s
let after_kill p = p.after_kill ()