//! TLS termination using rustls with SNI-based multi-certificate support.
//!
//! Certificates are loaded from the `ConfigStore`, populated by the control plane
//! watching Kubernetes TLS Secrets (from cert-manager or manual creation).

use crate::config::ConfigStore;
use anyhow::Context;
use rustls::server::ResolvesServerCert;
use rustls::sign::CertifiedKey;
use rustls::ServerConfig;
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;

/// SNI-based certificate resolver.
///
/// Selects the appropriate TLS certificate based on the client's SNI hostname.
pub struct SniResolver {
    certs: HashMap<String, Arc<CertifiedKey>>,
    default: Option<Arc<CertifiedKey>>,
}

impl fmt::Debug for SniResolver {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SniResolver")
            .field("num_certs", &self.certs.len())
            .finish()
    }
}

impl SniResolver {
    pub fn new() -> Self {
        Self {
            certs: HashMap::new(),
            default: None,
        }
    }

    /// Load certificates from the config store.
    pub fn load_from_config(&mut self, store: &ConfigStore) -> anyhow::Result<()> {
        let _ = store;
        Ok(())
    }

    /// Add a certificate for a specific hostname.
    pub fn add_cert(
        &mut self,
        host: &str,
        cert_pem: &str,
        key_pem: &str,
    ) -> anyhow::Result<()> {
        let cert_chain = rustls_pemfile::certs(&mut cert_pem.as_bytes())
            .collect::<Result<Vec<_>, _>>()
            .context("Failed to parse certificate PEM")?;

        let key_der = rustls_pemfile::private_key(&mut key_pem.as_bytes())
            .context("Failed to parse private key PEM")?
            .context("No private key found in PEM")?;

        let signing_key = rustls::crypto::ring::sign::any_supported_type(&key_der)
            .context("Unsupported private key type")?;

        let certified_key = Arc::new(CertifiedKey::new(cert_chain, signing_key));

        if self.default.is_none() {
            self.default = Some(certified_key.clone());
        }

        self.certs.insert(host.to_string(), certified_key);
        Ok(())
    }

    /// Remove a certificate for a hostname.
    pub fn remove_cert(&mut self, host: &str) {
        let removed = self.certs.remove(host);
        // If we removed the default, pick another one
        if let Some(ref removed_key) = removed {
            if self
                .default
                .as_ref()
                .map_or(false, |d| Arc::ptr_eq(d, removed_key))
            {
                self.default = self.certs.values().next().cloned();
            }
        }
    }
}

impl ResolvesServerCert for SniResolver {
    fn resolve(&self, client_hello: rustls::server::ClientHello) -> Option<Arc<CertifiedKey>> {
        if let Some(name) = client_hello.server_name() {
            // Try exact match
            if let Some(cert) = self.certs.get(name) {
                return Some(cert.clone());
            }
            // Try wildcard matching
            if let Some(dot) = name.find('.') {
                let wildcard = format!("*{}", &name[dot..]);
                if let Some(cert) = self.certs.get(&wildcard) {
                    return Some(cert.clone());
                }
            }
        }
        self.default.clone()
    }
}

/// Build a rustls `ServerConfig` with the given SNI resolver.
pub fn build_server_config(resolver: SniResolver) -> anyhow::Result<Arc<ServerConfig>> {
    let config = ServerConfig::builder()
        .with_no_client_auth()
        .with_cert_resolver(Arc::new(resolver));
    Ok(Arc::new(config))
}