defmodule ToweropsWeb.Plugs.RateLimit do
  @moduledoc """
  Rate limiting plug using Hammer to prevent brute-force attacks.

  Applies different rate limits based on the type of endpoint:
  - Auth endpoints (login, register, password reset): 10 requests per minute per IP
  - API endpoints: 100 requests per minute per IP

  Uses the client's real IP address from X-Forwarded-For or X-Real-IP headers
  when behind a proxy/load balancer.
  """

  import Phoenix.Controller, only: [json: 2]
  import Plug.Conn

  require Logger

  @auth_limit 10
  @auth_period to_timeout(minute: 1)

  @api_limit 1000
  @api_period to_timeout(minute: 1)

  @type limit_type :: :auth | :api

  def init(opts) do
    type = Keyword.get(opts, :type, :api)

    if type not in [:auth, :api] do
      raise ArgumentError, "RateLimit plug :type must be :auth or :api, got: #{inspect(type)}"
    end

    type
  end

  def call(conn, type) do
    if rate_limiting_enabled?() do
      check_rate_limit(conn, type)
    else
      conn
    end
  end

  defp rate_limiting_enabled? do
    Application.get_env(:towerops, :rate_limiting_enabled, true)
  end

  defp check_rate_limit(conn, type) do
    remote_ip = get_remote_ip(conn)
    bucket_key = "#{type}:#{remote_ip}"

    {limit, period} = limits_for(type)

    case Towerops.RateLimit.hit(bucket_key, period, limit) do
      {:allow, _count} ->
        conn

      {:deny, retry_after} ->
        Logger.warning("Rate limit exceeded for #{bucket_key}, retry_after: #{retry_after}ms")

        conn
        |> put_resp_header("retry-after", Integer.to_string(div(retry_after, 1000)))
        |> put_status(:too_many_requests)
        |> json(%{error: "Too many requests. Please try again later."})
        |> halt()
    end
  end

  defp limits_for(:auth), do: {@auth_limit, @auth_period}
  defp limits_for(:api), do: {@api_limit, @api_period}

  defp get_remote_ip(conn) do
    # Try X-Forwarded-For first (set by Traefik and other proxies)
    case get_req_header(conn, "x-forwarded-for") do
      [forwarded | _] ->
        forwarded
        |> String.split(",")
        |> List.first()
        |> String.trim()

      [] ->
        get_fallback_ip(conn)
    end
  end

  defp get_fallback_ip(conn) do
    case get_req_header(conn, "x-real-ip") do
      [real_ip | _] ->
        real_ip

      [] ->
        format_remote_ip(conn.remote_ip)
    end
  end

  defp format_remote_ip({a, b, c, d}), do: "#{a}.#{b}.#{c}.#{d}"

  defp format_remote_ip({a, b, c, d, e, f, g, h}) do
    Enum.map_join([a, b, c, d, e, f, g, h], ":", &Integer.to_string(&1, 16))
  end
end