defmodule ToweropsWeb.Plugs.BruteForceProtection do
  @moduledoc """
  Plug that detects and blocks brute force scanning behavior.

  ## Detection Strategy

  Tracks unique 404 paths per IP address using Redis. When an unauthenticated
  IP hits 5+ unique 404s within 60 seconds, a ban is created/escalated.

  ## Ban Escalation

  - 1st offense: 5-minute ban
  - 2nd offense (within 30 days): 1-hour ban
  - 3rd+ offense (within 30 days): Permanent ban + Cloudflare WAF push

  ## Exemptions

  - Authenticated users (conn.assigns[:current_user] present)
  - Whitelisted IPs (checked via BruteForce context)

  ## Placement

  Must be placed in the endpoint.ex pipeline AFTER:
  - RemoteIpLogger (needs remote_ip in Logger metadata)
  - Static file serving (no need to track static assets)

  Must be placed BEFORE:
  - Router (needs to intercept requests before routing)
  """

  import Plug.Conn

  alias Towerops.Security.BruteForce

  require Logger

  def init(opts), do: opts

  def call(conn, _opts) do
    ip_address = get_remote_ip(conn)

    # Check whitelist and ban status before routing
    conn = check_and_block(conn, ip_address)

    # Only register 404 tracking if we didn't already halt the request
    # (halted requests have already been sent, can't register before_send)
    if conn.halted do
      conn
    else
      maybe_track_404(conn, ip_address)
    end
  end

  defp check_and_block(conn, ip_address) do
    # Whitelist bypass
    if BruteForce.check_whitelist(ip_address) do
      conn
    else
      # Check if IP is currently banned
      case BruteForce.check_ban_status(ip_address) do
        :allowed ->
          conn

        {:blocked, banned_until} ->
          Logger.warning("Blocked request from banned IP: #{ip_address}")
          block_request(conn, banned_until)
      end
    end
  end

  defp maybe_track_404(conn, ip_address) do
    # Register before_send callback to track 404s after routing
    register_before_send(conn, fn conn ->
      track_404_if_needed(conn, ip_address)
      conn
    end)
  end

  defp track_404_if_needed(conn, ip_address) do
    # Only track 404s from unauthenticated users
    if conn.status == 404 && is_nil(Map.get(conn.assigns, :current_user)) do
      path = conn.request_path
      Logger.debug("Tracking 404 for IP #{ip_address}: #{path}")
      BruteForce.record_404_path(ip_address, path)
    end
  end

  defp block_request(conn, banned_until) do
    retry_after =
      if banned_until do
        DateTime.diff(banned_until, DateTime.utc_now(), :second)
      else
        # Permanent ban - use a large value
        86_400
      end

    conn
    |> put_resp_header("retry-after", to_string(max(0, retry_after)))
    |> send_resp(403, "Forbidden - Too many invalid requests")
    |> halt()
  end

  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 | _] ->
        # X-Forwarded-For can be a comma-separated list; take the first (original client)
        forwarded
        |> String.split(",")
        |> List.first()
        |> String.trim()

      [] ->
        get_fallback_ip(conn)
    end
  end

  defp get_fallback_ip(conn) do
    # Try X-Real-IP header
    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
