//! Retrieves direct PostgreSQL connection headers.
//!
//! Athena accepts multiple compatibility headers for direct Postgres routing:
//! - `x-pg-uri` (preferred)
//! - `x-athena-jdbc-url` (legacy compatibility)
//! - `x-jdbc-url` / `X-JDBC-URL`
//!
//! JDBC-style URLs (`jdbc:postgresql://...`) are converted to `postgres://...`
//! for sqlx.

use std::net::IpAddr;
use std::str::FromStr;

use actix_web::HttpRequest;
use sha256::digest;
use tokio_postgres::Config as PgConfig;
use tokio_postgres::config::Host;

use super::request_context::disallow_jdbc_routing;

const DIRECT_POSTGRES_HEADERS: [&str; 4] =
    ["x-pg-uri", "x-athena-jdbc-url", "X-JDBC-URL", "x-jdbc-url"];

fn read_direct_postgres_header(req: &HttpRequest) -> Option<String> {
    DIRECT_POSTGRES_HEADERS
        .iter()
        .filter_map(|name| req.headers().get(*name))
        .filter_map(|value| value.to_str().ok())
        .map(str::trim)
        .find(|value| !value.is_empty())
        .map(str::to_string)
}

/// Extracts the JDBC URL from the request, if present.
pub fn x_jdbc_url(req: &HttpRequest) -> Option<String> {
    if disallow_jdbc_routing(req) {
        return None;
    }

    read_direct_postgres_header(req)
}

/// Converts a JDBC PostgreSQL URL to the `postgres://` format expected by sqlx.
pub fn jdbc_to_postgres_url(jdbc_url: &str) -> Option<String> {
    let trimmed: &str = jdbc_url.trim();
    if trimmed.starts_with("jdbc:postgresql://") {
        Some(trimmed.replacen("jdbc:postgresql://", "postgres://", 1))
    } else if trimmed.starts_with("postgres://") || trimmed.starts_with("postgresql://") {
        Some(trimmed.to_string())
    } else {
        None
    }
}

/// True when a parsed Postgres URL contains both username and password credentials.
pub fn direct_postgres_url_has_user_and_password(postgres_url: &str) -> bool {
    let Ok(parsed) = PgConfig::from_str(postgres_url) else {
        return false;
    };
    let has_user = parsed
        .get_user()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .is_some();
    let has_password = parsed
        .get_password()
        .map(|value| !value.is_empty())
        .unwrap_or(false);
    has_user && has_password
}

/// True when the request contains a direct Postgres URI header and embedded credentials.
pub fn request_has_credentialed_direct_postgres_uri(req: &HttpRequest) -> bool {
    let Some(raw_uri) = x_jdbc_url(req) else {
        return false;
    };
    let Some(postgres_url) = jdbc_to_postgres_url(&raw_uri) else {
        return false;
    };
    direct_postgres_url_has_user_and_password(&postgres_url)
}

/// Builds a deterministic opaque routing token for direct Postgres headers.
pub fn direct_postgres_client_token(req: &HttpRequest) -> Option<String> {
    let raw_uri = x_jdbc_url(req)?;
    let postgres_url = jdbc_to_postgres_url(&raw_uri)?;
    let hashed = digest(postgres_url);
    Some(format!("pg_direct_{}", &hashed[..16]))
}

fn host_matches_allowlist(host: &str, allowed_hosts: &[String]) -> bool {
    if allowed_hosts.is_empty() {
        return true;
    }
    let normalized: String = host.trim().to_ascii_lowercase();
    allowed_hosts
        .iter()
        .any(|allowed| normalized == *allowed || normalized.ends_with(&format!(".{}", allowed)))
}

fn host_is_private_or_local(host: &Host) -> bool {
    match host {
        Host::Tcp(hostname) => {
            let normalized: String = hostname.trim().to_ascii_lowercase();
            if normalized == "localhost" {
                return true;
            }
            match IpAddr::from_str(&normalized) {
                Ok(IpAddr::V4(ip)) => {
                    ip.is_loopback()
                        || ip.is_private()
                        || ip.is_link_local()
                        || ip.is_unspecified()
                        || ip.is_multicast()
                }
                Ok(IpAddr::V6(ip)) => {
                    ip.is_loopback()
                        || ip.is_unspecified()
                        || ip.is_multicast()
                        || ip.is_unique_local()
                        || ip.is_unicast_link_local()
                }
                Err(_) => false,
            }
        }
        #[cfg(unix)]
        Host::Unix(_) => true,
    }
}

/// Validates a direct Postgres URI against private-host and allowlist policies.
pub fn validate_postgres_target(
    postgres_url: &str,
    allow_private_hosts: bool,
    allowed_hosts: &[String],
) -> Result<(), String> {
    let parsed = PgConfig::from_str(postgres_url)
        .map_err(|_| "Direct PostgreSQL URI could not be parsed as a connection string")?;

    let hosts = parsed.get_hosts();
    if hosts.is_empty() {
        return Err("Direct PostgreSQL URI must include a host".to_string());
    }

    for host in hosts {
        match host {
            Host::Tcp(hostname) => {
                if !host_matches_allowlist(hostname, allowed_hosts) {
                    return Err(
                        "Direct PostgreSQL host is not in the configured JDBC host allowlist"
                            .to_string(),
                    );
                }

                if !allow_private_hosts && host_is_private_or_local(host) {
                    return Err(
                        "Direct PostgreSQL private or local hosts are blocked by policy"
                            .to_string(),
                    );
                }
            }
            #[cfg(unix)]
            Host::Unix(_) => {
                return Err("Direct PostgreSQL unix socket hosts are not allowed".to_string());
            }
        }
    }

    Ok(())
}

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

    #[test]
    fn test_jdbc_to_postgres_url() {
        assert_eq!(
            jdbc_to_postgres_url("jdbc:postgresql://localhost:5432/mydb"),
            Some("postgres://localhost:5432/mydb".to_string())
        );
        assert_eq!(
            jdbc_to_postgres_url("jdbc:postgresql://user:pass@host/db"),
            Some("postgres://user:pass@host/db".to_string())
        );
        assert_eq!(
            jdbc_to_postgres_url("postgres://localhost/db"),
            Some("postgres://localhost/db".to_string())
        );
        assert_eq!(jdbc_to_postgres_url("jdbc:mysql://localhost/db"), None);
    }

    #[test]
    fn request_parser_prefers_x_pg_uri_header() {
        let req = actix_web::test::TestRequest::default()
            .insert_header(("x-pg-uri", "postgres://user:pass@host/db"))
            .insert_header(("x-jdbc-url", "jdbc:postgresql://ignored:ignored@other/db"))
            .to_http_request();
        assert_eq!(
            x_jdbc_url(&req),
            Some("postgres://user:pass@host/db".to_string())
        );
    }

    #[test]
    fn request_parser_accepts_x_athena_jdbc_url() {
        let req = actix_web::test::TestRequest::default()
            .insert_header(("x-athena-jdbc-url", "jdbc:postgresql://user:pass@host/db"))
            .to_http_request();
        assert_eq!(
            x_jdbc_url(&req),
            Some("jdbc:postgresql://user:pass@host/db".to_string())
        );
    }

    #[test]
    fn credential_detection_requires_user_and_password() {
        assert!(direct_postgres_url_has_user_and_password(
            "postgres://user:pass@host:5432/db"
        ));
        assert!(!direct_postgres_url_has_user_and_password(
            "postgres://user@host:5432/db"
        ));
        assert!(!direct_postgres_url_has_user_and_password(
            "postgres://host:5432/db"
        ));
    }

    #[test]
    fn direct_client_token_is_stable_for_same_uri() {
        let req_a = actix_web::test::TestRequest::default()
            .insert_header(("x-pg-uri", "postgres://user:pass@host/db"))
            .to_http_request();
        let req_b = actix_web::test::TestRequest::default()
            .insert_header(("x-pg-uri", "postgres://user:pass@host/db"))
            .to_http_request();
        assert_eq!(
            direct_postgres_client_token(&req_a),
            direct_postgres_client_token(&req_b)
        );
    }

    #[test]
    fn test_validate_postgres_target_rejects_local_when_private_blocked() {
        let result: Result<(), String> =
            validate_postgres_target("postgres://localhost:5432/test", false, &[]);
        assert!(result.is_err());
    }

    #[test]
    fn test_validate_postgres_target_allows_allowlisted_host() {
        let allowed: Vec<String> = vec!["db.example.com".to_string()];
        let result = validate_postgres_target(
            "postgres://user:pass@db.example.com:5432/test",
            false,
            &allowed,
        );
        assert!(result.is_ok());
    }

    #[test]
    fn test_validate_postgres_target_rejects_non_allowlisted_host() {
        let allowed: Vec<String> = vec!["db.example.com".to_string()];
        let result = validate_postgres_target(
            "postgres://user:pass@other.example.com:5432/test",
            true,
            &allowed,
        );
        assert!(result.is_err());
    }
}
