use std::net::{SocketAddr, TcpListener};
use std::path::{Path, PathBuf};
use std::sync::Once;

use axum::http::header::HeaderName;
use axum::http::HeaderValue;
use axum::response::Response;
use axum::Router;
use axum_server::tls_rustls::RustlsConfig;

const DEFAULT_HSTS_MAX_AGE_SECONDS: u64 = 31_536_000;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HstsConfig {
    pub enabled: bool,
    pub max_age_seconds: u64,
    pub include_subdomains: bool,
    pub preload: bool,
}

impl HstsConfig {
    pub fn enabled() -> Self {
        Self {
            enabled: true,
            max_age_seconds: DEFAULT_HSTS_MAX_AGE_SECONDS,
            include_subdomains: false,
            preload: false,
        }
    }

    pub fn disabled() -> Self {
        Self {
            enabled: false,
            max_age_seconds: 0,
            include_subdomains: false,
            preload: false,
        }
    }

    fn header_value(&self) -> Option<HeaderValue> {
        if !self.enabled {
            return None;
        }
        let mut value = format!("max-age={}", self.max_age_seconds);
        if self.include_subdomains {
            value.push_str("; includeSubDomains");
        }
        if self.preload {
            value.push_str("; preload");
        }
        HeaderValue::from_str(&value).ok()
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum HttpTlsConfig {
    #[default]
    Plain,
    EdgeTerminated {
        hsts: HstsConfig,
    },
    PemFiles {
        cert: PathBuf,
        key: PathBuf,
        hsts: HstsConfig,
    },
    SelfSignedDev {
        hosts: Vec<String>,
    },
}

impl HttpTlsConfig {
    pub fn plain() -> Self {
        Self::Plain
    }

    pub fn edge_terminated() -> Self {
        Self::EdgeTerminated {
            hsts: HstsConfig::enabled(),
        }
    }

    pub fn pem_files(cert: impl Into<PathBuf>, key: impl Into<PathBuf>) -> Self {
        Self::PemFiles {
            cert: cert.into(),
            key: key.into(),
            hsts: HstsConfig::enabled(),
        }
    }

    pub fn self_signed_dev() -> Self {
        Self::SelfSignedDev {
            hosts: vec!["localhost".to_string(), "127.0.0.1".to_string()],
        }
    }

    pub fn listener_scheme(&self) -> &'static str {
        match self {
            Self::Plain | Self::EdgeTerminated { .. } => "http",
            Self::PemFiles { .. } | Self::SelfSignedDev { .. } => "https",
        }
    }

    pub fn advertised_scheme(&self) -> &'static str {
        match self {
            Self::Plain => "http",
            Self::EdgeTerminated { .. } | Self::PemFiles { .. } | Self::SelfSignedDev { .. } => {
                "https"
            }
        }
    }

    pub fn is_edge_terminated(&self) -> bool {
        matches!(self, Self::EdgeTerminated { .. })
    }

    pub fn security_headers(&self) -> Vec<(HeaderName, HeaderValue)> {
        let hsts = match self {
            Self::EdgeTerminated { hsts } | Self::PemFiles { hsts, .. } => hsts,
            Self::Plain | Self::SelfSignedDev { .. } => return Vec::new(),
        };
        hsts.header_value()
            .map(|value| (HeaderName::from_static("strict-transport-security"), value))
            .into_iter()
            .collect()
    }

    async fn rustls_config(&self) -> Result<Option<RustlsConfig>, String> {
        match self {
            Self::Plain | Self::EdgeTerminated { .. } => Ok(None),
            Self::PemFiles { cert, key, .. } => load_rustls_config(cert, key).await.map(Some),
            Self::SelfSignedDev { hosts } => self_signed_rustls_config(hosts).await.map(Some),
        }
    }
}

pub async fn serve_router_from_tcp(
    listener: TcpListener,
    router: Router,
    tls: &HttpTlsConfig,
) -> Result<(), String> {
    // `with_connect_info` stashes the peer `SocketAddr` in each request's
    // extensions so handlers can recover the real transport peer (e.g. the
    // site adapter's `req.remote_addr` and trusted-proxy `client_ip`
    // resolution). Adapters that don't read it pay nothing.
    match tls.rustls_config().await? {
        Some(config) => axum_server::from_tcp_rustls(listener, config)
            .map_err(|error| format!("HTTPS listener setup failed: {error}"))?
            .serve(router.into_make_service_with_connect_info::<SocketAddr>())
            .await
            .map_err(|error| format!("HTTPS listener failed: {error}")),
        None => axum_server::from_tcp(listener)
            .map_err(|error| format!("HTTP listener setup failed: {error}"))?
            .serve(router.into_make_service_with_connect_info::<SocketAddr>())
            .await
            .map_err(|error| format!("HTTP listener failed: {error}")),
    }
}

pub fn bind_listener(bind: SocketAddr) -> Result<TcpListener, String> {
    let listener =
        TcpListener::bind(bind).map_err(|error| format!("failed to bind {bind}: {error}"))?;
    listener
        .set_nonblocking(true)
        .map_err(|error| format!("failed to enable nonblocking listener mode: {error}"))?;
    Ok(listener)
}

pub fn apply_security_headers(router: Router, tls: &HttpTlsConfig) -> Router {
    let headers = tls.security_headers();
    if headers.is_empty() {
        return router;
    }
    router.layer(axum::middleware::map_response(
        move |mut response: Response| {
            let headers = headers.clone();
            async move {
                for (name, value) in headers {
                    response.headers_mut().insert(name, value);
                }
                response
            }
        },
    ))
}

async fn load_rustls_config(cert: &Path, key: &Path) -> Result<RustlsConfig, String> {
    install_crypto_provider();
    if !cert.is_file() {
        return Err(format!("TLS certificate not found: {}", cert.display()));
    }
    if !key.is_file() {
        return Err(format!("TLS private key not found: {}", key.display()));
    }
    RustlsConfig::from_pem_file(cert, key)
        .await
        .map_err(|error| {
            format!(
                "failed to load TLS certificate {} and key {}: {error}",
                cert.display(),
                key.display()
            )
        })
}

async fn self_signed_rustls_config(hosts: &[String]) -> Result<RustlsConfig, String> {
    install_crypto_provider();
    let hosts = if hosts.is_empty() {
        vec!["localhost".to_string(), "127.0.0.1".to_string()]
    } else {
        hosts.to_vec()
    };
    let cert = rcgen::generate_simple_self_signed(hosts)
        .map_err(|error| format!("failed to generate self-signed dev certificate: {error}"))?;
    RustlsConfig::from_pem(
        cert.cert.pem().into_bytes(),
        cert.signing_key.serialize_pem().into_bytes(),
    )
    .await
    .map_err(|error| format!("failed to load self-signed dev certificate: {error}"))
}

fn install_crypto_provider() {
    static INSTALL: Once = Once::new();
    INSTALL.call_once(|| {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
    });
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::routing::get;
    use tempfile::TempDir;
    use tokio::sync::oneshot;

    #[tokio::test]
    async fn pem_config_reports_missing_files_before_serving() {
        let temp = TempDir::new().expect("tempdir");
        let tls =
            HttpTlsConfig::pem_files(temp.path().join("missing.pem"), temp.path().join("key.pem"));

        let error = tls
            .rustls_config()
            .await
            .expect_err("missing files should fail");

        assert!(error.contains("TLS certificate not found"), "error={error}");
    }

    #[tokio::test]
    async fn pem_config_reports_invalid_pem_before_serving() {
        let temp = TempDir::new().expect("tempdir");
        let cert = temp.path().join("cert.pem");
        let key = temp.path().join("key.pem");
        std::fs::write(&cert, "not a cert").expect("write cert");
        std::fs::write(&key, "not a key").expect("write key");
        let tls = HttpTlsConfig::pem_files(&cert, &key);

        let error = tls
            .rustls_config()
            .await
            .expect_err("invalid files should fail");

        assert!(
            error.contains("failed to load TLS certificate"),
            "error={error}"
        );
    }

    #[tokio::test]
    async fn self_signed_dev_starts_https_listener() {
        let listener = bind_listener("127.0.0.1:0".parse().unwrap()).expect("listener");
        let addr = listener.local_addr().expect("addr");
        let (ready_tx, ready_rx) = oneshot::channel();
        let app = Router::new().route("/health", get(|| async { "ok" }));
        let tls = HttpTlsConfig::self_signed_dev();

        let task = tokio::spawn(async move {
            let _ = ready_tx.send(());
            serve_router_from_tcp(listener, app, &tls).await
        });
        ready_rx.await.expect("ready");

        let client = reqwest::Client::builder()
            .danger_accept_invalid_certs(true)
            .build()
            .expect("client");
        let response = client
            .get(format!("https://{addr}/health"))
            .send()
            .await
            .expect("request");
        assert!(response.status().is_success());

        task.abort();
    }

    #[tokio::test]
    async fn plain_and_edge_modes_start_plain_listener() {
        for tls in [HttpTlsConfig::plain(), HttpTlsConfig::edge_terminated()] {
            let listener = bind_listener("127.0.0.1:0".parse().unwrap()).expect("listener");
            let addr = listener.local_addr().expect("addr");
            let (ready_tx, ready_rx) = oneshot::channel();
            let app = apply_security_headers(
                Router::new().route("/health", get(|| async { "ok" })),
                &tls,
            );
            let is_edge = tls.is_edge_terminated();

            let task = tokio::spawn(async move {
                let _ = ready_tx.send(());
                serve_router_from_tcp(listener, app, &tls).await
            });
            ready_rx.await.expect("ready");

            let response = reqwest::get(format!("http://{addr}/health"))
                .await
                .expect("request");
            assert!(response.status().is_success());
            if is_edge {
                assert!(response.headers().contains_key("strict-transport-security"));
            } else {
                assert!(!response.headers().contains_key("strict-transport-security"));
            }

            task.abort();
        }
    }
}
