defmodule Oban.Pro.Limiters.Rate.TokenBucket do
  @moduledoc false

  # Tokens refill continuously at a rate of `allowed / period` per second. Available demand is the
  # number of tokens in the bucket.

  @behaviour Oban.Pro.Limiters.Rate.Algorithm

  @impl Oban.Pro.Limiters.Rate.Algorithm
  def demand(windows, allowed, _period, _curr_time, _prev_time) do
    total_consumed = sum_consumed(windows)
    max(allowed - total_consumed, 0)
  end

  @impl Oban.Pro.Limiters.Rate.Algorithm
  def partition_demands(windows, allowed, _period, _curr_time, _prev_time, available_keys) do
    untracked =
      for key <- available_keys,
          not is_map_key(windows, key),
          into: %{},
          do: {key, allowed}

    Enum.reduce(windows, untracked, fn {key, val}, acc ->
      case max(allowed - val["consumed"], 0) do
        demand when demand > 0 -> Map.put(acc, key, demand)
        _zero -> acc
      end
    end)
  end

  @impl Oban.Pro.Limiters.Rate.Algorithm
  def flatten(windows) do
    consumed = Enum.reduce(windows, 0, fn {_key, val}, acc -> acc + val["consumed"] end)

    [%{"consumed" => consumed}]
  end

  @impl Oban.Pro.Limiters.Rate.Algorithm
  def merge(windows_a, windows_b) do
    Map.merge(windows_a, windows_b, fn _key, v1, v2 ->
      %{"consumed" => v1["consumed"] + v2["consumed"]}
    end)
  end

  @doc """
  Track job execution by updating consumed tokens.

  Refills tokens based on elapsed time, then adds new consumption.
  Returns the updated windows map and the next window_time.
  """
  @impl Oban.Pro.Limiters.Rate.Algorithm
  def track(windows, new_counts, period, prev_time, curr_time, allowed) do
    refill_amount = calculate_refill(allowed, period, curr_time, prev_time)

    new_consumed = Map.new(new_counts, fn {key, count} -> {key, %{"consumed" => count}} end)

    all_windows =
      Map.merge(windows, new_consumed, fn _key, existing, new ->
        updated_consumed = max(existing["consumed"] - refill_amount, 0) + new["consumed"]
        %{"consumed" => trunc(updated_consumed)}
      end)

    all_windows =
      Enum.reduce(new_consumed, all_windows, fn {key, val}, acc ->
        Map.put_new(acc, key, val)
      end)

    all_windows =
      for {key, val} <- all_windows,
          val["consumed"] > 0,
          into: %{},
          do: {key, val}

    {all_windows, curr_time}
  end

  defp calculate_refill(allowed, period, curr_time, prev_time) do
    elapsed = curr_time - prev_time
    refill_rate = allowed / period
    min(elapsed * refill_rate, allowed)
  end

  defp sum_consumed(windows) do
    Enum.reduce(windows, 0, fn {_key, val}, acc -> acc + val["consumed"] end)
  end
end
