defmodule Belay do
  @moduledoc """
  A standalone, agent-native durable job engine on Postgres (or in-memory for
  tests). No Oban, no Ecto — Postgrex, Jason, and telemetry only.

  ## Start an instance

  Configure in application env (keyed by the instance name, like `Ecto.Repo`
  / `Phoenix.Endpoint`) and pull it in with `otp_app`:

      # config/config.exs
      config :my_app, MyBelay,
        queues: [default: 10, ai: [limit: 5, global_limit: 2, rate: [allowed: 60, period: 60]]],
        crons: [[name: "digest", expr: "0 8 * * *", worker: MyApp.Digest]]

      # config/runtime.exs
      config :my_app, MyBelay,
        storage: [adapter: :postgres, url: System.fetch_env!("DATABASE_URL")]

      # application.ex — inline opts override the app-env base
      children = [{Belay, otp_app: :my_app, name: MyBelay}]

  Or pass every option inline on the child spec and skip `otp_app` — both
  forms take the same keys.

  ## Define work

      defmodule MyApp.Agent do
        use Belay.Worker, queue: :ai, max_attempts: 10

        @impl Belay.Worker
        def run(ctx) do
          text = Belay.step(ctx, :fetch, fn -> fetch!(ctx.job.input["url"]) end)

          summary =
            Belay.step(ctx, :summarize, fn -> llm!(text) end,
              cost: [usd: 0.02, tokens: 1200]
            )

          %{"approved" => true} = Belay.await(ctx, :approval)

          {:ok, summary}
        end
      end

      Belay.insert(MyBelay, MyApp.Agent.new(%{"url" => url}, budget: [usd: 1.0]))

  Committed step results are memoized per job — retries replay past journaled
  work. Step bodies remain at-least-once until the journal write commits.
  Budgets fail the job with `:budget_exceeded` after a step's declared cost
  crosses the limit. Signals
  (`Belay.signal_job/4`) wake awaiting jobs instantly; `steer_job/3` injects
  guidance readable via `steering/1` at step boundaries.

  ## Instance options

  | Option | Default | Purpose |
  |---|---|---|
  | `:name` | `Belay` | instance name (atom); first argument to every API call |
  | `:otp_app` | `nil` | read the options below from `config :otp_app, name, ...`; inline opts override |
  | `:storage` | required | `[adapter: :postgres, url: ...]` or `[adapter: :memory]` |
  | `:queues` | `[]` | `queue: limit` or `queue: [limit:, global_limit:, rate:, partition:, manual:]` |
  | `:crons` | `[]` | `[[name:, expr:, worker:, input:, opts:]]`; merged with `Belay.Crons` rows |
  | `:notifiers` | `[:local]` | add `:postgres` for cross-fleet `pg_notify` wake-ups |
  | `:poll_interval` | `500` | idle polling ceiling (ms) |
  | `:busy_poll` | `25` | hot polling cadence (ms) |
  | `:lease_ttl` | `30_000` | running-job lease (ms); crash-orphan recovery window |
  | `:sweep_interval` | `5_000` | reclaim/retention cadence (ms) |
  | `:cron_interval` | `20_000` | cron tick (ms; slots dedup regardless) |
  | `:dynamic_sync` | `5_000` | runtime-queue reconciliation cadence (ms) |
  | `:shutdown_grace` | `15_000` | time running jobs get on shutdown (ms) |
  | `:retention` | 1d/7d/7d | per-terminal-state pruning: `[succeeded:, failed:, cancelled:]` seconds or `:infinity` |
  | `:signal_ttl` | `604_800` | seconds before undelivered signals are pruned |
  | `:encryption` | `nil` | `[key: {mod, fun, args}]` returning a 32-byte key |
  | `:clock` | system | `Belay.Clock` implementation (tests use `Belay.Clock.Sim`) |
  | `:node_id` | derived | stable identity for leases |
  """

  use Supervisor

  alias Belay.{Config, Ctx, Job, Runner}

  @typedoc "An instance name, as given in `:name`."
  @type instance :: atom()

  @typedoc "A job build produced by `YourWorker.new/2`."
  @type buildable :: {Belay.Worker, module(), map(), keyword()}

  # -- Supervision --------------------------------------------------------------

  @doc false
  @spec start_link(keyword()) :: Supervisor.on_start()
  def start_link(opts) do
    config = Config.new(opts)

    Supervisor.start_link(__MODULE__, config, name: Module.concat(config.name, "Supervisor"))
  end

  @doc false
  @spec child_spec(keyword()) :: Supervisor.child_spec()
  def child_spec(opts) do
    name = Keyword.get(opts, :name, Belay)

    %{id: {__MODULE__, name}, start: {__MODULE__, :start_link, [opts]}, type: :supervisor}
  end

  @impl Supervisor
  def init(config) do
    {storage_mod, storage_opts} = config.storage

    config = %{config | storage_ref: {storage_mod, storage_mod.ref(config.name)}}

    Config.put(config)

    children =
      [
        {Registry, keys: :unique, name: registry(config.name)},
        {Registry, keys: :duplicate, name: run_registry(config.name)},
        %{id: :pg, start: {:pg, :start_link, [pg_scope(config.name)]}},
        storage_mod.child_spec({config, storage_opts})
      ] ++
        Belay.Notifier.children(config) ++
        [
          {Task.Supervisor, name: task_sup(config.name)},
          {Belay.LeaseKeeper, config},
          {Belay.Sweeper, config},
          {Belay.CronScheduler, config},
          # Producers live under a DynamicSupervisor so runtime queue CRUD
          # (Belay.Queues) can start and stop them; QueueSync reconciles
          # static config + dynamic rows on every node.
          {DynamicSupervisor, name: producer_sup(config.name), strategy: :one_for_one},
          {Belay.QueueSync, config}
        ]

    # Lenient restart budget: transient storage failures must never take the
    # tree down — the loops themselves also rescue and back off (see Producer,
    # LeaseKeeper, Sweeper, CronScheduler).
    Supervisor.init(children, strategy: :one_for_one, max_restarts: 20, max_seconds: 30)
  end

  @doc false
  def registry(name), do: Module.concat(name, "Registry")

  @doc false
  def run_registry(name), do: Module.concat(name, "RunRegistry")

  @doc false
  def task_sup(name), do: Module.concat(name, "Tasks")

  @doc false
  def producer_sup(name), do: Module.concat(name, "Producers")

  # -- Inserting ----------------------------------------------------------------

  @doc """
  Insert a job built with `WorkerModule.new(input, opts)`.

  Options on `new/2`: `:queue`, `:priority`, `:max_attempts`, `:schedule_in`,
  `:partition_key`, `:meta`, `:budget` (`[usd: 5.0, tokens: 100_000]`), and
  `:unique` — either `"key"` (dedupe while incomplete) or
  `[key: k, within: seconds]` (dedupe per time window). A deduped insert
  returns the existing job flagged `duplicate?: true`.
  """
  @spec insert(instance(), buildable()) :: {:ok, Job.t()}
  def insert(name, {Belay.Worker, worker, input, opts}) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref
    now = Config.now(config)

    row =
      Job.new(
        worker,
        input,
        opts
        |> Keyword.put(:now, now)
        |> Keyword.put(:encryption_key, Config.encryption_key(config)),
        worker.__belay_defaults__()
      )

    case storage.insert_jobs(ref, [row], now) do
      {:ok, [job]} ->
        poke(config, job.queue)

        {:ok, job}

      {:ok, []} when row.unique_key != nil ->
        # Deduped by uniqueness — hand back the live holder of the key.
        {:ok, existing} = storage.get_by_unique_key(ref, row.unique_key)

        {:ok, %{existing | duplicate?: true}}
    end
  end

  @doc """
  Insert many jobs at once. Rows deduped by uniqueness or cron slots are
  silently skipped; only inserted jobs are returned.
  """
  @spec insert_all(instance(), [buildable()]) :: [Job.t()]
  def insert_all(name, buildables) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref
    now = Config.now(config)

    key = Config.encryption_key(config)

    rows =
      Enum.map(buildables, fn {Belay.Worker, worker, input, opts} ->
        Job.new(
          worker,
          input,
          opts |> Keyword.put(:now, now) |> Keyword.put(:encryption_key, key),
          worker.__belay_defaults__()
        )
      end)

    {:ok, jobs} = storage.insert_jobs(ref, rows, now)

    jobs |> Enum.map(& &1.queue) |> Enum.uniq() |> Enum.each(&poke(config, &1))

    jobs
  end

  # -- Introspection & control --------------------------------------------------

  @doc "Fetch a job by id."
  @spec get_job(instance(), integer()) :: {:ok, Job.t()} | {:error, :not_found}
  def get_job(name, id) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref

    case storage.get_job(ref, id) do
      {:ok, job} -> {:ok, job}
      :error -> {:error, :not_found}
    end
  end

  @doc "Cancel a job: immediate for parked states, cooperative for running."
  @spec cancel_job(instance(), integer()) :: {:ok, :cancelled | :requested | :noop}
  def cancel_job(name, id) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref

    {:ok, result} = storage.request_cancel(ref, id, Config.now(config))

    {:ok, result.status}
  end

  @doc "List a job's recorded steps with costs."
  @spec steps(instance(), integer()) :: {:ok, [map()]}
  def steps(name, id) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref

    storage.list_steps(ref, id)
  end

  @doc "Per-queue, per-state job counts."
  @spec stats(instance()) :: %{
          optional(String.t()) => %{optional(String.t()) => non_neg_integer()}
        }
  def stats(name) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref
    {:ok, rows} = storage.queue_stats(ref)

    Enum.reduce(rows, %{}, fn %{queue: q, state: s, count: c}, acc ->
      Map.update(acc, q, %{s => c}, &Map.put(&1, s, c))
    end)
  end

  @doc """
  List jobs newest-first. Filters: `:queue`, `:state`, `:worker`,
  `:workflow_id`, `:parent_id`, `:before_id` (pagination cursor), `:limit`
  (default 50).
  """
  @spec list_jobs(instance(), Enumerable.t()) :: [Job.t()]
  def list_jobs(name, filters \\ %{}) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref
    {:ok, jobs} = storage.list_jobs(ref, Map.new(filters))

    jobs
  end

  @doc "Resurrect a failed or cancelled job (grants one more attempt if exhausted)."
  @spec retry_job(instance(), integer()) :: {:ok, Job.t()} | {:error, :not_retryable | :not_found}
  def retry_job(name, id) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref

    with {:ok, job} <- storage.retry(ref, id, Config.now(config)) do
      poke(config, job.queue)

      {:ok, job}
    end
  end

  @doc "Stop a queue's local producer from claiming (running jobs finish)."
  @spec pause_queue(instance(), atom() | String.t()) :: :ok | {:error, :no_producer}
  def pause_queue(name, queue), do: call_producer(name, queue, :belay_pause)

  @doc "Resume a paused queue."
  @spec resume_queue(instance(), atom() | String.t()) :: :ok | {:error, :no_producer}
  def resume_queue(name, queue), do: call_producer(name, queue, :belay_resume)

  defp call_producer(name, queue, message) do
    case Registry.lookup(registry(name), {:producer, to_string(queue)}) do
      [{pid, _}] ->
        GenServer.call(pid, message)

      _ ->
        {:error, :no_producer}
    end
  end

  @doc """
  Await a job's terminal result. Returns `{:ok, value}` for success,
  `{:error, {:job, state}}` for failure/cancellation, `{:error, :timeout}`.
  """
  @spec await_result(instance(), integer(), timeout()) ::
          {:ok, term()} | {:error, {:job, :failed | :cancelled} | :not_found | :timeout}
  def await_result(name, id, timeout \\ 5_000) do
    config = Config.fetch!(name)

    Registry.register(run_registry(name), {:result, id}, nil)

    check = fn ->
      case get_job(name, id) do
        {:ok, %Job{state: "succeeded"} = job} ->
          {:ok, Job.result(job)}

        {:ok, %Job{state: state}} when state in ~w(failed cancelled) ->
          {:error, {:job, String.to_atom(state)}}

        {:error, :not_found} ->
          {:error, :not_found}

        _ ->
          :pending
      end
    end

    result = await_loop(check, config, id, monotonic_ms() + timeout)

    Registry.unregister(run_registry(name), {:result, id})

    result
  end

  # Wake instantly on a result notification; otherwise re-check on an
  # exponential 5ms → 200ms backoff, so short agent tasks return in
  # milliseconds even when the completion happened on an unconnected node.
  defp await_loop(check, config, id, deadline, interval \\ 5) do
    case check.() do
      :pending ->
        remaining = deadline - monotonic_ms()

        if remaining <= 0 do
          {:error, :timeout}
        else
          receive do
            {:belay_result, ^id, _info} -> await_loop(check, config, id, deadline)
          after
            min(remaining, interval) ->
              await_loop(check, config, id, deadline, min(interval * 2, 200))
          end
        end

      result ->
        result
    end
  end

  defp monotonic_ms, do: System.monotonic_time(:millisecond)

  # -- Signals ------------------------------------------------------------------

  @doc "Deliver a durable signal to a scope, waking any awaiting jobs."
  @spec signal(instance(), String.t() | atom(), String.t() | atom(), map()) :: :ok
  def signal(name, scope, signal_name, payload \\ %{}) when is_map(payload) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref

    {:ok, woken} =
      storage.put_signal(
        ref,
        to_string(scope),
        to_string(signal_name),
        payload,
        Config.now(config)
      )

    woken |> Enum.map(& &1.queue) |> Enum.uniq() |> Enum.each(&poke(config, &1))

    :ok
  end

  @doc "Deliver a durable signal scoped to one job."
  @spec signal_job(instance(), integer(), String.t() | atom(), map()) :: :ok
  def signal_job(name, job_id, signal_name, payload \\ %{}) do
    signal(name, "job:#{job_id}", signal_name, payload)
  end

  @doc "Inject steering guidance readable by the running job via `steering/1`."
  @spec steer_job(instance(), integer(), map()) :: :ok
  def steer_job(name, job_id, payload) when is_map(payload) do
    signal_job(name, job_id, "$steer", payload)
  end

  @doc "Delete a previously delivered signal from a scope."
  @spec clear_signal(instance(), String.t() | atom(), String.t() | atom()) :: :ok
  def clear_signal(name, scope, signal_name) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref

    storage.clear_signal(ref, to_string(scope), to_string(signal_name))
  end

  # -- In-job APIs (take the ctx) -----------------------------------------------

  @doc """
  Run and memoize `fun` under a per-job step name, with optional
  `cost: [usd:, tokens:]`.

  A committed result is replayed without re-running `fun`. The body is
  at-least-once until that journal write commits, so external effects should
  be idempotent across a crash-before-journal window.
  """
  @spec step(Ctx.t(), String.t() | atom(), (-> term()), keyword()) :: term()
  def step(%Ctx{} = ctx, step_name, fun, opts \\ []) when is_function(fun, 0) do
    Runner.step(ctx, step_name, fun, opts)
  end

  @doc """
  Wait for a signal. Returns the payload, or `{:error, :timeout}` after
  `:timeout` seconds. Parks the job (no process held) until signalled.
  """
  @spec await(Ctx.t(), String.t() | atom(), keyword()) :: map() | {:error, :timeout}
  def await(%Ctx{} = ctx, signal_name, opts \\ []) do
    Runner.await(ctx, signal_name, opts)
  end

  @doc """
  Durably sleep: parks the job (freeing its slot) and resumes after the
  target. The wake time is memoized under `name`, so replays skip past it.
  """
  @spec sleep(Ctx.t(), String.t() | atom(), non_neg_integer()) :: :ok
  def sleep(%Ctx{} = ctx, name, seconds), do: Runner.sleep(ctx, name, seconds)

  @doc "Read the latest steering payload, or nil."
  @spec steering(Ctx.t()) :: map() | nil
  def steering(%Ctx{} = ctx), do: Runner.steering(ctx)

  @doc """
  Spawn a child job from inside a running job. Memoized under `name`, so a
  crash after spawning cannot duplicate the child. Returns the child job id.
  """
  @spec spawn(Ctx.t(), String.t() | atom(), buildable()) :: integer()
  def spawn(%Ctx{} = ctx, name, buildable), do: Runner.spawn_child(ctx, name, buildable)

  @doc "Spawn many children as one memoized step. Returns their ids in order."
  @spec spawn_many(Ctx.t(), String.t() | atom(), [buildable()]) :: [integer()]
  def spawn_many(%Ctx{} = ctx, name, buildables), do: Runner.spawn_many(ctx, name, buildables)

  @doc """
  Park (at zero cost) until every spawned child reaches a terminal state.
  Returns the children ordered by id.
  """
  @spec await_children(Ctx.t()) :: [Job.t()]
  def await_children(%Ctx{} = ctx), do: Runner.await_children(ctx)

  @doc """
  Fan out one child per input and wait for all of them — the map phase of a
  map-reduce inside a single durable job. Returns child jobs in input order.

      results =
        Belay.map_children(ctx, :chunks, MyApp.Summarize, chunk_inputs)
        |> Enum.map(&Belay.Job.result/1)
  """
  @spec map_children(Ctx.t(), String.t() | atom(), module(), [map()], keyword()) :: [Job.t()]
  def map_children(%Ctx{} = ctx, name, worker, inputs, opts \\ []) do
    ids = Runner.spawn_many(ctx, name, Enum.map(inputs, &{Belay.Worker, worker, &1, opts}))
    by_id = ctx |> Runner.await_children() |> Map.new(&{&1.id, &1})

    Enum.map(ids, &Map.fetch!(by_id, &1))
  end

  @doc """
  Append to the job's durable event stream (progress, tokens, partial output).
  Subscribers receive `{:belay_event, job_id, seq, payload}` live;
  `events/3` replays from any offset.
  """
  @spec emit(Ctx.t(), map()) :: {:ok, integer() | :replayed | :no_registry}
  def emit(%Ctx{} = ctx, payload), do: Runner.emit(ctx, payload)

  @doc "Record actual resource usage against the queue's rate resource bucket."
  @spec debit(Ctx.t(), String.t(), integer()) :: :ok
  def debit(%Ctx{} = ctx, resource, units), do: Runner.debit(ctx, resource, units)

  @doc "Subscribe the calling process to a job's event stream."
  @spec subscribe_events(instance(), integer()) :: :ok
  def subscribe_events(name, job_id) do
    {:ok, _} = Registry.register(run_registry(name), {:events, job_id}, nil)

    :ok
  end

  @doc "Unsubscribe the calling process from a job's live event stream."
  @spec unsubscribe_events(instance(), integer()) :: :ok
  def unsubscribe_events(name, job_id) do
    Registry.unregister(run_registry(name), {:events, job_id})
  end

  @doc "Replay a job's event stream from an offset (0 for everything)."
  @spec events(instance(), integer(), non_neg_integer()) :: [map()]
  def events(name, job_id, after_seq \\ 0) do
    config = Config.fetch!(name)
    {storage, ref} = config.storage_ref
    {:ok, events} = storage.list_events(ref, job_id, after_seq)

    events
  end

  # -- Internal: producer poking ------------------------------------------------

  @doc false
  def poke(%Config{} = config, queue) do
    Belay.Notifier.broadcast_all(config, {:poke, to_string(queue)})
  end

  @doc false
  def pg_scope(name), do: Module.concat(name, "PG")
end