use once_cell::sync::Lazy;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use regex::Regex;

static LEGITIMATE_BOTS: Lazy<Vec<Regex>> = Lazy::new(|| {
    vec![
        Regex::new(r"googlebot").unwrap(),
        Regex::new(r"bingbot").unwrap(),
        Regex::new(r"slurp").unwrap(),
        Regex::new(r"duckduckbot").unwrap(),
        Regex::new(r"baiduspider").unwrap(),
        Regex::new(r"yandexbot").unwrap(),
        Regex::new(r"facebookexternalhit").unwrap(),
        Regex::new(r"twitterbot").unwrap(),
        Regex::new(r"linkedinbot").unwrap(),
        Regex::new(r"whatsapp").unwrap(),
        Regex::new(r"telegrambot").unwrap(),
        Regex::new(r"applebot").unwrap(),
        Regex::new(r"pingdom").unwrap(),
        Regex::new(r"uptimerobot").unwrap(),
        Regex::new(r"statuscake").unwrap(),
        Regex::new(r"site24x7").unwrap(),
    ]
});

static SUSPICIOUS_UA: Lazy<Vec<(&'static str, Regex)>> = Lazy::new(|| {
    vec![
        (r"bot", Regex::new(r"bot").unwrap()),
        (r"crawler", Regex::new(r"crawler").unwrap()),
        (r"spider", Regex::new(r"spider").unwrap()),
        (r"scraper", Regex::new(r"scraper").unwrap()),
        (r"curl", Regex::new(r"curl").unwrap()),
        (r"wget", Regex::new(r"wget").unwrap()),
        (r"python", Regex::new(r"python").unwrap()),
        (r"java", Regex::new(r"java").unwrap()),
        (r"node", Regex::new(r"node").unwrap()),
        (r"go-http", Regex::new(r"go-http").unwrap()),
        (r"axios", Regex::new(r"axios").unwrap()),
        (r"okhttp", Regex::new(r"okhttp").unwrap()),
        (r"libwww", Regex::new(r"libwww").unwrap()),
        (r"lwp-trivial", Regex::new(r"lwp-trivial").unwrap()),
        (r"mechanize", Regex::new(r"mechanize").unwrap()),
        (r"requests", Regex::new(r"requests").unwrap()),
        (r"urllib", Regex::new(r"urllib").unwrap()),
        (r"httpie", Regex::new(r"httpie").unwrap()),
        (r"postman", Regex::new(r"postman").unwrap()),
        (r"insomnia", Regex::new(r"insomnia").unwrap()),
        (r"^$", Regex::new(r"^$").unwrap()),
        (r"mozilla/4\.0$", Regex::new(r"mozilla/4\.0$").unwrap()),
    ]
});

fn get_header(headers: &Bound<'_, PyDict>, key: &str) -> Option<String> {
    headers
        .get_item(key)
        .ok()
        .flatten()
        .and_then(|v| v.str().ok())
        .and_then(|s| s.to_str().ok().map(|v| v.to_string()))
}

fn has_header(headers: &Bound<'_, PyDict>, key: &str) -> bool {
    match get_header(headers, key) {
        Some(value) => !value.is_empty(),
        None => false,
    }
}

fn check_user_agent(user_agent: &str) -> Option<String> {
    if user_agent.is_empty() {
        return Some("Empty user agent".to_string());
    }

    let ua_lower = user_agent.to_lowercase();

    for legit in LEGITIMATE_BOTS.iter() {
        if legit.is_match(&ua_lower) {
            return None;
        }
    }

    for (pattern, regex) in SUSPICIOUS_UA.iter() {
        if regex.is_match(&ua_lower) {
            return Some(format!("Pattern: {}", pattern));
        }
    }

    if user_agent.len() < 10 {
        return Some("Too short".to_string());
    }
    if user_agent.len() > 500 {
        return Some("Too long".to_string());
    }

    None
}

#[pyfunction]
fn validate_headers(headers: Bound<'_, PyDict>) -> PyResult<Option<String>> {
    let mut missing = Vec::new();
    if !has_header(&headers, "HTTP_USER_AGENT") {
        missing.push("user-agent".to_string());
    }
    if !has_header(&headers, "HTTP_ACCEPT") {
        missing.push("accept".to_string());
    }

    if !missing.is_empty() {
        return Ok(Some(format!(
            "Missing required headers: {}",
            missing.join(", ")
        )));
    }

    let user_agent = get_header(&headers, "HTTP_USER_AGENT").unwrap_or_default();
    if let Some(reason) = check_user_agent(&user_agent) {
        return Ok(Some(format!("Suspicious user agent: {}", reason)));
    }

    let server_protocol = get_header(&headers, "SERVER_PROTOCOL").unwrap_or_default();
    let accept = get_header(&headers, "HTTP_ACCEPT").unwrap_or_default();
    let accept_language = get_header(&headers, "HTTP_ACCEPT_LANGUAGE").unwrap_or_default();
    let accept_encoding = get_header(&headers, "HTTP_ACCEPT_ENCODING").unwrap_or_default();
    let connection = get_header(&headers, "HTTP_CONNECTION").unwrap_or_default();

    if server_protocol.starts_with("HTTP/2") && user_agent.to_lowercase().contains("mozilla/4.0") {
        return Ok(Some(
            "Suspicious headers: HTTP/2 with old browser user agent".to_string(),
        ));
    }
    if !user_agent.is_empty() && accept.is_empty() {
        return Ok(Some(
            "Suspicious headers: User-Agent present but no Accept header".to_string(),
        ));
    }
    if accept == "*/*" && accept_language.is_empty() && accept_encoding.is_empty() {
        return Ok(Some(
            "Suspicious headers: Generic Accept header without language/encoding".to_string(),
        ));
    }
    if !user_agent.is_empty()
        && accept_language.is_empty()
        && accept_encoding.is_empty()
        && connection.is_empty()
    {
        return Ok(Some(
            "Suspicious headers: Missing all browser-standard headers".to_string(),
        ));
    }
    if !user_agent.is_empty()
        && server_protocol == "HTTP/1.0"
        && user_agent.to_lowercase().contains("chrome")
    {
        return Ok(Some(
            "Suspicious headers: Modern browser with HTTP/1.0".to_string(),
        ));
    }

    let mut score = 0;
    if has_header(&headers, "HTTP_USER_AGENT") {
        score += 2;
    }
    if has_header(&headers, "HTTP_ACCEPT") {
        score += 2;
    }

    for header in [
        "HTTP_ACCEPT_LANGUAGE",
        "HTTP_ACCEPT_ENCODING",
        "HTTP_CONNECTION",
        "HTTP_CACHE_CONTROL",
    ] {
        if has_header(&headers, header) {
            score += 1;
        }
    }

    if !accept_language.is_empty() && !accept_encoding.is_empty() {
        score += 1;
    }
    if connection == "keep-alive" {
        score += 1;
    }
    if accept.contains("text/html") && accept.contains("application/xml") {
        score += 1;
    }

    if score < 3 {
        return Ok(Some(format!("Low header quality score: {}", score)));
    }

    Ok(None)
}

#[pymodule]
fn aiwaf_rust(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(validate_headers, m)?)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use pyo3::types::PyDict;

    fn dict_from_pairs(py: Python<'_>, pairs: &[(&str, &str)]) -> Bound<'_, PyDict> {
        let dict = PyDict::new_bound(py);
        for (k, v) in pairs {
            dict.set_item(*k, *v).unwrap();
        }
        dict
    }


    #[test]
    fn validate_headers_blocks_missing_required() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[("HTTP_USER_AGENT", "Mozilla/5.0")]);
            let result = validate_headers(headers).unwrap();
            assert!(matches!(result, Some(msg) if msg.contains("Missing required headers")));
        });
    }

    #[test]
    fn validate_headers_blocks_suspicious_user_agent() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[
                ("HTTP_USER_AGENT", "python-requests/2.25.1"),
                ("HTTP_ACCEPT", "*/*"),
            ]);
            let result = validate_headers(headers).unwrap();
            assert!(matches!(result, Some(msg) if msg.contains("Suspicious user agent")));
        });
    }

    #[test]
    fn validate_headers_blocks_suspicious_combinations() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[
                ("HTTP_USER_AGENT", "Mozilla/4.0"),
                ("HTTP_ACCEPT", "text/html"),
                ("SERVER_PROTOCOL", "HTTP/2"),
            ]);
            let result = validate_headers(headers).unwrap();
            assert!(matches!(result, Some(msg) if msg.contains("Suspicious headers")));
        });
    }

    #[test]
    fn validate_headers_allows_legit_browser() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[
                ("HTTP_USER_AGENT", "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"),
                ("HTTP_ACCEPT", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"),
                ("HTTP_ACCEPT_LANGUAGE", "en-US,en;q=0.5"),
                ("HTTP_ACCEPT_ENCODING", "gzip, deflate"),
                ("HTTP_CONNECTION", "keep-alive"),
            ]);
            let result = validate_headers(headers).unwrap();
            assert!(result.is_none());
        });
    }

    #[test]
    fn validate_headers_allows_legit_bot() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[
                ("HTTP_USER_AGENT", "Googlebot/2.1 (+http://www.google.com/bot.html)"),
                ("HTTP_ACCEPT", "*/*"),
                ("HTTP_ACCEPT_LANGUAGE", "en-US"),
            ]);
            let result = validate_headers(headers).unwrap();
            assert!(result.is_none());
        });
    }

    #[test]
    fn validate_headers_blocks_accept_star_missing_lang_encoding() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[
                ("HTTP_USER_AGENT", "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"),
                ("HTTP_ACCEPT", "*/*"),
            ]);
            let result = validate_headers(headers).unwrap();
            assert!(matches!(result, Some(msg) if msg.contains("Generic Accept header")));
        });
    }

    #[test]
    fn validate_headers_blocks_http10_chrome() {
        Python::with_gil(|py| {
            let headers = dict_from_pairs(py, &[
                ("HTTP_USER_AGENT", "Mozilla/5.0 Chrome/120.0.0.0"),
                ("HTTP_ACCEPT", "text/html"),
                ("HTTP_ACCEPT_LANGUAGE", "en-US"),
                ("SERVER_PROTOCOL", "HTTP/1.0"),
            ]);
            let result = validate_headers(headers).unwrap();
            assert!(matches!(result, Some(msg) if msg.contains("HTTP/1.0")));
        });
    }

}
