use axum::extract::Request;
use std::net::{IpAddr, SocketAddr};

/// Extract the client IP from a request.
///
/// When the `TRUST_PROXY` environment variable is set to `1` or `true`, the
/// `X-Forwarded-For` and `X-Real-IP` headers are consulted first, which is
/// appropriate when the service runs behind a trusted reverse proxy (e.g.
/// Caddy). Otherwise — or if those headers are absent/empty — the direct TCP
/// connection address from `ConnectInfo` is used.
///
/// **Important**: only enable `TRUST_PROXY` when the application is guaranteed
/// to receive traffic exclusively through a controlled reverse proxy. Enabling
/// it on a directly-exposed port allows clients to spoof their IP address and
/// bypass per-IP rate limiting.
pub fn extract_client_ip(req: &Request) -> String {
    if trust_proxy_enabled() {
        if let Some(ip) = forwarded_for_ip(req.headers()) {
            return ip;
        }
        if let Some(ip) = real_ip(req.headers()) {
            return ip;
        }
    }

    connect_info_ip(req).unwrap_or_else(|| "unknown".to_string())
}

fn trust_proxy_enabled() -> bool {
    static CACHE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *CACHE.get_or_init(|| {
        matches!(
            std::env::var("TRUST_PROXY").as_deref(),
            Ok("1") | Ok("true") | Ok("yes")
        )
    })
}

fn forwarded_for_ip(headers: &axum::http::HeaderMap) -> Option<String> {
    let value = headers.get("x-forwarded-for")?.to_str().ok()?;
    let first = value.split(',').next()?.trim();
    if first.is_empty() {
        None
    } else {
        validate_ip(first)
    }
}

fn real_ip(headers: &axum::http::HeaderMap) -> Option<String> {
    let value = headers.get("x-real-ip")?.to_str().ok()?;
    let ip = value.trim();
    if ip.is_empty() { None } else { validate_ip(ip) }
}

/// Validate that a string is a valid IP address.
/// Returns Some(ip) if valid, None otherwise.
fn validate_ip(s: &str) -> Option<String> {
    s.parse::<IpAddr>().ok().map(|ip| ip.to_string())
}

fn connect_info_ip(req: &Request) -> Option<String> {
    req.extensions()
        .get::<axum::extract::ConnectInfo<SocketAddr>>()
        .map(|c| c.0.ip().to_string())
}