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

  @behaviour Oban.Pro.Limiter

  alias Oban.Pro.Partition

  @impl Oban.Pro.Limiter
  def check(%{producer: producer} = changes) when is_map(producer.meta.global_limit) do
    tracked =
      [producer | changes.all_producers]
      |> Enum.uniq_by(& &1.uuid)
      |> Enum.map(& &1.meta.global_limit)
      |> Enum.filter(&is_map/1)
      |> Enum.reduce(%{}, &merge/2)

    allowed = producer.meta.global_limit.allowed
    tracked = update_legacy_tracked(tracked)

    if is_nil(producer.meta.global_limit.partition) do
      {:ok, combined_demands(allowed, tracked)}
    else
      {:ok, separate_demands(changes.conf, allowed, producer, tracked)}
    end
  end

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

  defp merge(global_limit, acc) do
    Map.merge(global_limit.tracked, acc, fn _, v1, v2 -> v1 + v2 end)
  end

  defp combined_demands(allowed, tracked) do
    total = tracked |> Map.values() |> Enum.sum()

    max(allowed - total, 0)
  end

  defp separate_demands(conf, allowed, producer, tracked) do
    untracked =
      for key <- Partition.available_keys(conf, producer.queue, producer.meta.local_limit),
          not is_map_key(tracked, key),
          into: %{},
          do: {key, allowed}

    allowed =
      if producer.meta.global_limit.burst do
        size =
          tracked
          |> Map.merge(untracked)
          |> Map.delete("none")
          |> map_size()
          |> max(1)

        producer.meta.local_limit
        |> div(size)
        |> max(allowed)
      else
        allowed
      end

    Enum.reduce(tracked, untracked, fn {key, total}, acc ->
      case max(allowed - total, 0) do
        demand when demand > 0 -> Map.put(acc, key, demand)
        _ -> acc
      end
    end)
  end

  # Tracking must be updated after acking and before fetching in order to have accurate limits for
  # the current fetch transaction.
  def prepare_tracked(%{global_limit: global} = meta, keys) do
    tracked =
      Enum.reduce(keys, global.tracked, fn key, acc ->
        if is_map_key(acc, key) do
          Map.update!(acc, key, &(&1 - 1))
        else
          acc
        end
      end)

    put_in(meta.global_limit.tracked, tracked)
  end

  @impl Oban.Pro.Limiter
  def track(%{global_limit: global} = meta, jobs) when is_map(global) do
    tracked = Enum.frequencies_by(jobs, &Partition.get_key(&1, "*"))

    put_in(meta.global_limit.tracked, tracked)
  end

  def track(meta, _jobs), do: meta

  defp update_legacy_tracked(%{"count" => count}), do: %{"none" => count}

  defp update_legacy_tracked(tracked) do
    Map.new(tracked, fn
      {key, %{"count" => count}} -> {key, count}
      tuple -> tuple
    end)
  end
end