//! OCSP stapling support.
//!
//! Extracts the OCSP responder URL from a certificate, fetches the OCSP
//! response via HTTP, and caches it for hot-attachment to TLS connections.

use anyhow::{bail, Context, Result};
use dashmap::DashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::task::JoinHandle;
use x509_parser::extensions::ParsedExtension;
use x509_parser::pem::parse_x509_pem;

/// How long a cached OCSP response is considered valid.
const OCSP_CACHE_TTL: Duration = Duration::from_secs(12 * 3600); // 12 hours

/// How often the background refresh task runs.
const REFRESH_INTERVAL: Duration = Duration::from_secs(6 * 3600); // 6 hours

/// A cached OCSP response with a validity timestamp.
#[derive(Debug, Clone)]
struct CachedOcspResponse {
    /// The raw OCSP response bytes (DER-encoded).
    response: Vec<u8>,
    /// When this response was fetched.
    fetched_at: Instant,
}

impl CachedOcspResponse {
    fn is_valid(&self) -> bool {
        self.fetched_at.elapsed() < OCSP_CACHE_TTL
    }
}

/// Certificate + optional issuer PEM for OCSP request construction.
type CertWithIssuer = (Vec<u8>, Option<Vec<u8>>);

/// OCSP stapling manager.
///
/// Fetches and caches OCSP responses for served certificates.
pub struct OcspStapler {
    /// domain → cached OCSP response.
    cache: Arc<DashMap<String, CachedOcspResponse>>,
    /// domain → (cert_pem, issuer_pem) for refetching.
    certs: Arc<DashMap<String, CertWithIssuer>>,
}

impl OcspStapler {
    /// Create a new OCSP stapler.
    pub fn new() -> Self {
        Self {
            cache: Arc::new(DashMap::new()),
            certs: Arc::new(DashMap::new()),
        }
    }

    /// Register a certificate for OCSP stapling.
    ///
    /// `cert_pem` is the end-entity certificate; `issuer_pem` is the
    /// issuing CA certificate (needed to build the OCSP request).  If
    /// `issuer_pem` is `None`, the stapler will try to find the issuer
    /// in the cert chain.
    pub fn register_cert(&self, domain: &str, cert_pem: &[u8], issuer_pem: Option<&[u8]>) {
        self.certs.insert(
            domain.to_lowercase(),
            (cert_pem.to_vec(), issuer_pem.map(|b| b.to_vec())),
        );
    }

    /// Unregister a domain from OCSP stapling.
    pub fn unregister(&self, domain: &str) {
        let key = domain.to_lowercase();
        self.cache.remove(&key);
        self.certs.remove(&key);
    }

    /// Fetch an OCSP response for the given PEM-encoded certificate.
    ///
    /// This makes an HTTP POST to the OCSP responder URL extracted from
    /// the certificate's Authority Information Access extension.
    ///
    /// `issuer_pem` is required for RFC 6960-compliant OCSP requests.
    /// Without the issuer certificate, the issuer key hash cannot be
    /// computed and this method returns an error.
    pub async fn fetch_ocsp_response(
        cert_pem: &[u8],
        issuer_pem: Option<&[u8]>,
    ) -> Result<Vec<u8>> {
        let responder_url = extract_ocsp_responder_url(cert_pem)?;
        tracing::debug!(url = %responder_url, "fetching OCSP response");

        // Build the OCSP request body.  RFC 6960 requires the issuer
        // certificate to compute the correct issuer key hash.
        let ocsp_request = build_ocsp_request(cert_pem, issuer_pem)?;

        let client = reqwest::Client::builder()
            .timeout(Duration::from_secs(10))
            .build()
            .context("build HTTP client for OCSP")?;

        let response = client
            .post(&responder_url)
            .header("Content-Type", "application/ocsp-request")
            .body(ocsp_request)
            .send()
            .await
            .with_context(|| format!("POST to OCSP responder: {responder_url}"))?;

        if !response.status().is_success() {
            bail!(
                "OCSP responder returned HTTP {}",
                response.status().as_u16()
            );
        }

        let body = response.bytes().await.context("read OCSP response body")?;

        if body.is_empty() {
            bail!("OCSP responder returned empty response");
        }

        Ok(body.to_vec())
    }

    /// Fetch and cache the OCSP response for a domain.
    ///
    /// Returns `Ok(())` on success. If no issuer certificate is registered
    /// for the domain, OCSP stapling is skipped gracefully (logged as info).
    pub async fn staple(&self, domain: &str) -> Result<()> {
        let key = domain.to_lowercase();
        let cert_entry = self
            .certs
            .get(&key)
            .with_context(|| format!("no cert registered for OCSP: {domain}"))?;

        let (ref cert_pem, ref issuer_pem) = *cert_entry;

        // Without the issuer certificate we cannot build a conformant
        // OCSP request (the issuer key hash would be wrong).  Skip
        // gracefully rather than sending a malformed request.
        if issuer_pem.is_none() {
            tracing::info!(
                domain,
                "OCSP stapling skipped: issuer certificate not available"
            );
            return Ok(());
        }

        let response = Self::fetch_ocsp_response(cert_pem, issuer_pem.as_deref()).await?;

        self.cache.insert(
            key,
            CachedOcspResponse {
                response,
                fetched_at: Instant::now(),
            },
        );

        tracing::info!(domain, "OCSP response stapled");
        Ok(())
    }

    /// Get the cached OCSP response for a domain, if available and valid.
    pub fn get_stapled_response(&self, domain: &str) -> Option<Vec<u8>> {
        let key = domain.to_lowercase();
        self.cache.get(&key).and_then(|entry| {
            if entry.is_valid() {
                Some(entry.response.clone())
            } else {
                None
            }
        })
    }

    /// Return the number of cached OCSP responses.
    pub fn cached_count(&self) -> usize {
        self.cache.len()
    }

    /// Start a background task that periodically refreshes OCSP responses.
    pub fn start_refresh(self: Arc<Self>) -> JoinHandle<()> {
        let stapler = Arc::clone(&self);
        tokio::spawn(async move {
            tracing::info!("OCSP stapling refresh task started");
            let mut interval = tokio::time::interval(REFRESH_INTERVAL);

            loop {
                interval.tick().await;
                stapler.refresh_all().await;
            }
        })
    }

    /// Refresh all registered OCSP responses.
    async fn refresh_all(&self) {
        let domains: Vec<String> = self.certs.iter().map(|entry| entry.key().clone()).collect();

        for domain in &domains {
            if let Err(e) = self.staple(domain).await {
                tracing::warn!(domain, error = %e, "OCSP refresh failed");
            }
        }

        if !domains.is_empty() {
            tracing::debug!(count = domains.len(), "OCSP refresh cycle complete");
        }
    }
}

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

/// Extract the OCSP responder URL from the Authority Information Access
/// extension of a PEM-encoded certificate.
pub fn extract_ocsp_responder_url(cert_pem: &[u8]) -> Result<String> {
    let (_, pem) = parse_x509_pem(cert_pem).map_err(|e| anyhow::anyhow!("parse PEM: {e}"))?;

    let (_, cert) = x509_parser::parse_x509_certificate(&pem.contents)
        .map_err(|e| anyhow::anyhow!("parse X.509: {e}"))?;

    for ext in cert.extensions() {
        if let ParsedExtension::AuthorityInfoAccess(aia) = ext.parsed_extension() {
            for desc in &aia.accessdescs {
                // OID 1.3.6.1.5.5.7.48.1 = OCSP
                if desc.access_method.to_string() == "1.3.6.1.5.5.7.48.1" {
                    if let x509_parser::extensions::GeneralName::URI(url) = desc.access_location {
                        return Ok(url.to_string());
                    }
                }
            }
        }
    }

    bail!("no OCSP responder URL found in certificate");
}

/// Build an RFC 6960-compliant OCSP request for the given certificate.
///
/// This constructs a DER-encoded OCSPRequest containing a single
/// CertID derived from the certificate's serial number, issuer DN
/// hash, and issuer public key hash.
///
/// When `issuer_pem` is `Some`, the issuer's SubjectPublicKeyInfo
/// is extracted and hashed to produce a correct `issuerKeyHash`.
/// When `issuer_pem` is `None`, this function returns an error
/// because the resulting request would violate RFC 6960.
fn build_ocsp_request(cert_pem: &[u8], issuer_pem: Option<&[u8]>) -> Result<Vec<u8>> {
    let (_, pem) =
        parse_x509_pem(cert_pem).map_err(|e| anyhow::anyhow!("parse PEM for OCSP req: {e}"))?;

    let (_, cert) = x509_parser::parse_x509_certificate(&pem.contents)
        .map_err(|e| anyhow::anyhow!("parse X.509 for OCSP req: {e}"))?;

    let serial = cert.raw_serial();
    let issuer_dn_raw = cert.issuer().as_raw();

    // SHA-1 hash of the issuer distinguished name (DER-encoded).
    let issuer_name_hash =
        ring::digest::digest(&ring::digest::SHA1_FOR_LEGACY_USE_ONLY, issuer_dn_raw);

    // Compute the issuer key hash from the issuer certificate's
    // SubjectPublicKeyInfo.  RFC 6960 §4.1.1 specifies this as the
    // SHA-1 hash of the value of the BIT STRING subjectPublicKey
    // (excluding tag and length) in the issuer's SubjectPublicKeyInfo.
    let issuer_key_hash = if let Some(issuer_pem_bytes) = issuer_pem {
        let (_, issuer_parsed) = parse_x509_pem(issuer_pem_bytes)
            .map_err(|e| anyhow::anyhow!("failed to parse issuer PEM: {e}"))?;
        let (_, issuer_cert) = x509_parser::parse_x509_certificate(&issuer_parsed.contents)
            .map_err(|e| anyhow::anyhow!("failed to parse issuer X.509: {e}"))?;
        let spki = &issuer_cert
            .tbs_certificate
            .subject_pki
            .subject_public_key
            .data;
        ring::digest::digest(&ring::digest::SHA1_FOR_LEGACY_USE_ONLY, spki)
    } else {
        // Without the issuer certificate we cannot compute the correct
        // key hash.  Bail rather than sending a malformed OCSP request.
        bail!("issuer certificate required for RFC 6960-compliant OCSP request");
    };

    // Build the DER-encoded OCSPRequest manually.
    // Structure (simplified ASN.1):
    //   OCSPRequest ::= SEQUENCE {
    //     tbsRequest TBSRequest }
    //   TBSRequest  ::= SEQUENCE {
    //     requestList SEQUENCE OF Request }
    //   Request     ::= SEQUENCE {
    //     reqCert CertID }
    //   CertID      ::= SEQUENCE {
    //     hashAlgorithm  AlgorithmIdentifier (SHA-1),
    //     issuerNameHash OCTET STRING,
    //     issuerKeyHash  OCTET STRING,
    //     serialNumber   INTEGER }
    let mut cert_id = Vec::new();
    // AlgorithmIdentifier for SHA-1: SEQUENCE { OID 1.3.14.3.2.26, NULL }
    let sha1_alg = &[
        0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00,
    ];
    cert_id.extend_from_slice(sha1_alg);
    // issuerNameHash OCTET STRING
    cert_id.push(0x04);
    cert_id.push(issuer_name_hash.as_ref().len() as u8);
    cert_id.extend_from_slice(issuer_name_hash.as_ref());
    // issuerKeyHash OCTET STRING
    cert_id.push(0x04);
    cert_id.push(issuer_key_hash.as_ref().len() as u8);
    cert_id.extend_from_slice(issuer_key_hash.as_ref());
    // serialNumber INTEGER
    cert_id.push(0x02);
    cert_id.push(serial.len() as u8);
    cert_id.extend_from_slice(serial);

    // Wrap CertID in SEQUENCE
    let cert_id_seq = der_sequence(&cert_id);
    // Request = SEQUENCE { CertID }
    let request = der_sequence(&cert_id_seq);
    // requestList = SEQUENCE OF Request
    let request_list = der_sequence(&request);
    // TBSRequest = SEQUENCE { requestList }
    let tbs_request = der_sequence(&request_list);
    // OCSPRequest = SEQUENCE { tbsRequest }
    let ocsp_request = der_sequence(&tbs_request);

    Ok(ocsp_request)
}

/// Wrap `data` in a DER SEQUENCE (tag 0x30).
fn der_sequence(data: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(2 + data.len());
    out.push(0x30);
    der_encode_length(data.len(), &mut out);
    out.extend_from_slice(data);
    out
}

/// Encode a DER length.
fn der_encode_length(len: usize, out: &mut Vec<u8>) {
    if len < 0x80 {
        out.push(len as u8);
    } else if len < 0x100 {
        out.push(0x81);
        out.push(len as u8);
    } else {
        out.push(0x82);
        out.push((len >> 8) as u8);
        out.push(len as u8);
    }
}

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

    #[test]
    fn extract_ocsp_url_from_self_signed_fails() {
        // Self-signed certs don't have AIA extension.
        let out = generate_self_signed(&["test.com".into()]).unwrap();
        let result = extract_ocsp_responder_url(&out.cert_pem);
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(msg.contains("no OCSP responder URL"), "got: {msg}");
    }

    #[test]
    fn extract_ocsp_url_invalid_pem() {
        let result = extract_ocsp_responder_url(b"not a cert");
        assert!(result.is_err());
    }

    #[test]
    fn build_ocsp_request_with_issuer() {
        // Self-signed cert is its own issuer, so we can use it as both.
        let out = generate_self_signed(&["test.com".into()]).unwrap();
        let request = build_ocsp_request(&out.cert_pem, Some(&out.cert_pem)).unwrap();

        // Should be a non-empty DER structure starting with SEQUENCE tag.
        assert!(!request.is_empty());
        assert_eq!(request[0], 0x30, "should start with SEQUENCE tag");
    }

    #[test]
    fn build_ocsp_request_without_issuer_fails() {
        let out = generate_self_signed(&["test.com".into()]).unwrap();
        let result = build_ocsp_request(&out.cert_pem, None);
        assert!(result.is_err());
        let msg = result.unwrap_err().to_string();
        assert!(msg.contains("issuer certificate required"), "got: {msg}");
    }

    #[test]
    fn build_ocsp_request_invalid_pem() {
        let result = build_ocsp_request(b"garbage", None);
        assert!(result.is_err());
    }

    #[test]
    fn issuer_key_hash_differs_from_name_hash() {
        // For a self-signed cert, the issuer key hash (hash of the public
        // key bits) must differ from the issuer name hash (hash of the
        // subject DN). They are computed from entirely different data.
        let out = generate_self_signed(&["test.com".into()]).unwrap();
        let (_, pem) = parse_x509_pem(&out.cert_pem).unwrap();
        let (_, cert) = x509_parser::parse_x509_certificate(&pem.contents).unwrap();

        let issuer_dn_raw = cert.issuer().as_raw();
        let issuer_name_hash =
            ring::digest::digest(&ring::digest::SHA1_FOR_LEGACY_USE_ONLY, issuer_dn_raw);

        // The issuer is itself (self-signed), so extract its own SPKI.
        let spki = &cert.tbs_certificate.subject_pki.subject_public_key.data;
        let issuer_key_hash = ring::digest::digest(&ring::digest::SHA1_FOR_LEGACY_USE_ONLY, spki);

        assert_ne!(
            issuer_name_hash.as_ref(),
            issuer_key_hash.as_ref(),
            "issuer_key_hash should differ from issuer_name_hash for a real cert"
        );
    }

    #[test]
    fn der_sequence_encoding() {
        let data = vec![0x01, 0x02, 0x03];
        let seq = der_sequence(&data);
        assert_eq!(seq, vec![0x30, 0x03, 0x01, 0x02, 0x03]);
    }

    #[test]
    fn der_length_short() {
        let mut out = Vec::new();
        der_encode_length(5, &mut out);
        assert_eq!(out, vec![0x05]);
    }

    #[test]
    fn der_length_medium() {
        let mut out = Vec::new();
        der_encode_length(200, &mut out);
        assert_eq!(out, vec![0x81, 200]);
    }

    #[test]
    fn der_length_long() {
        let mut out = Vec::new();
        der_encode_length(300, &mut out);
        assert_eq!(out, vec![0x82, 0x01, 0x2C]);
    }

    #[test]
    fn ocsp_stapler_new() {
        let stapler = OcspStapler::new();
        assert_eq!(stapler.cached_count(), 0);
    }

    #[test]
    fn ocsp_stapler_default() {
        let stapler = OcspStapler::default();
        assert_eq!(stapler.cached_count(), 0);
    }

    #[test]
    fn register_and_unregister() {
        let stapler = OcspStapler::new();
        let out = generate_self_signed(&["test.com".into()]).unwrap();

        stapler.register_cert("test.com", &out.cert_pem, None);
        assert!(stapler.certs.contains_key("test.com"));

        stapler.unregister("test.com");
        assert!(!stapler.certs.contains_key("test.com"));
    }

    #[test]
    fn get_stapled_response_missing() {
        let stapler = OcspStapler::new();
        assert!(stapler.get_stapled_response("missing.com").is_none());
    }

    #[test]
    fn cached_response_is_returned() {
        let stapler = OcspStapler::new();

        // Manually insert a cached response.
        stapler.cache.insert(
            "cached.com".to_string(),
            CachedOcspResponse {
                response: vec![1, 2, 3, 4, 5],
                fetched_at: Instant::now(),
            },
        );

        let resp = stapler.get_stapled_response("cached.com");
        assert!(resp.is_some());
        assert_eq!(resp.unwrap(), vec![1, 2, 3, 4, 5]);
    }

    #[test]
    fn expired_cache_returns_none() {
        let _stapler = OcspStapler::new();

        // Insert a response that's "old" by subtracting from Instant.
        // Since we can't easily create an Instant in the past, we'll
        // test the is_valid logic directly.
        let old = CachedOcspResponse {
            response: vec![1, 2, 3],
            fetched_at: Instant::now() - OCSP_CACHE_TTL - Duration::from_secs(1),
        };
        assert!(!old.is_valid());

        let fresh = CachedOcspResponse {
            response: vec![4, 5, 6],
            fetched_at: Instant::now(),
        };
        assert!(fresh.is_valid());
    }

    #[test]
    fn cache_ttl_is_12_hours() {
        assert_eq!(OCSP_CACHE_TTL, Duration::from_secs(43200));
    }

    #[test]
    fn refresh_interval_is_6_hours() {
        assert_eq!(REFRESH_INTERVAL, Duration::from_secs(21600));
    }

    #[test]
    fn register_cert_case_insensitive() {
        let stapler = OcspStapler::new();
        let out = generate_self_signed(&["TEST.COM".into()]).unwrap();

        stapler.register_cert("TEST.COM", &out.cert_pem, None);
        assert!(stapler.certs.contains_key("test.com"));
    }

    #[tokio::test]
    async fn staple_unregistered_domain_fails() {
        let stapler = OcspStapler::new();
        let result = stapler.staple("unknown.com").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn refresh_all_with_no_certs() {
        let stapler = OcspStapler::new();
        // Should complete without errors.
        stapler.refresh_all().await;
    }
}
