use std::collections::HashMap;
use std::sync::Mutex;
use std::time::Instant;

use axum::http::{HeaderValue, Request, Response, StatusCode};
use axum::middleware::Next;
use axum::body::Body;

// per-key rate limit state
struct RateWindow {
    count: u32,
    window_start: Instant,
    window_secs: u64,
}

static RATE_LIMITS: std::sync::LazyLock<Mutex<HashMap<String, RateWindow>>> =
    std::sync::LazyLock::new(|| Mutex::new(HashMap::new()));

const RATE_LIMIT_MAX_ENTRIES: usize = 4096;

pub fn check_rate_limit(session: &str, key: &str, max: u32, window_secs: u64) -> Result<(), StatusCode> {
    check_bucket(&format!("{}:{}", session, key), max, window_secs)
}

pub fn check_rate_limit_ip(ip: &str, key: &str, max: u32, window_secs: u64) -> Result<(), StatusCode> {
    check_bucket(&format!("ip:{}:{}", ip, key), max, window_secs)
}

fn check_bucket(bucket: &str, max: u32, window_secs: u64) -> Result<(), StatusCode> {
    let mut map = RATE_LIMITS.lock().unwrap();
    let now = Instant::now();

    // bound memory: prune expired windows occasionally instead of growing
    // forever with per-session and per-puzzle bucket keys
    if map.len() > RATE_LIMIT_MAX_ENTRIES {
        map.retain(|_, w| now.duration_since(w.window_start).as_secs() < w.window_secs);
    }

    let entry = map.entry(bucket.to_string()).or_insert(RateWindow { count: 0, window_start: now, window_secs });
    if now.duration_since(entry.window_start).as_secs() >= window_secs {
        entry.count = 1;
        entry.window_start = now;
        Ok(())
    } else if entry.count < max {
        entry.count += 1;
        Ok(())
    } else {
        Err(StatusCode::TOO_MANY_REQUESTS)
    }
}

pub fn extract_client_ip(headers: &axum::http::HeaderMap) -> String {
    // proxy headers are only meaningful when an explicit trusted reverse proxy
    // is in front; otherwise a client can supply whatever it wants
    if std::env::var("TRUST_PROXY").as_deref() != Ok("true") {
        return "unknown".to_string();
    }
    headers.get("x-forwarded-for")
        .or_else(|| headers.get("x-real-ip"))
        .and_then(|v| v.to_str().ok())
        .unwrap_or("unknown")
        .split(',').next().unwrap_or("unknown").trim().to_string()
}

pub async fn security_headers(
    request: Request<Body>,
    next: Next,
) -> Response<Body> {
    let mut response = next.run(request).await;
    let headers = response.headers_mut();

    headers.insert("X-Content-Type-Options", HeaderValue::from_static("nosniff"));
    headers.insert("X-Frame-Options", HeaderValue::from_static("DENY"));
    headers.insert("X-XSS-Protection", HeaderValue::from_static("1; mode=block"));
    headers.insert("Referrer-Policy", HeaderValue::from_static("strict-origin-when-cross-origin"));
    headers.insert("Permissions-Policy", HeaderValue::from_static(
        "camera=(), microphone=(), geolocation=(), payment=()"
    ));
    headers.insert("Content-Security-Policy", HeaderValue::from_static(
        "default-src 'self'; script-src 'self' https://cdn.jsdelivr.net; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com https://cdn.jsdelivr.net; font-src 'self' https://fonts.gstatic.com; img-src 'self' data:; connect-src 'self'; frame-ancestors 'none'"
    ));

    if std::env::var("HSTS").as_deref() == Ok("true") {
        headers.insert(
            "Strict-Transport-Security",
            HeaderValue::from_static("max-age=31536000; includeSubDomains"),
        );
    }

    response
}

pub fn is_internal_url(url: &str) -> bool {
    let url_lower = url.to_lowercase();
    let host = url_lower
        .strip_prefix("http://").or_else(|| url_lower.strip_prefix("https://"))
        .unwrap_or(&url_lower)
        .split('/').next().unwrap_or("")
        .split(':').next().unwrap_or("");

    // strip IPv6 brackets
    let host = host.strip_prefix('[').and_then(|h| h.strip_suffix(']')).unwrap_or(host);

    // block non-dotted numeric hosts (decimal IP like 2130706433 = 127.0.0.1)
    if !host.is_empty() && host.chars().all(|c| c.is_ascii_digit()) { return true; }

    // block octal IPs (0177.0.0.1 = 127.0.0.1)
    if host.starts_with("0") && host.contains('.') && host.chars().all(|c| c.is_ascii_digit() || c == '.') {
        return true;
    }

    // block URLs with @ (http://user@internal/path)
    if host.contains('@') { return true; }

    // loopback
    host == "localhost"
        || host.starts_with("127.")            // full 127.0.0.0/8 range
        || host == "::1"
        || host.starts_with("::ffff:127.")     // IPv4-mapped IPv6 loopback
        // current network
        || host.starts_with("0.")              // 0.0.0.0/8
        || host == "0.0.0.0"
        // RFC 1918 private ranges
        || host.starts_with("10.")             // 10.0.0.0/8
        || is_172_private(host)                // 172.16.0.0/12
        || host.starts_with("192.168.")        // 192.168.0.0/16
        // carrier-grade NAT (CVE-2026-22245 - Mastodon missed this)
        || is_cgnat(host)                      // 100.64.0.0/10
        // link-local
        || host.starts_with("169.254.")        // 169.254.0.0/16
        // benchmark testing
        || host.starts_with("198.18.")         // 198.18.0.0/15
        || host.starts_with("198.19.")
        // IPv6 private/local
        || host.starts_with("fc")              // fc00::/7 unique local
        || host.starts_with("fd")
        || host.starts_with("fe80")            // fe80::/10 link-local
        || host.starts_with("::ffff:10.")      // IPv4-mapped private
        || host.starts_with("::ffff:192.168.")
        || host.starts_with("::ffff:0.")
        // DNS rebinding / special hostnames
        || host.ends_with(".local")
        || host.ends_with(".internal")
        || host.ends_with(".localhost")
        || host == "metadata.google.internal"
        // cloud metadata endpoints
        || host == "169.254.169.254"
        || host == "metadata.google.internal"
        || host == "100.100.100.200"           // Alibaba cloud metadata
}

fn is_172_private(host: &str) -> bool {
    if !host.starts_with("172.") { return false; }
    let second = host.strip_prefix("172.").and_then(|s| s.split('.').next())
        .and_then(|s| s.parse::<u8>().ok()).unwrap_or(0);
    (16..=31).contains(&second)
}

fn is_cgnat(host: &str) -> bool {
    if !host.starts_with("100.") { return false; }
    let second = host.strip_prefix("100.").and_then(|s| s.split('.').next())
        .and_then(|s| s.parse::<u8>().ok()).unwrap_or(0);
    (64..=127).contains(&second)
}

pub fn sanitize_html(input: &str) -> String {
    input
        .replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&#x27;")
}

/// Redirect policy that refuses to follow hops into internal or loopback
/// addresses, so an initial URL check cannot be undone by a redirect chain.
pub fn safe_redirect_policy() -> reqwest::redirect::Policy {
    reqwest::redirect::Policy::custom(|attempt| {
        if is_internal_url(attempt.url().as_str()) {
            attempt.stop()
        } else {
            attempt.follow()
        }
    })
}

/// Sanitize ActivityPub note content for rendering on HTML permalink pages.
/// Only the tags the daily-post template actually emits survive.
pub fn sanitize_note_html(html: &str) -> String {
    let tags: std::collections::HashSet<&str> = ["a", "br", "p", "span"].into_iter().collect();
    ammonia::Builder::default()
        .tags(tags)
        .link_rel(Some("nofollow noopener noreferrer"))
        .clean(html)
        .to_string()
}
