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

  @behaviour Oban.Pro.Limiter

  alias Oban.Pro.Limiters.Rate.Algorithm
  alias Oban.Pro.Partition

  @impl Oban.Pro.Limiter
  def check(%{producer: producer} = changes) when is_map(producer.meta.rate_limit) do
    %{window_time: prev_time} = producer.meta.rate_limit

    curr_time = unix_now()

    windows =
      [producer | changes.all_producers]
      |> Enum.uniq_by(& &1.uuid)
      |> Enum.map(& &1.meta.rate_limit)
      |> Enum.filter(&(is_map(&1) and &1.window_time >= curr_time - &1.period))
      |> Enum.reduce(%{}, &merge_windows/2)

    windows = update_legacy_windows(windows)

    demand = calculate_demand(producer.meta.rate_limit, windows, curr_time, prev_time, changes)

    {:ok, demand}
  end

  def check(_changes), do: {:ok, nil}

  defp calculate_demand(rate_limit, windows, curr_time, prev_time, changes) do
    %{allowed: allowed, period: period} = rate_limit

    algorithm = Algorithm.callback(rate_limit)

    if is_nil(rate_limit.partition) do
      algorithm.demand(windows, allowed, period, curr_time, prev_time)
    else
      producer = changes.producer

      available_keys =
        Partition.available_keys(changes.conf, producer.queue, producer.meta.local_limit)

      algorithm.partition_demands(
        windows,
        allowed,
        period,
        curr_time,
        prev_time,
        available_keys
      )
    end
  end

  defp merge_windows(rate_limit, acc) do
    Algorithm.callback(rate_limit).merge(rate_limit.windows, acc)
  end

  defp update_legacy_windows(%{"args" => _} = window), do: %{"*" => window}
  defp update_legacy_windows(windows), do: windows

  def flatten(%{rate_limit: rate_limit}) do
    Algorithm.callback(rate_limit).flatten(rate_limit.windows)
  end

  @impl Oban.Pro.Limiter
  def track(%{rate_limit: rate_limit} = meta, jobs) when is_map(rate_limit) do
    new_counts = build_counts(jobs)
    curr_time = unix_now()

    {all_windows, next_time} =
      Algorithm.callback(rate_limit).track(
        rate_limit.windows,
        new_counts,
        rate_limit.period,
        rate_limit.window_time,
        curr_time,
        rate_limit.allowed
      )

    %{meta | rate_limit: %{rate_limit | windows: all_windows, window_time: next_time}}
  end

  def track(meta, _jobs), do: meta

  defp build_counts(jobs) do
    Enum.reduce(jobs, %{}, fn job, acc ->
      key = Partition.get_key(job, "*")
      weight = Oban.Pro.Worker.get_weight(job)

      Map.update(acc, key, weight, &(&1 + weight))
    end)
  end

  defp unix_now do
    DateTime.to_unix(DateTime.utc_now(), :second)
  end
end