//! Inbound federation handler: implements `FederationService` Cap'n Proto interface.
//!
//! Delegates all operations to the local [`Store`], acting as a trusted relay
//! from authenticated peer servers.
//!
//! **Security:** Each handler validates the request's `origin` field against
//! the `verified_peer_domain` extracted from the mTLS client certificate at
//! connection time. Per-peer rate limits prevent abuse.

use std::sync::Arc;

use capnp::capability::Promise;
use dashmap::DashMap;
use quicproquo_proto::federation_capnp::federation_service;
use tokio::sync::Notify;

use crate::auth::RateEntry;
use crate::storage::Store;

/// Per-peer federation rate limit: max requests within a 60-second window.
const FED_RATE_LIMIT_WINDOW_SECS: u64 = 60;
const FED_RATE_LIMIT_MAX: u32 = 200;

/// Inbound federation RPC handler.
pub struct FederationServiceImpl {
    pub store: Arc<dyn Store>,
    pub waiters: Arc<DashMap<Vec<u8>, Arc<Notify>>>,
    pub local_domain: String,
    /// The peer domain extracted from the mTLS client certificate's CN/SAN
    /// at connection time. All requests must declare an `origin` matching this.
    pub verified_peer_domain: Option<String>,
    /// Per-peer rate limiter (keyed by peer domain).
    pub rate_limits: Arc<DashMap<String, RateEntry>>,
}

/// Validate that the request's `origin` matches the mTLS-verified peer domain.
fn validate_origin(
    verified: &Option<String>,
    declared: &str,
) -> Result<(), capnp::Error> {
    match verified {
        Some(ref expected) if expected == declared => Ok(()),
        Some(ref expected) => Err(capnp::Error::failed(format!(
            "federation auth: origin '{}' does not match mTLS cert '{}'",
            declared, expected
        ))),
        None => Err(capnp::Error::failed(
            "federation auth: no verified peer domain (mTLS required)".into(),
        )),
    }
}

/// Extract and validate the origin string from the request's auth field.
fn extract_and_validate_origin(
    service: &FederationServiceImpl,
    get_auth: Result<quicproquo_proto::federation_capnp::federation_auth::Reader<'_>, capnp::Error>,
) -> Result<String, capnp::Error> {
    let auth = get_auth
        .map_err(|_| capnp::Error::failed("federation auth: missing auth field".into()))?;
    let origin_reader = auth.get_origin()
        .map_err(|_| capnp::Error::failed("federation auth: missing origin".into()))?;
    let origin = origin_reader.to_str()
        .map_err(|_| capnp::Error::failed("federation auth: origin is not valid UTF-8".into()))?;

    if origin.is_empty() {
        return Err(capnp::Error::failed("federation auth: origin must not be empty".into()));
    }

    validate_origin(&service.verified_peer_domain, origin)?;
    check_federation_rate_limit(&service.rate_limits, origin)?;

    Ok(origin.to_string())
}

/// Per-peer federation rate limiter.
fn check_federation_rate_limit(
    rate_limits: &DashMap<String, RateEntry>,
    peer_domain: &str,
) -> Result<(), capnp::Error> {
    let now = crate::auth::current_timestamp();
    let mut entry = rate_limits.entry(peer_domain.to_string()).or_insert(RateEntry {
        count: 0,
        window_start: now,
    });

    if now - entry.window_start >= FED_RATE_LIMIT_WINDOW_SECS {
        entry.count = 1;
        entry.window_start = now;
    } else {
        entry.count += 1;
        if entry.count > FED_RATE_LIMIT_MAX {
            return Err(capnp::Error::failed(format!(
                "federation rate limit exceeded for peer '{peer_domain}'"
            )));
        }
    }
    Ok(())
}

impl federation_service::Server for FederationServiceImpl {
    fn relay_enqueue(
        &mut self,
        params: federation_service::RelayEnqueueParams,
        mut results: federation_service::RelayEnqueueResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        // Validate origin against mTLS cert and apply rate limit.
        let origin = match extract_and_validate_origin(self, p.get_auth()) {
            Ok(o) => o,
            Err(e) => return Promise::err(e),
        };

        let recipient_key = match p.get_recipient_key() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad recipient_key: {e}"))),
        };
        let payload = match p.get_payload() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad payload: {e}"))),
        };
        let channel_id = p.get_channel_id().unwrap_or_default().to_vec();

        if recipient_key.len() != 32 {
            return Promise::err(capnp::Error::failed("recipient_key must be 32 bytes".into()));
        }
        if payload.is_empty() {
            return Promise::err(capnp::Error::failed("payload must not be empty".into()));
        }

        let seq = match self.store.enqueue(&recipient_key, &channel_id, payload, None) {
            Ok(s) => s,
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        };

        results.get().set_seq(seq);

        // Wake any waiting fetchWait clients.
        if let Some(waiter) = self.waiters.get(&recipient_key) {
            waiter.notify_waiters();
        }

        tracing::info!(
            origin = %origin,
            recipient_prefix = %hex::encode(&recipient_key[..4]),
            seq = seq,
            "federation: relayed enqueue"
        );

        Promise::ok(())
    }

    fn relay_batch_enqueue(
        &mut self,
        params: federation_service::RelayBatchEnqueueParams,
        mut results: federation_service::RelayBatchEnqueueResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        // Validate origin against mTLS cert and apply rate limit.
        let _origin = match extract_and_validate_origin(self, p.get_auth()) {
            Ok(o) => o,
            Err(e) => return Promise::err(e),
        };

        let recipient_keys = match p.get_recipient_keys() {
            Ok(v) => v,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad recipient_keys: {e}"))),
        };
        let payload = match p.get_payload() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad payload: {e}"))),
        };
        let channel_id = p.get_channel_id().unwrap_or_default().to_vec();

        let mut seqs = Vec::with_capacity(recipient_keys.len() as usize);
        for i in 0..recipient_keys.len() {
            let rk = match recipient_keys.get(i) {
                Ok(v) => v.to_vec(),
                Err(e) => return Promise::err(capnp::Error::failed(format!("bad key[{i}]: {e}"))),
            };
            if rk.len() != 32 {
                return Promise::err(capnp::Error::failed(
                    format!("recipient_key[{i}] must be 32 bytes"),
                ));
            }
            let seq = match self.store.enqueue(&rk, &channel_id, payload.clone(), None) {
                Ok(s) => s,
                Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
            };
            seqs.push(seq);
            if let Some(waiter) = self.waiters.get(&rk) {
                waiter.notify_waiters();
            }
        }

        let mut list = results.get().init_seqs(seqs.len() as u32);
        for (i, seq) in seqs.iter().enumerate() {
            list.set(i as u32, *seq);
        }

        tracing::info!(
            recipient_count = recipient_keys.len(),
            "federation: relayed batch_enqueue"
        );

        Promise::ok(())
    }

    fn proxy_fetch_key_package(
        &mut self,
        params: federation_service::ProxyFetchKeyPackageParams,
        mut results: federation_service::ProxyFetchKeyPackageResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        // Validate origin against mTLS cert and apply rate limit.
        if let Err(e) = extract_and_validate_origin(self, p.get_auth()) {
            return Promise::err(e);
        }

        let identity_key = match p.get_identity_key() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad identity_key: {e}"))),
        };

        match self.store.fetch_key_package(&identity_key) {
            Ok(Some(pkg)) => results.get().set_package(&pkg),
            Ok(None) => results.get().set_package(&[]),
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        }

        Promise::ok(())
    }

    fn proxy_fetch_hybrid_key(
        &mut self,
        params: federation_service::ProxyFetchHybridKeyParams,
        mut results: federation_service::ProxyFetchHybridKeyResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        // Validate origin against mTLS cert and apply rate limit.
        if let Err(e) = extract_and_validate_origin(self, p.get_auth()) {
            return Promise::err(e);
        }

        let identity_key = match p.get_identity_key() {
            Ok(v) => v.to_vec(),
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad identity_key: {e}"))),
        };

        match self.store.fetch_hybrid_key(&identity_key) {
            Ok(Some(pk)) => results.get().set_hybrid_public_key(&pk),
            Ok(None) => results.get().set_hybrid_public_key(&[]),
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        }

        Promise::ok(())
    }

    fn proxy_resolve_user(
        &mut self,
        params: federation_service::ProxyResolveUserParams,
        mut results: federation_service::ProxyResolveUserResults,
    ) -> Promise<(), capnp::Error> {
        let p = match params.get() {
            Ok(p) => p,
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        // Validate origin against mTLS cert and apply rate limit.
        if let Err(e) = extract_and_validate_origin(self, p.get_auth()) {
            return Promise::err(e);
        }

        let username = match p.get_username() {
            Ok(u) => match u.to_str() {
                Ok(s) => s.to_string(),
                Err(e) => return Promise::err(capnp::Error::failed(format!("bad utf-8: {e}"))),
            },
            Err(e) => return Promise::err(capnp::Error::failed(format!("bad params: {e}"))),
        };

        match self.store.get_user_identity_key(&username) {
            Ok(Some(key)) => results.get().set_identity_key(&key),
            Ok(None) => results.get().set_identity_key(&[]),
            Err(e) => return Promise::err(capnp::Error::failed(format!("store error: {e}"))),
        }

        Promise::ok(())
    }

    fn federation_health(
        &mut self,
        _params: federation_service::FederationHealthParams,
        mut results: federation_service::FederationHealthResults,
    ) -> Promise<(), capnp::Error> {
        // Health check does not require origin validation (diagnostic endpoint).
        results.get().set_status("ok");
        results.get().set_server_domain(&self.local_domain);
        Promise::ok(())
    }
}

/// Extract the peer domain from the mTLS client certificate's first SAN (DNS name)
/// or CN, given the QUIC connection's peer identity (a certificate chain).
pub fn extract_peer_domain(conn: &quinn::Connection) -> Option<String> {
    let identity = conn.peer_identity()?;
    let certs = identity.downcast::<Vec<rustls::pki_types::CertificateDer<'static>>>().ok()?;
    let first_cert = certs.first()?;

    // Parse the DER certificate to extract SAN DNS names or CN.
    let (_, parsed) = x509_parser::parse_x509_certificate(first_cert.as_ref()).ok()?;

    // Prefer SAN DNS names.
    if let Ok(Some(san)) = parsed.subject_alternative_name() {
        for name in &san.value.general_names {
            if let x509_parser::extensions::GeneralName::DNSName(dns) = name {
                return Some(dns.to_string());
            }
        }
    }

    // Fall back to CN.
    for rdn in parsed.subject().iter() {
        for attr in rdn.iter() {
            if attr.attr_type() == &x509_parser::oid_registry::OID_X509_COMMON_NAME {
                if let Ok(cn) = attr.as_str() {
                    return Some(cn.to_string());
                }
            }
        }
    }

    None
}
