//! ACME client for automatic certificate provisioning.
//!
//! Supports Let's Encrypt (production & staging) and ZeroSSL via the
//! `instant-acme` crate.  Handles account creation, order placement,
//! HTTP-01 and TLS-ALPN-01 challenges, order finalization, and
//! certificate download.

use crate::config::{AcmeCa, ChallengeType, TlsConfig};
use crate::store::CertStore;
use anyhow::{bail, Context, Result};
use instant_acme::{
    Account, AccountCredentials, AuthorizationStatus, ChallengeType as AcmeChallengeType,
    Identifier, NewAccount, NewOrder, OrderStatus,
};
use rcgen::{
    CertificateParams, DistinguishedName, DnType, Ia5String, KeyPair, SanType,
    PKCS_ECDSA_P256_SHA256,
};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::time::sleep;

/// Represents a pending ACME challenge that the caller must serve before
/// the ACME server will validate domain ownership.
#[derive(Debug, Clone)]
pub struct AcmeChallenge {
    /// The domain this challenge is for.
    pub domain: String,
    /// The challenge type.
    pub challenge_type: ChallengeType,
    /// For HTTP-01: the token (the path component after `/.well-known/acme-challenge/`).
    pub token: String,
    /// For HTTP-01: the key authorization string to serve as the response body.
    /// For TLS-ALPN-01: the key authorization hash (SHA-256) used in the self-signed cert.
    pub key_authorization: String,
    /// The challenge URL to POST to after setting up the response.
    pub url: String,
}

/// Minimum cooldown between certificate issuance attempts for the same domain.
const ACME_PER_DOMAIN_COOLDOWN: Duration = Duration::from_secs(3600); // 1 hour

/// ACME client wrapping `instant-acme` for certificate provisioning.
pub struct AcmeClient {
    /// The ACME CA directory URL.
    directory_url: String,
    /// Contact email for registration.
    email: Option<String>,
    /// Preferred challenge type.
    challenge_type: ChallengeType,
    /// Optional certificate store for persisting/loading ACME account credentials.
    store: Option<Arc<Mutex<CertStore>>>,
    /// Per-domain cooldown tracking to prevent abuse (last issuance time).
    domain_cooldowns: dashmap::DashMap<String, std::time::Instant>,
}

impl AcmeClient {
    /// Create a new ACME client from the TLS configuration.
    pub fn new(config: &TlsConfig) -> Result<Self> {
        Ok(Self {
            directory_url: config.acme_ca.directory_url().to_string(),
            email: config.acme_email.clone(),
            challenge_type: config.challenge.clone(),
            store: None,
            domain_cooldowns: dashmap::DashMap::new(),
        })
    }

    /// Create an ACME client targeting a specific directory URL (for testing).
    pub fn with_directory_url(url: &str, email: Option<String>) -> Self {
        Self {
            directory_url: url.to_string(),
            email,
            challenge_type: ChallengeType::Http01,
            store: None,
            domain_cooldowns: dashmap::DashMap::new(),
        }
    }

    /// Attach a certificate store for persisting/loading ACME account credentials.
    pub fn with_store(mut self, store: Arc<Mutex<CertStore>>) -> Self {
        self.store = Some(store);
        self
    }

    /// Return the configured directory URL.
    pub fn directory_url(&self) -> &str {
        &self.directory_url
    }

    /// Provision a certificate for the given domains.
    ///
    /// This performs the full ACME flow:
    /// 1. Create/register an account
    /// 2. Create an order for the domains
    /// 3. Process authorizations and challenges
    /// 4. Finalize with a CSR
    /// 5. Download the certificate chain
    ///
    /// The `challenge_handler` callback is invoked for each challenge that must
    /// be served. The caller must arrange for the challenge response to be
    /// available (e.g., via an HTTP endpoint) before the callback returns.
    ///
    /// Returns (cert_chain_pem, private_key_pem).
    pub async fn provision_cert<F>(
        &self,
        domains: &[String],
        challenge_handler: F,
    ) -> Result<(Vec<u8>, Vec<u8>)>
    where
        F: Fn(AcmeChallenge) -> Result<()>,
    {
        if domains.is_empty() {
            bail!("at least one domain is required for ACME provisioning");
        }

        // Per-domain rate limiting: reject if any domain was issued within the cooldown period.
        let now_instant = std::time::Instant::now();
        for domain in domains {
            if let Some(last) = self.domain_cooldowns.get(domain) {
                let elapsed = now_instant.duration_since(*last);
                if elapsed < ACME_PER_DOMAIN_COOLDOWN {
                    let remaining = ACME_PER_DOMAIN_COOLDOWN - elapsed;
                    bail!(
                        "ACME rate limit: domain '{}' was provisioned {:.0}s ago, \
                         cooldown is {}s (retry in {:.0}s)",
                        domain,
                        elapsed.as_secs_f64(),
                        ACME_PER_DOMAIN_COOLDOWN.as_secs(),
                        remaining.as_secs_f64()
                    );
                }
            }
        }

        tracing::info!(domains = ?domains, ca = %self.directory_url, "starting ACME provisioning");

        // 1. Create account.
        let account = self.create_account().await?;

        // 2. Create order.
        let identifiers: Vec<Identifier> =
            domains.iter().map(|d| Identifier::Dns(d.clone())).collect();

        let mut order = account
            .new_order(&NewOrder {
                identifiers: &identifiers,
            })
            .await
            .context("create ACME order")?;

        // 3. Process authorizations.
        let authorizations = order
            .authorizations()
            .await
            .context("fetch authorizations")?;

        for auth in &authorizations {
            match auth.status {
                AuthorizationStatus::Valid => {
                    tracing::debug!("authorization already valid, skipping");
                    continue;
                }
                AuthorizationStatus::Pending => {}
                _ => {
                    bail!("unexpected authorization status: {:?}", auth.status);
                }
            }

            let domain = match &auth.identifier {
                Identifier::Dns(d) => d.clone(),
            };

            // Find the matching challenge.
            let desired_type = match &self.challenge_type {
                ChallengeType::Http01 => AcmeChallengeType::Http01,
                ChallengeType::TlsAlpn01 => AcmeChallengeType::TlsAlpn01,
            };

            let challenge = auth
                .challenges
                .iter()
                .find(|c| c.r#type == desired_type)
                .with_context(|| {
                    format!(
                        "no {:?} challenge available for {domain}",
                        self.challenge_type
                    )
                })?;

            let key_auth = order.key_authorization(challenge);

            let acme_challenge = AcmeChallenge {
                domain: domain.clone(),
                challenge_type: self.challenge_type.clone(),
                token: challenge.token.clone(),
                key_authorization: key_auth.as_str().to_string(),
                url: challenge.url.clone(),
            };

            // Let the caller set up the challenge response.
            challenge_handler(acme_challenge)
                .with_context(|| format!("set up challenge for {domain}"))?;

            // Tell the ACME server we're ready.
            order
                .set_challenge_ready(&challenge.url)
                .await
                .with_context(|| format!("set challenge ready for {domain}"))?;
        }

        // 4. Wait for order to become ready.
        let status = self.poll_order_ready(&mut order).await?;
        if status != OrderStatus::Ready {
            bail!("order not ready after polling, status: {status:?}");
        }

        // 5. Generate a CSR key pair and finalize.
        let csr_key_pair =
            KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256).context("generate CSR key pair")?;

        let mut csr_params = CertificateParams::default();
        let mut dn = DistinguishedName::new();
        dn.push(DnType::CommonName, &domains[0]);
        csr_params.distinguished_name = dn;

        let mut sans = Vec::new();
        for d in domains {
            let ia5: Ia5String = d
                .as_str()
                .try_into()
                .with_context(|| format!("invalid domain for SAN: {d}"))?;
            sans.push(SanType::DnsName(ia5));
        }
        csr_params.subject_alt_names = sans;

        let csr = csr_params
            .serialize_request(&csr_key_pair)
            .context("serialize CSR")?;

        order
            .finalize(csr.der())
            .await
            .context("finalize ACME order")?;

        // 6. Wait for order to become valid (cert issued).
        let cert_chain_pem = self
            .poll_order_certificate(&mut order)
            .await
            .context("download certificate")?;

        let key_pem = csr_key_pair.serialize_pem();

        tracing::info!(domains = ?domains, "ACME certificate provisioned successfully");

        // Record issuance time for per-domain rate limiting.
        let now_instant = std::time::Instant::now();
        for domain in domains {
            self.domain_cooldowns.insert(domain.clone(), now_instant);
        }

        Ok((cert_chain_pem.into_bytes(), key_pem.into_bytes()))
    }

    /// Register a new ACME account, or restore one from persisted credentials.
    async fn create_account(&self) -> Result<Account> {
        // Try to load existing credentials from the store (lock is dropped before await).
        let saved_json: Option<String> = self.store.as_ref().and_then(|store_arc| {
            let store = store_arc.lock().ok()?;
            store.get_acme_credentials(&self.directory_url)
        });

        if let Some(creds_json) = saved_json {
            match serde_json::from_str::<AccountCredentials>(&creds_json) {
                Ok(credentials) => match Account::from_credentials(credentials).await {
                    Ok(account) => {
                        tracing::info!("restored ACME account from stored credentials");
                        return Ok(account);
                    }
                    Err(e) => {
                        tracing::warn!(
                            error = %e,
                            "failed to restore ACME account from stored credentials, creating new"
                        );
                    }
                },
                Err(e) => {
                    tracing::warn!(
                        error = %e,
                        "failed to deserialize stored ACME credentials, creating new"
                    );
                }
            }
        }

        let contact = self.email.as_ref().map(|e| format!("mailto:{e}"));

        let contact_refs: Vec<&str> = contact.iter().map(String::as_str).collect();

        let (account, credentials) = Account::create(
            &NewAccount {
                contact: &contact_refs,
                terms_of_service_agreed: true,
                only_return_existing: false,
            },
            &self.directory_url,
            None,
        )
        .await
        .context("create ACME account")?;

        // Persist credentials to the store for reuse across restarts.
        if let Some(ref store_arc) = self.store {
            if let Ok(store) = store_arc.lock() {
                match serde_json::to_string(&credentials) {
                    Ok(creds_json) => {
                        if let Err(e) =
                            store.save_acme_credentials(&self.directory_url, &creds_json)
                        {
                            tracing::warn!(error = %e, "failed to persist ACME account credentials");
                        } else {
                            tracing::debug!("persisted ACME account credentials to store");
                        }
                    }
                    Err(e) => {
                        tracing::warn!(error = %e, "failed to serialize ACME account credentials");
                    }
                }
            }
        }

        tracing::info!("ACME account registered");
        Ok(account)
    }

    /// Poll the order until it reaches `Ready` status (or timeout).
    async fn poll_order_ready(&self, order: &mut instant_acme::Order) -> Result<OrderStatus> {
        let max_attempts = 20;
        let delay = Duration::from_secs(3);

        for attempt in 1..=max_attempts {
            let state = order.refresh().await.context("refresh order state")?;
            tracing::debug!(attempt, status = ?state.status, "polling order");

            match state.status {
                OrderStatus::Ready | OrderStatus::Valid => return Ok(state.status),
                OrderStatus::Invalid => bail!("ACME order became invalid"),
                OrderStatus::Pending | OrderStatus::Processing => {
                    sleep(delay).await;
                }
            }
        }

        bail!("ACME order did not become ready after {max_attempts} attempts");
    }

    /// Poll until the certificate is available, then download it.
    async fn poll_order_certificate(&self, order: &mut instant_acme::Order) -> Result<String> {
        let max_attempts = 20;
        let delay = Duration::from_secs(3);

        for attempt in 1..=max_attempts {
            match order.certificate().await {
                Ok(Some(cert_chain_pem)) => {
                    tracing::debug!(attempt, "certificate downloaded");
                    return Ok(cert_chain_pem);
                }
                Ok(None) => {
                    tracing::debug!(attempt, "certificate not yet available, polling...");
                    sleep(delay).await;
                }
                Err(e) => {
                    bail!("error downloading certificate: {e}");
                }
            }
        }

        bail!("certificate not available after {max_attempts} attempts");
    }
}

/// Resolve the ACME directory URL for a given CA enum value.
pub fn resolve_directory_url(ca: &AcmeCa) -> &str {
    ca.directory_url()
}

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

    #[test]
    fn acme_client_creation() {
        let config = TlsConfig {
            acme_email: Some("test@example.com".into()),
            acme_ca: AcmeCa::LetsencryptStaging,
            ..Default::default()
        };
        let client = AcmeClient::new(&config).unwrap();
        assert_eq!(
            client.directory_url(),
            "https://acme-staging-v02.api.letsencrypt.org/directory"
        );
    }

    #[test]
    fn acme_client_with_custom_url() {
        let client = AcmeClient::with_directory_url(
            "https://custom-acme.example/dir",
            Some("me@example.com".into()),
        );
        assert_eq!(client.directory_url(), "https://custom-acme.example/dir");
    }

    #[test]
    fn resolve_directory_urls() {
        assert_eq!(
            resolve_directory_url(&AcmeCa::Letsencrypt),
            "https://acme-v02.api.letsencrypt.org/directory"
        );
        assert_eq!(
            resolve_directory_url(&AcmeCa::LetsencryptStaging),
            "https://acme-staging-v02.api.letsencrypt.org/directory"
        );
        assert_eq!(
            resolve_directory_url(&AcmeCa::Zerossl),
            "https://acme.zerossl.com/v2/DV90"
        );
        let custom = AcmeCa::Custom("https://internal/dir".into());
        assert_eq!(resolve_directory_url(&custom), "https://internal/dir");
    }

    #[test]
    fn challenge_struct_fields() {
        let ch = AcmeChallenge {
            domain: "example.com".into(),
            challenge_type: ChallengeType::Http01,
            token: "abc123".into(),
            key_authorization: "abc123.thumbprint".into(),
            url: "https://acme/chall/1".into(),
        };
        assert_eq!(ch.domain, "example.com");
        assert_eq!(ch.token, "abc123");
        assert_eq!(ch.key_authorization, "abc123.thumbprint");
    }

    #[test]
    fn challenge_type_variants() {
        let http = ChallengeType::Http01;
        let alpn = ChallengeType::TlsAlpn01;
        assert_ne!(http, alpn);
    }

    #[tokio::test]
    async fn provision_empty_domains_fails() {
        let config = TlsConfig {
            acme_ca: AcmeCa::LetsencryptStaging,
            ..Default::default()
        };
        let client = AcmeClient::new(&config).unwrap();
        let result = client.provision_cert(&[], |_| Ok(())).await;
        assert!(result.is_err());
        let err_msg = result.unwrap_err().to_string();
        assert!(err_msg.contains("at least one domain"), "got: {err_msg}");
    }

    #[test]
    fn default_challenge_type_is_http01() {
        let config = TlsConfig::default();
        let client = AcmeClient::new(&config).unwrap();
        assert_eq!(client.challenge_type, ChallengeType::Http01);
    }

    #[test]
    fn acme_challenge_clone() {
        let ch = AcmeChallenge {
            domain: "test.com".into(),
            challenge_type: ChallengeType::TlsAlpn01,
            token: "tok".into(),
            key_authorization: "ka".into(),
            url: "https://acme/c/1".into(),
        };
        let ch2 = ch.clone();
        assert_eq!(ch.domain, ch2.domain);
        assert_eq!(ch.token, ch2.token);
    }

    #[test]
    fn acme_client_stores_email() {
        let config = TlsConfig {
            acme_email: Some("user@example.com".into()),
            ..Default::default()
        };
        let client = AcmeClient::new(&config).unwrap();
        assert_eq!(client.email.as_deref(), Some("user@example.com"));
    }

    #[test]
    fn acme_client_no_email() {
        let config = TlsConfig::default();
        let client = AcmeClient::new(&config).unwrap();
        assert!(client.email.is_none());
    }
}
