//! SNI-based certificate resolver for multi-domain TLS.
//!
//! Implements the `rustls::server::ResolvesServerCert` trait so that
//! different domains served by the same listener can present their own
//! certificates.  Supports exact-match, wildcard (`*.example.com`),
//! and a default fallback certificate.

use anyhow::{Context, Result};
use dashmap::DashMap;
use rustls::server::{ClientHello, ResolvesServerCert};
use rustls::sign::CertifiedKey;
use rustls_pki_types::{CertificateDer, PrivateKeyDer};
use std::sync::{Arc, RwLock};

/// SNI-based certificate resolver.
///
/// Thread-safe; certificates can be hot-swapped at runtime.
///
/// When an `OnDemandProvisioner` is attached, unknown domains that pass
/// validation will be queued for background ACME provisioning.  The
/// default fallback cert is returned immediately while provisioning
/// happens asynchronously; once complete, `add_cert()` installs the
/// real cert so subsequent requests use it.
pub struct SniCertResolver {
    /// Exact domain → certified key.
    certs: DashMap<String, Arc<CertifiedKey>>,
    /// Default fallback certificate for unmatched domains.
    default: RwLock<Option<Arc<CertifiedKey>>>,
    /// Optional on-demand provisioner for unknown domains.
    on_demand: Option<Arc<crate::on_demand::OnDemandProvisioner>>,
}

impl std::fmt::Debug for SniCertResolver {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SniCertResolver")
            .field("certs_count", &self.certs.len())
            .field(
                "has_default",
                &self
                    .default
                    .read()
                    .ok()
                    .map(|g| g.is_some())
                    .unwrap_or(false),
            )
            .field("has_on_demand", &self.on_demand.is_some())
            .finish()
    }
}

impl SniCertResolver {
    /// Create a new empty resolver.
    pub fn new() -> Self {
        Self {
            certs: DashMap::new(),
            default: RwLock::new(None),
            on_demand: None,
        }
    }

    /// Attach an on-demand TLS provisioner.
    ///
    /// When `resolve()` cannot find a cert for a domain, the provisioner
    /// is consulted.  If the domain passes validation, a background ACME
    /// task is triggered and the default cert is returned in the meantime.
    pub fn with_on_demand(
        mut self,
        provisioner: Arc<crate::on_demand::OnDemandProvisioner>,
    ) -> Self {
        self.on_demand = Some(provisioner);
        self
    }

    /// Add (or replace) a certificate for `domain`.
    ///
    /// The domain may be a wildcard pattern like `*.example.com`.
    pub fn add_cert(&self, domain: &str, cert_pem: &[u8], key_pem: &[u8]) -> Result<()> {
        let certified_key = build_certified_key(cert_pem, key_pem)
            .with_context(|| format!("build certified key for {domain}"))?;

        let key = Arc::new(certified_key);
        self.certs.insert(domain.to_lowercase(), key);

        tracing::info!(domain, "certificate loaded into SNI resolver");
        Ok(())
    }

    /// Remove the certificate for `domain`.
    pub fn remove_cert(&self, domain: &str) {
        self.certs.remove(&domain.to_lowercase());
        tracing::info!(domain, "certificate removed from SNI resolver");
    }

    /// Set the default fallback certificate for unmatched domains.
    pub fn set_default(&self, cert_pem: &[u8], key_pem: &[u8]) -> Result<()> {
        let certified_key =
            build_certified_key(cert_pem, key_pem).context("build default certified key")?;

        let mut default = self.default.write().map_err(|e| anyhow::anyhow!("{e}"))?;
        *default = Some(Arc::new(certified_key));

        tracing::info!("default fallback certificate set");
        Ok(())
    }

    /// Check whether a certificate exists for the given domain.
    pub fn has_cert(&self, domain: &str) -> bool {
        self.certs.contains_key(&domain.to_lowercase())
    }

    /// Return the number of loaded certificates (excluding default).
    pub fn len(&self) -> usize {
        self.certs.len()
    }

    /// Check if the resolver has no certificates (excluding default).
    pub fn is_empty(&self) -> bool {
        self.certs.is_empty()
    }

    /// Resolve a server name to a `CertifiedKey`.
    ///
    /// Resolution order:
    /// 1. Exact match (case-insensitive)
    /// 2. Wildcard match (`*.example.com`)
    /// 3. Default fallback
    /// 4. On-demand provisioning trigger (returns default while provisioning)
    fn resolve(&self, server_name: &str) -> Option<Arc<CertifiedKey>> {
        let name = server_name.to_lowercase();

        // 1. Exact match.
        if let Some(key) = self.certs.get(&name) {
            return Some(Arc::clone(key.value()));
        }

        // 2. Wildcard match: if name is `foo.example.com`, check `*.example.com`.
        if let Some(dot_pos) = name.find('.') {
            let wildcard = format!("*.{}", &name[dot_pos + 1..]);
            if let Some(key) = self.certs.get(&wildcard) {
                return Some(Arc::clone(key.value()));
            }
        }

        // 3. Default fallback.
        let default = self
            .default
            .read()
            .ok()
            .and_then(|guard| guard.as_ref().map(Arc::clone));

        // 4. On-demand: trigger background provisioning for unknown domains.
        //    The actual ACME provisioning is async and cannot happen during
        //    the synchronous resolve() call.  We return the default cert
        //    immediately and the provisioner spawns a background task.
        //    Once provisioning completes, add_cert() installs the real cert
        //    so future resolve() calls get an exact match.
        if default.is_some() {
            if let Some(ref provisioner) = self.on_demand {
                if provisioner.should_provision(&name) && provisioner.start_provisioning(&name) {
                    tracing::info!(domain = %name, "on-demand TLS: queued for ACME provisioning");
                }
            }
        }

        default
    }
}

impl Default for SniCertResolver {
    fn default() -> Self {
        Self::new()
    }
}

impl ResolvesServerCert for SniCertResolver {
    fn resolve(&self, client_hello: ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
        let server_name = client_hello.server_name()?;
        self.resolve(server_name)
    }
}

/// Parse PEM-encoded certificate(s) and private key into a `CertifiedKey`.
fn build_certified_key(cert_pem: &[u8], key_pem: &[u8]) -> Result<CertifiedKey> {
    // Parse certificates.
    let certs = parse_pem_certs(cert_pem).context("parse certificate PEM")?;
    if certs.is_empty() {
        anyhow::bail!("no certificates found in PEM data");
    }

    // Parse private key.
    let key = parse_private_key(key_pem).context("parse private key PEM")?;

    // Build the signing key.
    let signing_key = rustls::crypto::ring::sign::any_supported_type(&key)
        .map_err(|e| anyhow::anyhow!("unsupported key type: {e}"))?;

    Ok(CertifiedKey::new(certs, signing_key))
}

/// Parse PEM-encoded certificates into DER blobs.
fn parse_pem_certs(pem: &[u8]) -> Result<Vec<CertificateDer<'static>>> {
    let mut reader = std::io::BufReader::new(pem);
    let certs: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut reader)
        .collect::<std::result::Result<Vec<_>, _>>()
        .context("read PEM certificates")?;
    Ok(certs)
}

/// Parse a PEM-encoded private key (PKCS#8 or PKCS#1 or SEC1).
fn parse_private_key(pem: &[u8]) -> Result<PrivateKeyDer<'static>> {
    let mut reader = std::io::BufReader::new(pem);

    // Try all key formats.
    loop {
        match rustls_pemfile::read_one(&mut reader).context("read PEM key item")? {
            Some(rustls_pemfile::Item::Pkcs8Key(key)) => {
                return Ok(PrivateKeyDer::Pkcs8(key));
            }
            Some(rustls_pemfile::Item::Pkcs1Key(key)) => {
                return Ok(PrivateKeyDer::Pkcs1(key));
            }
            Some(rustls_pemfile::Item::Sec1Key(key)) => {
                return Ok(PrivateKeyDer::Sec1(key));
            }
            Some(_) => {
                // Skip non-key items (e.g., certificates mixed in).
                continue;
            }
            None => {
                anyhow::bail!("no private key found in PEM data");
            }
        }
    }
}

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

    fn gen_cert(domains: &[&str]) -> (Vec<u8>, Vec<u8>) {
        let ds: Vec<String> = domains.iter().map(|s| s.to_string()).collect();
        let out = generate_self_signed(&ds).unwrap();
        (out.cert_pem, out.key_pem)
    }

    #[test]
    fn new_resolver_is_empty() {
        let resolver = SniCertResolver::new();
        assert!(resolver.is_empty());
        assert_eq!(resolver.len(), 0);
    }

    #[test]
    fn add_and_resolve_exact() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["example.com"]);
        resolver.add_cert("example.com", &cert, &key).unwrap();

        assert!(!resolver.is_empty());
        assert_eq!(resolver.len(), 1);
        assert!(resolver.has_cert("example.com"));

        let resolved = resolver.resolve("example.com");
        assert!(resolved.is_some());
    }

    #[test]
    fn case_insensitive_match() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["Example.COM"]);
        resolver.add_cert("Example.COM", &cert, &key).unwrap();

        // Stored as lowercase.
        assert!(resolver.has_cert("example.com"));
        let resolved = resolver.resolve("EXAMPLE.COM");
        assert!(resolved.is_some());
    }

    #[test]
    fn wildcard_match() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["*.example.com"]);
        resolver.add_cert("*.example.com", &cert, &key).unwrap();

        // Should match subdomain.
        let resolved = resolver.resolve("foo.example.com");
        assert!(resolved.is_some());

        let resolved = resolver.resolve("bar.example.com");
        assert!(resolved.is_some());
    }

    #[test]
    fn wildcard_does_not_match_base() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["*.example.com"]);
        resolver.add_cert("*.example.com", &cert, &key).unwrap();

        // Wildcard should NOT match the base domain itself.
        let resolved = resolver.resolve("example.com");
        assert!(resolved.is_none());
    }

    #[test]
    fn wildcard_does_not_match_nested_subdomain() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["*.example.com"]);
        resolver.add_cert("*.example.com", &cert, &key).unwrap();

        // `a.b.example.com` — the wildcard `*.example.com` should match
        // `b.example.com` but for `a.b.example.com` the first label is `a`
        // and the parent is `b.example.com`, not `example.com`.
        // Our implementation strips the first label and checks `*.b.example.com`.
        let resolved = resolver.resolve("a.b.example.com");
        assert!(resolved.is_none());
    }

    #[test]
    fn exact_match_takes_priority_over_wildcard() {
        let resolver = SniCertResolver::new();
        let (cert1, key1) = gen_cert(&["app.example.com"]);
        let (cert2, key2) = gen_cert(&["*.example.com"]);

        resolver.add_cert("app.example.com", &cert1, &key1).unwrap();
        resolver.add_cert("*.example.com", &cert2, &key2).unwrap();

        // Should resolve via exact match.
        let resolved = resolver.resolve("app.example.com");
        assert!(resolved.is_some());
    }

    #[test]
    fn fallback_default() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["fallback.local"]);
        resolver.set_default(&cert, &key).unwrap();

        // No exact or wildcard match → should fall back to default.
        let resolved = resolver.resolve("unknown.domain.com");
        assert!(resolved.is_some());
    }

    #[test]
    fn no_match_no_default_returns_none() {
        let resolver = SniCertResolver::new();
        let resolved = resolver.resolve("anything.com");
        assert!(resolved.is_none());
    }

    #[test]
    fn remove_cert() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["remove.me"]);
        resolver.add_cert("remove.me", &cert, &key).unwrap();
        assert!(resolver.has_cert("remove.me"));

        resolver.remove_cert("remove.me");
        assert!(!resolver.has_cert("remove.me"));
        assert!(resolver.is_empty());
    }

    #[test]
    fn replace_cert() {
        let resolver = SniCertResolver::new();
        let (cert1, key1) = gen_cert(&["replace.me"]);
        let (cert2, key2) = gen_cert(&["replace.me"]);

        resolver.add_cert("replace.me", &cert1, &key1).unwrap();
        resolver.add_cert("replace.me", &cert2, &key2).unwrap();

        // Still only one entry.
        assert_eq!(resolver.len(), 1);
        assert!(resolver.resolve("replace.me").is_some());
    }

    #[test]
    fn multiple_domains() {
        let resolver = SniCertResolver::new();
        let (cert1, key1) = gen_cert(&["a.com"]);
        let (cert2, key2) = gen_cert(&["b.com"]);
        let (cert3, key3) = gen_cert(&["c.com"]);

        resolver.add_cert("a.com", &cert1, &key1).unwrap();
        resolver.add_cert("b.com", &cert2, &key2).unwrap();
        resolver.add_cert("c.com", &cert3, &key3).unwrap();

        assert_eq!(resolver.len(), 3);
        assert!(resolver.resolve("a.com").is_some());
        assert!(resolver.resolve("b.com").is_some());
        assert!(resolver.resolve("c.com").is_some());
        assert!(resolver.resolve("d.com").is_none());
    }

    #[test]
    fn invalid_cert_pem_errors() {
        let resolver = SniCertResolver::new();
        let result = resolver.add_cert("bad.com", b"not a cert", b"not a key");
        assert!(result.is_err());
    }

    #[test]
    fn invalid_key_pem_errors() {
        let (cert, _key) = gen_cert(&["test.com"]);
        let resolver = SniCertResolver::new();
        let result = resolver.add_cert("test.com", &cert, b"not a real key");
        assert!(result.is_err());
    }

    #[test]
    fn default_trait_impl() {
        let resolver = SniCertResolver::default();
        assert!(resolver.is_empty());
    }

    #[test]
    fn wildcard_case_insensitive() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["*.EXAMPLE.COM"]);
        resolver.add_cert("*.EXAMPLE.COM", &cert, &key).unwrap();

        let resolved = resolver.resolve("foo.example.com");
        assert!(resolved.is_some());
    }

    #[test]
    fn set_default_replaces() {
        let resolver = SniCertResolver::new();
        let (cert1, key1) = gen_cert(&["default1"]);
        let (cert2, key2) = gen_cert(&["default2"]);

        resolver.set_default(&cert1, &key1).unwrap();
        resolver.set_default(&cert2, &key2).unwrap();

        // Should still resolve to *some* default.
        let resolved = resolver.resolve("unknown.com");
        assert!(resolved.is_some());
    }

    #[test]
    fn resolve_empty_server_name() {
        let resolver = SniCertResolver::new();
        let (cert, key) = gen_cert(&["default"]);
        resolver.set_default(&cert, &key).unwrap();

        // Empty string — no exact/wildcard match, falls back to default.
        let resolved = resolver.resolve("");
        assert!(resolved.is_some());
    }

    // --- On-demand provisioning integration tests ---

    #[test]
    fn on_demand_triggers_provisioning() {
        let provisioner = Arc::new(
            crate::on_demand::OnDemandProvisioner::new(100, None)
                .with_validator(|d| d.ends_with(".example.com")),
        );
        let resolver = SniCertResolver::new().with_on_demand(Arc::clone(&provisioner));

        let (cert, key) = gen_cert(&["fallback.local"]);
        resolver.set_default(&cert, &key).unwrap();

        // No exact match for "new.example.com", should trigger provisioning
        // and return the default cert.
        let resolved = resolver.resolve("new.example.com");
        assert!(resolved.is_some(), "should return default cert");

        // Provisioner should have "new.example.com" in-flight.
        assert_eq!(provisioner.in_flight_count(), 1);
    }

    #[test]
    fn on_demand_skips_rejected_domains() {
        let provisioner = Arc::new(
            crate::on_demand::OnDemandProvisioner::new(100, None)
                .with_validator(|d| d.ends_with(".allowed.com")),
        );
        let resolver = SniCertResolver::new().with_on_demand(Arc::clone(&provisioner));

        let (cert, key) = gen_cert(&["fallback.local"]);
        resolver.set_default(&cert, &key).unwrap();

        // Domain doesn't pass validation — still returns default but
        // should NOT trigger provisioning.
        let resolved = resolver.resolve("evil.attacker.com");
        assert!(resolved.is_some(), "should return default cert");
        assert_eq!(
            provisioner.in_flight_count(),
            0,
            "should not be provisioned"
        );
    }

    #[test]
    fn on_demand_no_trigger_without_default() {
        let provisioner = Arc::new(crate::on_demand::OnDemandProvisioner::new(100, None));
        let resolver = SniCertResolver::new().with_on_demand(Arc::clone(&provisioner));

        // No default cert set — on-demand should NOT trigger (nothing
        // to return to the client).
        let resolved = resolver.resolve("any.com");
        assert!(resolved.is_none());
        assert_eq!(provisioner.in_flight_count(), 0);
    }

    #[test]
    fn on_demand_exact_match_skips_provisioning() {
        let provisioner = Arc::new(crate::on_demand::OnDemandProvisioner::new(100, None));
        let resolver = SniCertResolver::new().with_on_demand(Arc::clone(&provisioner));

        let (cert, key) = gen_cert(&["known.com"]);
        resolver.add_cert("known.com", &cert, &key).unwrap();

        // Exact match exists — should NOT trigger provisioning.
        let resolved = resolver.resolve("known.com");
        assert!(resolved.is_some());
        assert_eq!(provisioner.in_flight_count(), 0);
    }

    #[test]
    fn debug_impl() {
        let resolver = SniCertResolver::new();
        let debug_str = format!("{resolver:?}");
        assert!(debug_str.contains("SniCertResolver"));
        assert!(debug_str.contains("certs_count"));
    }
}
