lockstep

defmodule Lockstep.Generator do
  @moduledoc """
  Deterministic, seeded streams of operations for Jepsen-style
  workloads.

  A generator emits a sequence of opaque "ops" — tuples or atoms
  describing what a worker should do next. Workers pull from
  generators in a loop; the operations themselves are recorded via
  `Lockstep.History`.

  ## Built-in generators

    * `random/2` — uniform random over a pool of ops, seeded so
      iterations are reproducible.
    * `cycle/1` — fixed sequence repeated forever.
    * `weighted/2` — biased random sampling.
    * `mix/2` — interleave outputs from multiple generators
      round-robin.
    * `take/2` — limit any generator to N ops total.

  ## Quick example

      gen =
        Lockstep.Generator.random(
          [:read, {:write, 1}, {:write, 2}, {:cas, 1, 2}],
          seed: 42
        )
        |> Lockstep.Generator.take(50)

      Lockstep.Generator.each(gen, fn op ->
        Lockstep.History.op(history, op_kind(op), op_args(op), fn ->
          apply_op(reg, op)
        end)
      end)

  Generators are immutable; `next/1` returns `{op, new_gen}` or
  `:done`. They DON'T allocate processes — pass a generator into a
  closure and the worker advances it locally with no controller
  involvement.

  ## Seeding

  When a generator is created without an explicit `:seed`, it
  derives one from `:erlang.unique_integer/1` — fine for one-off
  use, but iterations may diverge. For Jepsen-style reproducibility,
  pass an explicit seed (typically derived from the iteration's seed
  via `phash2({iter_seed, worker_id})`).
  """

  @enforce_keys [:kind]
  defstruct [:kind, :ops, :seed, :remaining, :weights, :total, :children, :next_child]

  @type t :: %__MODULE__{
          kind: :random | :cycle | :weighted | :mix,
          ops: [any()] | nil,
          seed: any(),
          remaining: non_neg_integer() | :infinity,
          weights: [non_neg_integer()] | nil,
          total: non_neg_integer() | nil,
          children: [t()] | nil,
          next_child: non_neg_integer() | nil
        }

  @doc """
  Uniform random sampling over `ops` with replacement. Pass `:seed`
  to make the stream reproducible across runs.
  """
  @spec random([any()], keyword()) :: t()
  def random(ops, opts \\ []) when is_list(ops) and ops != [] do
    seed = Keyword.get(opts, :seed, :erlang.unique_integer([:positive]))

    %__MODULE__{
      kind: :random,
      ops: ops,
      seed: :rand.seed_s(:exs1024s, {seed, seed * 7, seed * 11}),
      remaining: :infinity
    }
  end

  @doc "Fixed sequence repeated forever (until `take/2`-limited)."
  @spec cycle([any()]) :: t()
  def cycle(ops) when is_list(ops) and ops != [] do
    %__MODULE__{kind: :cycle, ops: ops, remaining: :infinity}
  end

  @doc """
  Weighted random sampling. `weighted_ops` is a list of `{op, weight}`
  tuples. Higher weight = more likely to be picked.
  """
  @spec weighted([{any(), pos_integer()}], keyword()) :: t()
  def weighted(weighted_ops, opts \\ [])
      when is_list(weighted_ops) and weighted_ops != [] do
    seed = Keyword.get(opts, :seed, :erlang.unique_integer([:positive]))
    {ops, weights} = Enum.unzip(weighted_ops)
    total = Enum.sum(weights)

    if total <= 0, do: raise(ArgumentError, "weights must sum to a positive integer")

    %__MODULE__{
      kind: :weighted,
      ops: ops,
      weights: weights,
      total: total,
      seed: :rand.seed_s(:exs1024s, {seed, seed * 13, seed * 17}),
      remaining: :infinity
    }
  end

  @doc """
  Round-robin interleave outputs from `gens`. When any child is
  exhausted, this generator becomes done.
  """
  @spec mix([t()]) :: t()
  def mix(gens) when is_list(gens) and gens != [] do
    %__MODULE__{
      kind: :mix,
      children: gens,
      next_child: 0,
      remaining: :infinity
    }
  end

  @doc """
  Cap any generator at `n` total ops. After `n` calls to `next/1`,
  the generator returns `:done`.
  """
  @spec take(t(), non_neg_integer()) :: t()
  def take(%__MODULE__{} = gen, n) when is_integer(n) and n >= 0 do
    %{gen | remaining: n}
  end

  @doc """
  Pull the next op. Returns `{op, new_gen}` or `:done` when the
  generator has been exhausted (only happens after `take/2` or when
  `mix/1`'s last child finishes).
  """
  @spec next(t()) :: {any(), t()} | :done
  def next(%__MODULE__{remaining: 0}), do: :done

  def next(%__MODULE__{kind: :random, ops: ops, seed: seed} = gen) do
    {idx, new_seed} = :rand.uniform_s(length(ops), seed)
    op = Enum.at(ops, idx - 1)
    {op, %{gen | seed: new_seed, remaining: dec(gen.remaining)}}
  end

  def next(%__MODULE__{kind: :cycle, ops: [head | tail]} = gen) do
    {head, %{gen | ops: tail ++ [head], remaining: dec(gen.remaining)}}
  end

  def next(%__MODULE__{kind: :weighted} = gen) do
    {target, new_seed} = :rand.uniform_s(gen.total, gen.seed)
    op = pick_weighted(gen.ops, gen.weights, target)
    {op, %{gen | seed: new_seed, remaining: dec(gen.remaining)}}
  end

  def next(%__MODULE__{kind: :mix, children: children, next_child: idx} = gen) do
    child = Enum.at(children, idx)

    case next(child) do
      :done ->
        :done

      {op, new_child} ->
        new_children = List.replace_at(children, idx, new_child)

        new_gen = %{
          gen
          | children: new_children,
            next_child: rem(idx + 1, length(children)),
            remaining: dec(gen.remaining)
        }

        {op, new_gen}
    end
  end

  defp dec(:infinity), do: :infinity
  defp dec(n) when is_integer(n) and n > 0, do: n - 1

  defp pick_weighted([op | _], [w | _], target) when target <= w, do: op

  defp pick_weighted([_ | ops], [w | weights], target) when target > w do
    pick_weighted(ops, weights, target - w)
  end

  @doc """
  Iterate `gen` to exhaustion, calling `fun.(op)` for each op.
  Returns the count of ops produced.
  """
  @spec each(t(), (any() -> any())) :: non_neg_integer()
  def each(%__MODULE__{} = gen, fun) when is_function(fun, 1) do
    do_each(gen, fun, 0)
  end

  defp do_each(gen, fun, count) do
    case next(gen) do
      :done ->
        count

      {op, gen2} ->
        fun.(op)
        do_each(gen2, fun, count + 1)
    end
  end

  @doc """
  Reduce over `gen`'s ops to exhaustion. `fun.(op, acc)` returns
  the new acc. Useful for accumulating results without an external
  side effect.
  """
  @spec reduce(t(), acc, (any(), acc -> acc)) :: acc when acc: var
  def reduce(%__MODULE__{} = gen, acc, fun) when is_function(fun, 2) do
    case next(gen) do
      :done -> acc
      {op, gen2} -> reduce(gen2, fun.(op, acc), fun)
    end
  end
end