//! TLS termination with SNI-based certificate selection
//!
//! This module provides TLS termination using rustls with support for:
//! - SNI-based certificate selection (exact and wildcard matching)
//! - Session resumption
//! - Modern cipher suites
//! - TLS 1.2 and 1.3
//! - ALPN negotiation for HTTP/1.1 and HTTP/2
//!
//! ## Security: Certificate Privacy
//!
//! Certificates are only returned when the SNI exactly matches a configured hostname.
//! This prevents information disclosure about which domains are hosted on the server.
//! Connections without matching SNI (including IP address connections) will fail the
//! TLS handshake rather than leak certificate information.

use dashmap::DashMap;
use parking_lot::RwLock;
use rustls::pki_types::{CertificateDer, PrivateKeyDer};
use rustls::server::{ClientHello, ResolvesServerCert};
use rustls::ServerConfig;
use std::fmt;
use std::fs::File;
use std::io::{BufReader, Cursor};
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use tracing::info;

/// ALPN protocol identifiers
pub const ALPN_H2: &[u8] = b"h2";
pub const ALPN_HTTP11: &[u8] = b"http/1.1";

/// Negotiated protocol after TLS handshake
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NegotiatedProtocol {
    H2,
    Http1,
}

impl NegotiatedProtocol {
    /// Determine protocol from ALPN negotiation result
    pub fn from_alpn(alpn: Option<&[u8]>) -> Self {
        match alpn {
            Some(ALPN_H2) => NegotiatedProtocol::H2,
            _ => NegotiatedProtocol::Http1,
        }
    }
}

/// TLS configuration manager
pub struct TlsManager {
    /// Certificate resolver
    resolver: Arc<SniResolver>,
    /// Statistics
    stats: TlsStats,
}

/// TLS statistics
#[derive(Default)]
pub struct TlsStats {
    pub handshakes_started: AtomicU64,
    pub handshakes_completed: AtomicU64,
    pub handshakes_failed: AtomicU64,
}

/// SNI-based certificate resolver
pub struct SniResolver {
    /// Certified keys indexed by hostname
    certs: DashMap<String, Arc<rustls::sign::CertifiedKey>>,
    /// Default certificate for unknown SNI (None by default for security)
    default_cert: RwLock<Option<Arc<rustls::sign::CertifiedKey>>>,
    /// Whether to use default certificate fallback (false by default to prevent info disclosure)
    use_default_fallback: bool,
}

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

impl SniResolver {
    /// Create a new SNI resolver with default fallback disabled (secure by default)
    pub fn new() -> Self {
        Self {
            certs: DashMap::new(),
            default_cert: RwLock::new(None),
            use_default_fallback: false,
        }
    }

    /// Create a new SNI resolver with configurable default fallback behavior
    pub fn with_default_fallback(use_fallback: bool) -> Self {
        Self {
            certs: DashMap::new(),
            default_cert: RwLock::new(None),
            use_default_fallback: use_fallback,
        }
    }

    /// Add a certificate for a hostname
    pub fn add_cert(&self, hostname: &str, certified_key: Arc<rustls::sign::CertifiedKey>) {
        self.certs.insert(hostname.to_lowercase(), certified_key);
    }

    /// Set the default certificate (only used if use_default_fallback is true)
    pub fn set_default(&self, certified_key: Arc<rustls::sign::CertifiedKey>) {
        *self.default_cert.write() = Some(certified_key);
    }
}

impl ResolvesServerCert for SniResolver {
    fn resolve(&self, client_hello: ClientHello<'_>) -> Option<Arc<rustls::sign::CertifiedKey>> {
        let sni = client_hello.server_name()?;
        let sni_lower = sni.to_lowercase();

        // Try exact match first
        if let Some(cert) = self.certs.get(&sni_lower) {
            return Some(cert.clone());
        }

        // Try wildcard match
        if let Some(dot_pos) = sni_lower.find('.') {
            let wildcard = format!("*{}", &sni_lower[dot_pos..]);
            if let Some(cert) = self.certs.get(&wildcard) {
                return Some(cert.clone());
            }
        }

        // Optionally fall back to default certificate
        // By default this is disabled to prevent information disclosure
        if self.use_default_fallback {
            self.default_cert.read().clone()
        } else {
            None
        }
    }
}

impl TlsManager {
    /// Create a new TLS manager with default fallback disabled (secure by default)
    pub fn new() -> Self {
        let resolver = Arc::new(SniResolver::new());

        Self {
            resolver,
            stats: TlsStats::default(),
        }
    }

    /// Create a new TLS manager with configurable default certificate fallback
    ///
    /// # Arguments
    /// * `use_default_fallback` - If true, connections with non-matching SNI will receive
    ///   the default certificate. If false (recommended), they will be rejected.
    ///
    /// # Security
    /// Setting this to `true` may leak information about which certificates exist on the server.
    /// Only enable if you need backward compatibility with clients that don't support SNI.
    pub fn with_default_fallback(use_fallback: bool) -> Self {
        let resolver = Arc::new(SniResolver::with_default_fallback(use_fallback));

        Self {
            resolver,
            stats: TlsStats::default(),
        }
    }

    /// Load a certificate and key for a hostname
    pub fn load_cert<P: AsRef<Path>>(
        &self,
        hostname: &str,
        cert_path: P,
        key_path: P,
    ) -> anyhow::Result<()> {
        let certified_key = load_certified_key(cert_path.as_ref(), key_path.as_ref())?;
        let certified_key = Arc::new(certified_key);

        self.resolver.add_cert(hostname, certified_key.clone());

        // Set as default if this is the first cert
        if self.resolver.default_cert.read().is_none() {
            self.resolver.set_default(certified_key);
        }

        info!("Loaded TLS certificate for {}", hostname);
        Ok(())
    }

    /// Load a certificate and key from PEM strings for a hostname
    pub fn load_cert_pem(
        &self,
        hostname: &str,
        cert_pem: &str,
        key_pem: &str,
    ) -> anyhow::Result<()> {
        let certified_key = load_certified_key_from_pem(cert_pem, key_pem)?;
        let certified_key = Arc::new(certified_key);

        self.resolver.add_cert(hostname, certified_key.clone());

        if self.resolver.default_cert.read().is_none() {
            self.resolver.set_default(certified_key);
        }

        info!("Loaded TLS certificate for {}", hostname);
        Ok(())
    }

    /// Build the server configuration with ALPN support for HTTP/2
    pub fn build_server_config(&self) -> anyhow::Result<Arc<ServerConfig>> {
        let mut config = ServerConfig::builder()
            .with_no_client_auth()
            .with_cert_resolver(self.resolver.clone());

        // Enable ALPN for HTTP/2 and HTTP/1.1
        // Order matters: h2 is preferred over http/1.1
        config.alpn_protocols = vec![ALPN_H2.to_vec(), ALPN_HTTP11.to_vec()];

        Ok(Arc::new(config))
    }

    /// Build the server configuration with HTTP/1.1 only (no HTTP/2)
    pub fn build_server_config_http1_only(&self) -> anyhow::Result<Arc<ServerConfig>> {
        let mut config = ServerConfig::builder()
            .with_no_client_auth()
            .with_cert_resolver(self.resolver.clone());

        // Only advertise HTTP/1.1
        config.alpn_protocols = vec![ALPN_HTTP11.to_vec()];

        Ok(Arc::new(config))
    }

    /// Record a handshake start
    pub fn record_handshake_start(&self) {
        self.stats.handshakes_started.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a successful handshake
    pub fn record_handshake_complete(&self) {
        self.stats.handshakes_completed.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a failed handshake
    pub fn record_handshake_failed(&self) {
        self.stats.handshakes_failed.fetch_add(1, Ordering::Relaxed);
    }
}

/// Load a certificate chain and private key
fn load_certified_key(
    cert_path: &Path,
    key_path: &Path,
) -> anyhow::Result<rustls::sign::CertifiedKey> {
    // Load certificates
    let cert_file = File::open(cert_path)
        .map_err(|e| anyhow::anyhow!("Failed to open cert file {:?}: {}", cert_path, e))?;
    let mut cert_reader = BufReader::new(cert_file);

    let certs: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut cert_reader)
        .filter_map(|r| r.ok())
        .collect();

    if certs.is_empty() {
        anyhow::bail!("No certificates found in {:?}", cert_path);
    }

    // Load private key
    let key_file = File::open(key_path)
        .map_err(|e| anyhow::anyhow!("Failed to open key file {:?}: {}", key_path, e))?;
    let mut key_reader = BufReader::new(key_file);

    let key = load_private_key(&mut key_reader)?;

    // Create signing key
    let signing_key = rustls::crypto::ring::sign::any_supported_type(&key)
        .map_err(|e| anyhow::anyhow!("Invalid private key: {:?}", e))?;

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

/// Load a certificate chain and private key from PEM strings
fn load_certified_key_from_pem(
    cert_pem: &str,
    key_pem: &str,
) -> anyhow::Result<rustls::sign::CertifiedKey> {
    let mut cert_reader = BufReader::new(Cursor::new(cert_pem.as_bytes()));
    let certs: Vec<CertificateDer<'static>> = rustls_pemfile::certs(&mut cert_reader)
        .filter_map(|r| r.ok())
        .collect();

    if certs.is_empty() {
        anyhow::bail!("No certificates found in PEM data");
    }

    let mut key_reader = BufReader::new(Cursor::new(key_pem.as_bytes()));
    let key = load_private_key(&mut key_reader)?;
    let signing_key = rustls::crypto::ring::sign::any_supported_type(&key)
        .map_err(|e| anyhow::anyhow!("Invalid private key: {:?}", e))?;

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

/// Load a private key from a PEM file
fn load_private_key<R: std::io::BufRead>(reader: &mut R) -> anyhow::Result<PrivateKeyDer<'static>> {
    loop {
        match rustls_pemfile::read_one(reader)? {
            Some(rustls_pemfile::Item::Pkcs1Key(key)) => {
                return Ok(PrivateKeyDer::Pkcs1(key));
            }
            Some(rustls_pemfile::Item::Pkcs8Key(key)) => {
                return Ok(PrivateKeyDer::Pkcs8(key));
            }
            Some(rustls_pemfile::Item::Sec1Key(key)) => {
                return Ok(PrivateKeyDer::Sec1(key));
            }
            None => break,
            _ => continue,
        }
    }

    anyhow::bail!("No private key found")
}

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

    // =====================================================================
    // TlsManager tests
    // =====================================================================

    #[test]
    fn test_tls_manager_creation() {
        let manager = TlsManager::new();
        assert_eq!(manager.stats.handshakes_started.load(Ordering::Relaxed), 0);
        assert_eq!(manager.stats.handshakes_completed.load(Ordering::Relaxed), 0);
        assert_eq!(manager.stats.handshakes_failed.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_tls_manager_stats_recording() {
        let manager = TlsManager::new();

        manager.record_handshake_start();
        manager.record_handshake_start();
        manager.record_handshake_start();

        assert_eq!(manager.stats.handshakes_started.load(Ordering::Relaxed), 3);

        manager.record_handshake_complete();
        manager.record_handshake_complete();

        assert_eq!(manager.stats.handshakes_completed.load(Ordering::Relaxed), 2);

        manager.record_handshake_failed();

        assert_eq!(manager.stats.handshakes_failed.load(Ordering::Relaxed), 1);
    }

    #[test]
    fn test_tls_manager_stats_concurrent() {
        let manager = TlsManager::new();

        // Simulate concurrent updates
        for _ in 0..1000 {
            manager.record_handshake_start();
        }

        assert_eq!(manager.stats.handshakes_started.load(Ordering::Relaxed), 1000);
    }

    // =====================================================================
    // SniResolver tests
    // =====================================================================

    #[test]
    fn test_sni_resolver_creation() {
        let resolver = SniResolver::new();
        assert!(resolver.certs.is_empty());
        assert!(resolver.default_cert.read().is_none());
    }

    #[test]
    fn test_sni_resolver_debug() {
        let resolver = SniResolver::new();
        let debug_str = format!("{:?}", resolver);
        assert!(debug_str.contains("SniResolver"));
        assert!(debug_str.contains("cert_count"));
        assert!(debug_str.contains("has_default"));
    }

    // =====================================================================
    // TlsStats tests
    // =====================================================================

    #[test]
    fn test_tls_stats_default() {
        let stats = TlsStats::default();
        assert_eq!(stats.handshakes_started.load(Ordering::Relaxed), 0);
        assert_eq!(stats.handshakes_completed.load(Ordering::Relaxed), 0);
        assert_eq!(stats.handshakes_failed.load(Ordering::Relaxed), 0);
    }

    #[test]
    fn test_tls_stats_atomic_operations() {
        let stats = TlsStats::default();

        stats.handshakes_started.fetch_add(10, Ordering::Relaxed);
        stats.handshakes_completed.fetch_add(8, Ordering::Relaxed);
        stats.handshakes_failed.fetch_add(2, Ordering::Relaxed);

        assert_eq!(stats.handshakes_started.load(Ordering::Relaxed), 10);
        assert_eq!(stats.handshakes_completed.load(Ordering::Relaxed), 8);
        assert_eq!(stats.handshakes_failed.load(Ordering::Relaxed), 2);
    }

    // =====================================================================
    // Error handling tests
    // =====================================================================

    #[test]
    fn test_load_cert_nonexistent_file() {
        let manager = TlsManager::new();
        let result = manager.load_cert(
            "test.example.com",
            "/nonexistent/path/cert.pem",
            "/nonexistent/path/key.pem",
        );

        assert!(result.is_err());
    }

    #[test]
    fn test_load_private_key_empty_reader() {
        let mut reader = std::io::BufReader::new(std::io::Cursor::new(Vec::<u8>::new()));
        let result = load_private_key(&mut reader);

        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("No private key found"));
    }

    #[test]
    fn test_load_private_key_invalid_pem() {
        let invalid_data = b"not a valid PEM file";
        let mut reader = std::io::BufReader::new(std::io::Cursor::new(invalid_data.to_vec()));
        let result = load_private_key(&mut reader);

        assert!(result.is_err());
    }

    // =====================================================================
    // Default fallback behavior tests
    // =====================================================================

    #[test]
    fn test_resolver_secure_by_default() {
        let resolver = SniResolver::new();
        assert!(!resolver.use_default_fallback);
    }

    #[test]
    fn test_resolver_with_fallback_enabled() {
        let resolver = SniResolver::with_default_fallback(true);
        assert!(resolver.use_default_fallback);
    }

    #[test]
    fn test_resolver_with_fallback_disabled() {
        let resolver = SniResolver::with_default_fallback(false);
        assert!(!resolver.use_default_fallback);
    }

    #[test]
    fn test_manager_secure_by_default() {
        let manager = TlsManager::new();
        assert!(!manager.resolver.use_default_fallback);
    }

    #[test]
    fn test_manager_with_fallback_enabled() {
        let manager = TlsManager::with_default_fallback(true);
        assert!(manager.resolver.use_default_fallback);
    }

    #[test]
    fn test_manager_with_fallback_disabled() {
        let manager = TlsManager::with_default_fallback(false);
        assert!(!manager.resolver.use_default_fallback);
    }

    // =====================================================================
    // Integration-style tests
    // =====================================================================

    #[test]
    fn test_manager_workflow() {
        let manager = TlsManager::new();

        // Simulate a series of handshakes
        for i in 0..100 {
            manager.record_handshake_start();

            if i % 10 == 0 {
                manager.record_handshake_failed();
            } else {
                manager.record_handshake_complete();
            }
        }

        assert_eq!(manager.stats.handshakes_started.load(Ordering::Relaxed), 100);
        assert_eq!(manager.stats.handshakes_completed.load(Ordering::Relaxed), 90);
        assert_eq!(manager.stats.handshakes_failed.load(Ordering::Relaxed), 10);
    }

    #[test]
    fn test_manager_build_server_config() {
        // Install the ring crypto provider for this test
        let _ = rustls::crypto::ring::default_provider().install_default();

        let manager = TlsManager::new();

        // Should be able to build config even without certs (will fail on actual use)
        let config = manager.build_server_config();
        assert!(config.is_ok());
    }
}
