let schema = 5
let name = "managed-" ^ string_of_int schema

type outcome =
  | Exited of int
  | Signaled of int
  | Timed_out of string
  | Refused of string

type ending = Ended of float | Killed
type member = { pid : int; command : string; ending : ending }
type stop = { grace : float; members : member list }

type facts = {
  nice_refused : string option;
  silence : float option;
  stop : stop option;
}

let no_facts = { nice_refused = None; silence = None; stop = None }

type retention = { keeper : int; group : int; leases : int }

let equal_retention a b =
  Int.equal a.keeper b.keeper
  && Int.equal a.group b.group
  && Int.equal a.leases b.leases

type t =
  | Outcome of outcome * facts
  | Retained of retention
  | Drained of retention

module Codec = struct
  open Json.Codec

  let refuse fmt = Json.Error.fail_msgf Json.Meta.none fmt

  let validate_positive n =
    if n <= 0 then refuse "process id must be positive" else n

  let positive = map int ~enc:validate_positive ~dec:validate_positive

  let validate_nonnegative n =
    if n < 0 then refuse "lease count must not be negative" else n

  let nonnegative = map int ~enc:validate_nonnegative ~dec:validate_nonnegative

  let validate_silence n =
    if (not (Float.is_finite n)) || n < 0. then
      refuse "silence must be finite and non-negative"
    else n

  let finite = map number ~enc:validate_silence ~dec:validate_silence

  (* A member's ending: the seconds after TERM it was gone by, or [killed]. *)
  let ending =
    let validate n =
      if (not (Float.is_finite n)) || n < 0. then
        refuse "an ending must be finite and non-negative"
      else n
    in
    let ended = map number ~enc:validate ~dec:validate in
    Object.map (fun ended killed ->
        match (ended, killed) with
        | Some s, None -> Ended s
        | None, Some true -> Killed
        | _ -> refuse "a member ends after TERM or is killed, one of them")
    |> Object.opt_member "ended" ended ~enc:(function
      | Ended s -> Some s
      | Killed -> None)
    |> Object.opt_member "killed" bool ~enc:(function
      | Killed -> Some true
      | Ended _ -> None)
    |> Object.seal

  let member =
    Object.map (fun pid command ending -> { pid; command; ending })
    |> Object.member "pid" positive ~enc:(fun m -> m.pid)
    |> Object.member "command" string ~enc:(fun m -> m.command)
    |> Object.member "ending" ending ~enc:(fun m -> m.ending)
    |> Object.seal

  let stop =
    Object.map (fun grace members -> { grace; members })
    |> Object.member "grace" finite ~enc:(fun s -> s.grace)
    |> Object.member "members" (list member) ~enc:(fun s -> s.members)
    |> Object.seal

  let retention =
    Object.map (fun keeper group leases -> { keeper; group; leases })
    |> Object.member "keeper" positive ~enc:(fun r -> r.keeper)
    |> Object.member "group" positive ~enc:(fun r -> r.group)
    |> Object.member "leases" nonnegative ~enc:(fun r -> r.leases)
    |> Object.seal

  let outcome_value ctor project codec =
    Object.map (fun value nice_refused silence stop ->
        (ctor value, { nice_refused; silence; stop }))
    |> Object.member "value" codec ~enc:(fun (outcome, _) -> project outcome)
    |> Object.opt_member "nice_refused" string ~enc:(fun (_, facts) ->
        facts.nice_refused)
    |> Object.opt_member "silence" finite ~enc:(fun (_, facts) -> facts.silence)
    |> Object.opt_member "stop" stop ~enc:(fun (_, facts) -> facts.stop)
    |> Object.seal

  let exited =
    Object.Case.map "exited"
      (outcome_value
         (fun n -> Exited n)
         (function Exited n -> n | _ -> invalid_arg "not exited")
         int)
      ~dec:(fun (o, f) -> Outcome (o, f))

  let signaled =
    Object.Case.map "signaled"
      (outcome_value
         (fun n -> Signaled n)
         (function Signaled n -> n | _ -> invalid_arg "not signaled")
         int)
      ~dec:(fun (o, f) -> Outcome (o, f))

  let timeout =
    Object.Case.map "timeout"
      (outcome_value
         (fun n -> Timed_out n)
         (function Timed_out n -> n | _ -> invalid_arg "not timeout")
         string)
      ~dec:(fun (o, f) -> Outcome (o, f))

  let refused =
    Object.Case.map "refused"
      (outcome_value
         (fun n -> Refused n)
         (function Refused n -> n | _ -> invalid_arg "not refused")
         string)
      ~dec:(fun (o, f) -> Outcome (o, f))

  let retained = Object.Case.map "retained" retention ~dec:(fun r -> Retained r)
  let drained = Object.Case.map "drained" retention ~dec:(fun r -> Drained r)

  let enc_case = function
    | Outcome ((Exited _ as o), f) -> Object.Case.value exited (o, f)
    | Outcome ((Signaled _ as o), f) -> Object.Case.value signaled (o, f)
    | Outcome ((Timed_out _ as o), f) -> Object.Case.value timeout (o, f)
    | Outcome ((Refused _ as o), f) -> Object.Case.value refused (o, f)
    | Retained r -> Object.Case.value retained r
    | Drained r -> Object.Case.value drained r

  let json =
    Object.map (fun read message ->
        if read <> schema then
          refuse "managed protocol requires schema %d" schema
        else message)
    |> Object.member "schema" int ~enc:(fun _ -> schema)
    |> Object.case_member "kind" string ~enc:Fun.id ~enc_case
         Object.Case.
           [ v exited; v signaled; v timeout; v refused; v retained; v drained ]
    |> Object.error_unknown |> Object.seal
end

let of_string = Json.of_string Codec.json
let of_string_exn = Json.of_string_exn Codec.json

let to_string ?buf ?indent ?preserve value =
  Json.to_string ?buf ?indent ?preserve Codec.json value

let pp = Json.pp_value Codec.json