//! Certificate Authority for MITM proxy.
//!
//! Generates a self-signed root CA at first run and caches it to disk.
//! Dynamically generates per-domain leaf certificates signed by this CA.

use rcgen::{
    BasicConstraints, CertificateParams, DistinguishedName, DnType, ExtendedKeyUsagePurpose, IsCa,
    KeyPair, KeyUsagePurpose, SanType,
};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::info;

/// MITM Certificate Authority.
pub struct MitmCa {
    /// Root CA certificate (DER-encoded for rustls).
    ca_cert_der: CertificateDer<'static>,
    /// Root CA private key.
    ca_key: KeyPair,
    /// Signed root CA cert (needed by rcgen to sign leaf certs).
    ca_signed: rcgen::Certificate,
    /// Cache of per-domain TLS configs.
    domain_cache: Arc<RwLock<HashMap<String, Arc<rustls::ServerConfig>>>>,
    /// Path to the CA PEM file (for SSL_CERT_FILE combined bundle).
    pub ca_pem_path: PathBuf,
}

impl MitmCa {
    /// Load or generate the MITM CA.
    ///
    /// The CA cert/key are stored at:
    ///   `<data_dir>/mitm-ca.pem`  (cert, for NODE_EXTRA_CA_CERTS)
    ///   `<data_dir>/mitm-ca.key`  (private key)
    pub fn load_or_generate(data_dir: &Path) -> Result<Self, String> {
        let cert_path = data_dir.join("mitm-ca.pem");
        let key_path = data_dir.join("mitm-ca.key");

        if cert_path.exists() && key_path.exists() {
            info!("Loading existing MITM CA from {}", cert_path.display());
            let cert_pem = std::fs::read_to_string(&cert_path)
                .map_err(|e| format!("Failed to read CA cert: {e}"))?;
            let key_pem = std::fs::read_to_string(&key_path)
                .map_err(|e| format!("Failed to read CA key: {e}"))?;

            let ca_key =
                KeyPair::from_pem(&key_pem).map_err(|e| format!("Failed to parse CA key: {e}"))?;

            // Re-create params and self-sign to get the rcgen Certificate object
            // (needed for signing leaf certs — rcgen 0.13 doesn't have from_ca_cert_pem).
            // The re-signed cert will have a different serial/notBefore, but that's fine
            // because we only use it for the rcgen signing API, NOT for the on-disk PEM.
            let params = Self::ca_params();
            let ca_signed = params
                .self_signed(&ca_key)
                .map_err(|e| format!("Failed to self-sign CA: {e}"))?;

            // Use the ORIGINAL on-disk PEM cert for DER — this is what the LS trusts
            // (via the combined CA bundle built by the wrapper script). Writing the
            // re-signed cert back would desync the LS's trust anchor.
            let ca_cert_der = Self::pem_to_der(&cert_pem)
                .unwrap_or_else(|| CertificateDer::from(ca_signed.der().to_vec()));

            Ok(Self {
                ca_cert_der,
                ca_key,
                ca_signed,
                domain_cache: Arc::new(RwLock::new(HashMap::new())),
                ca_pem_path: cert_path,
            })
        } else {
            info!("Generating new MITM CA at {}", cert_path.display());

            // Ensure data dir exists
            std::fs::create_dir_all(data_dir)
                .map_err(|e| format!("Failed to create data dir: {e}"))?;

            let ca_key =
                KeyPair::generate().map_err(|e| format!("Failed to generate CA key: {e}"))?;

            let params = Self::ca_params();
            let ca_signed = params
                .self_signed(&ca_key)
                .map_err(|e| format!("Failed to self-sign CA: {e}"))?;

            // Write cert and key to disk
            std::fs::write(&cert_path, ca_signed.pem())
                .map_err(|e| format!("Failed to write CA cert: {e}"))?;
            std::fs::write(&key_path, ca_key.serialize_pem())
                .map_err(|e| format!("Failed to write CA key: {e}"))?;

            #[cfg(unix)]
            {
                use std::os::unix::fs::PermissionsExt;
                let _ = std::fs::set_permissions(&key_path, std::fs::Permissions::from_mode(0o600));
            }

            let ca_cert_der = CertificateDer::from(ca_signed.der().to_vec());

            Ok(Self {
                ca_cert_der,
                ca_key,
                ca_signed,
                domain_cache: Arc::new(RwLock::new(HashMap::new())),
                ca_pem_path: cert_path,
            })
        }
    }

    /// Build the CA certificate parameters (reusable for both generate and load).
    fn ca_params() -> CertificateParams {
        let mut params = CertificateParams::default();

        let mut dn = DistinguishedName::new();
        dn.push(DnType::CommonName, "Antigravity MITM CA");
        dn.push(DnType::OrganizationName, "ZeroGravity");
        params.distinguished_name = dn;

        params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
        params.key_usages = vec![KeyUsagePurpose::KeyCertSign, KeyUsagePurpose::CrlSign];

        // Valid for 10 years
        let now = time::OffsetDateTime::now_utc();
        params.not_before = now;
        params.not_after = now + time::Duration::days(3650);

        params
    }

    /// Parse a PEM certificate into a DER-encoded CertificateDer.
    fn pem_to_der(pem: &str) -> Option<CertificateDer<'static>> {
        // Extract base64 content between BEGIN/END markers
        let mut in_cert = false;
        let mut b64 = String::new();
        for line in pem.lines() {
            if line.contains("BEGIN CERTIFICATE") {
                in_cert = true;
                continue;
            }
            if line.contains("END CERTIFICATE") {
                break;
            }
            if in_cert {
                b64.push_str(line.trim());
            }
        }
        if b64.is_empty() {
            return None;
        }
        use base64::Engine;
        let der = base64::engine::general_purpose::STANDARD
            .decode(&b64)
            .ok()?;
        Some(CertificateDer::from(der))
    }

    /// Get or create a TLS ServerConfig for the given domain.
    pub async fn server_config_for_domain(
        &self,
        domain: &str,
    ) -> Result<Arc<rustls::ServerConfig>, String> {
        // Check cache first
        {
            let cache = self.domain_cache.read().await;
            if let Some(config) = cache.get(domain) {
                return Ok(config.clone());
            }
        }

        // Generate leaf cert for this domain
        let mut params = CertificateParams::default();

        let mut dn = DistinguishedName::new();
        dn.push(DnType::CommonName, domain);
        params.distinguished_name = dn;

        params.subject_alt_names = vec![SanType::DnsName(
            domain
                .try_into()
                .map_err(|e| format!("Invalid domain: {e}"))?,
        )];
        params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth];
        params.key_usages = vec![
            KeyUsagePurpose::DigitalSignature,
            KeyUsagePurpose::KeyEncipherment,
        ];

        // Valid for 1 year
        let now = time::OffsetDateTime::now_utc();
        params.not_before = now;
        params.not_after = now + time::Duration::days(365);

        let leaf_key =
            KeyPair::generate().map_err(|e| format!("Failed to generate leaf key: {e}"))?;

        let leaf_cert = params
            .signed_by(&leaf_key, &self.ca_signed, &self.ca_key)
            .map_err(|e| format!("Failed to sign leaf cert for {domain}: {e}"))?;

        // Build rustls ServerConfig
        let leaf_cert_der = CertificateDer::from(leaf_cert.der().to_vec());
        let leaf_key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(leaf_key.serialize_der()));

        let mut config = rustls::ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(vec![leaf_cert_der, self.ca_cert_der.clone()], leaf_key_der)
            .map_err(|e| format!("Failed to build ServerConfig for {domain}: {e}"))?;

        // Advertise both h2 and http/1.1 so gRPC clients can negotiate HTTP/2
        config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];

        let config = Arc::new(config);

        // Cache it (cap at 64 entries — in practice only ~5-10 domains are seen)
        {
            let mut cache = self.domain_cache.write().await;
            if cache.len() >= 64 {
                cache.clear();
            }
            cache.insert(domain.to_string(), config.clone());
        }

        Ok(config)
    }
}
