firebird

defmodule Firebird.Phoenix.RateLimiter do
  @moduledoc """
  Token bucket rate limiter for Phoenix WASM request handling.

  Implements a token bucket algorithm for rate limiting HTTP requests.
  Tokens are replenished at a configurable rate, and each request
  consumes one token. When the bucket is empty, requests are rejected.

  Can be used as middleware with `Firebird.Phoenix.RequestHandler`.

  ## Usage

      # Create a rate limiter (100 requests per minute)
      limiter = RateLimiter.new(rate: 100, period: :minute)

      # Check and consume a token
      case RateLimiter.check(limiter, "client-ip") do
        {:ok, limiter, remaining} -> # allow request
        {:error, :rate_limited, limiter, retry_after} -> # reject
      end

      # Use as middleware
      handler = RequestHandler.use_middleware(handler,
        RateLimiter.middleware(rate: 10, period: :second))
  """

  @type t :: %__MODULE__{
          rate: pos_integer(),
          period_ms: pos_integer(),
          buckets: %{String.t() => bucket()},
          max_tokens: pos_integer()
        }

  @type bucket :: %{
          tokens: float(),
          last_refill: integer()
        }

  defstruct rate: 100,
            period_ms: 60_000,
            buckets: %{},
            max_tokens: 100

  @doc """
  Create a new rate limiter.

  ## Options

  - `:rate` - Maximum requests per period (default: 100)
  - `:period` - Time period: `:second`, `:minute`, `:hour` (default: :minute)
  - `:burst` - Maximum burst size (default: same as rate)
  """
  @spec new(keyword()) :: t()
  def new(opts \\ []) do
    rate = Keyword.get(opts, :rate, 100)
    period = Keyword.get(opts, :period, :minute)
    burst = Keyword.get(opts, :burst, rate)

    period_ms =
      case period do
        :second -> 1_000
        :minute -> 60_000
        :hour -> 3_600_000
        ms when is_integer(ms) -> ms
      end

    %__MODULE__{
      rate: rate,
      period_ms: period_ms,
      buckets: %{},
      max_tokens: burst
    }
  end

  @doc """
  Check if a request is allowed and consume a token.

  ## Parameters

  - `limiter` - Rate limiter state
  - `key` - Client identifier (IP, user ID, API key, etc.)

  ## Returns

  - `{:ok, new_limiter, remaining}` - Request allowed, tokens remaining
  - `{:error, :rate_limited, new_limiter, retry_after_ms}` - Rate limited
  """
  @spec check(t(), String.t()) ::
          {:ok, t(), non_neg_integer()} | {:error, :rate_limited, t(), non_neg_integer()}
  def check(%__MODULE__{} = limiter, key) when is_binary(key) do
    now = System.monotonic_time(:millisecond)
    bucket = get_bucket(limiter, key, now)

    if bucket.tokens >= 1.0 do
      new_bucket = %{bucket | tokens: bucket.tokens - 1.0}
      new_limiter = put_bucket(limiter, key, new_bucket)
      {:ok, new_limiter, trunc(new_bucket.tokens)}
    else
      # Calculate when next token will be available
      tokens_per_ms = limiter.rate / limiter.period_ms

      retry_after =
        if tokens_per_ms > 0,
          do: trunc((1.0 - bucket.tokens) / tokens_per_ms),
          else: limiter.period_ms

      {:error, :rate_limited, limiter, retry_after}
    end
  end

  @doc """
  Get the current token count for a key without consuming.
  """
  @spec remaining(t(), String.t()) :: non_neg_integer()
  def remaining(%__MODULE__{} = limiter, key) do
    now = System.monotonic_time(:millisecond)
    bucket = get_bucket(limiter, key, now)
    trunc(bucket.tokens)
  end

  @doc """
  Reset the rate limiter for a specific key.
  """
  @spec reset(t(), String.t()) :: t()
  def reset(%__MODULE__{} = limiter, key) do
    %{limiter | buckets: Map.delete(limiter.buckets, key)}
  end

  @doc """
  Reset all rate limiter state.
  """
  @spec reset_all(t()) :: t()
  def reset_all(%__MODULE__{} = limiter) do
    %{limiter | buckets: %{}}
  end

  @doc """
  Create a rate limiting middleware for use with `RequestHandler`.

  Uses the client's IP or a configurable header as the rate limit key.

  ## Options

  - `:rate` - Max requests per period
  - `:period` - Time period
  - `:key_header` - Header to use as key (default: uses assigns[:remote_ip] or "unknown")
  - `:error_body` - Custom error body

  ## Examples

      middleware = RateLimiter.middleware(rate: 100, period: :minute)
      handler = RequestHandler.use_middleware(handler, middleware)
  """
  @spec middleware(keyword()) :: (Firebird.Phoenix.Conn.t() -> Firebird.Phoenix.Conn.t())
  def middleware(opts \\ []) do
    key_header = Keyword.get(opts, :key_header, nil)
    error_body = Keyword.get(opts, :error_body, "Rate limit exceeded")

    # Use process dictionary for state (middleware is stateless per-call)
    limiter_key = {__MODULE__, :limiter}
    initial = new(opts)

    fn conn ->
      alias Firebird.Phoenix.Conn

      limiter = Process.get(limiter_key) || initial

      client_key =
        cond do
          key_header -> Map.get(conn.req_headers, key_header, "unknown")
          conn.assigns[:remote_ip] -> conn.assigns[:remote_ip]
          true -> "unknown"
        end

      case check(limiter, client_key) do
        {:ok, new_limiter, remaining} ->
          Process.put(limiter_key, new_limiter)

          conn
          |> Conn.put_resp_header("x-ratelimit-remaining", to_string(remaining))
          |> Conn.put_resp_header("x-ratelimit-limit", to_string(limiter.rate))

        {:error, :rate_limited, new_limiter, retry_after} ->
          Process.put(limiter_key, new_limiter)

          conn
          |> Conn.put_status(429)
          |> Conn.put_resp_header("retry-after", to_string(div(retry_after, 1000)))
          |> Conn.put_resp_header("x-ratelimit-remaining", "0")
          |> Conn.put_resp_body(error_body)
          |> Conn.halt()
      end
    end
  end

  # Internal: get or create bucket with token refill
  defp get_bucket(%__MODULE__{} = limiter, key, now) do
    case Map.get(limiter.buckets, key) do
      nil ->
        %{tokens: limiter.max_tokens * 1.0, last_refill: now}

      bucket ->
        elapsed = now - bucket.last_refill
        tokens_per_ms = limiter.rate / limiter.period_ms

        new_tokens =
          min(
            bucket.tokens + elapsed * tokens_per_ms,
            limiter.max_tokens * 1.0
          )

        %{bucket | tokens: new_tokens, last_refill: now}
    end
  end

  defp put_bucket(%__MODULE__{} = limiter, key, bucket) do
    %{limiter | buckets: Map.put(limiter.buckets, key, bucket)}
  end
end