//! qpq-server — unified Authentication + Delivery service.
//!
//! The server hosts Authentication + Delivery services over QUIC + Cap'n Proto.

use std::{net::IpAddr, net::SocketAddr, path::PathBuf, sync::Arc};

use anyhow::Context;
use clap::Parser;
use dashmap::DashMap;
use opaque_ke::ServerSetup;
use quicproquo_core::opaque_auth::OpaqueSuite;
use quicproquo_kt::MerkleLog;
use quinn::Endpoint;
use rand::rngs::OsRng;
use tokio::sync::Notify;
use tokio::task::LocalSet;

pub mod audit;
mod auth;
mod config;
pub mod domain;
mod error_codes;
mod federation;
pub mod hooks;
mod metrics;
mod node_service;
#[allow(unsafe_code)] // FFI: C-ABI plugin interaction requires unsafe blocks
mod plugin_loader;
mod sql_store;
mod tls;
mod storage;
pub mod v2_handlers;
mod ws_bridge;
#[cfg(feature = "webtransport")]
mod webtransport;

use auth::{AuthConfig, PendingLogin, RateEntry, SessionInfo};
use config::{
    load_config, merge_config, validate_production_config, DEFAULT_DATA_DIR, DEFAULT_DB_PATH,
    DEFAULT_LISTEN, DEFAULT_STORE_BACKEND, DEFAULT_TLS_CERT, DEFAULT_TLS_KEY,
};
use node_service::{handle_node_connection, spawn_cleanup_task};
use sql_store::SqlStore;
use storage::{FileBackedStore, Store};
use tls::build_server_config;

// ── CLI ───────────────────────────────────────────────────────────────────────

#[derive(Debug, Parser)]
#[command(
    name = "qpq-server",
    about = "quicproquo Delivery Service + Authentication Service",
    version
)]
struct Args {
    /// Optional path to a TOML config file (fields map to CLI flags).
    #[arg(long, env = "QPQ_CONFIG")]
    config: Option<PathBuf>,

    /// QUIC listen address (host:port).
    #[arg(long, default_value = DEFAULT_LISTEN, env = "QPQ_LISTEN")]
    listen: String,

    /// Directory for persisted server data (KeyPackages + delivery queues).
    #[arg(long, default_value = DEFAULT_DATA_DIR, env = "QPQ_DATA_DIR")]
    data_dir: String,

    /// TLS certificate path (generated automatically if missing).
    #[arg(long, default_value = DEFAULT_TLS_CERT, env = "QPQ_TLS_CERT")]
    tls_cert: PathBuf,

    /// TLS private key path (generated automatically if missing).
    #[arg(long, default_value = DEFAULT_TLS_KEY, env = "QPQ_TLS_KEY")]
    tls_key: PathBuf,

    /// Required bearer token for auth.version=1 requests. Use --allow-insecure-auth to run without it (dev only).
    #[arg(long, env = "QPQ_AUTH_TOKEN")]
    auth_token: Option<String>,

    /// Allow running without QPQ_AUTH_TOKEN (development only).
    #[arg(long, env = "QPQ_ALLOW_INSECURE_AUTH", default_value_t = false)]
    allow_insecure_auth: bool,

    /// Enable Sealed Sender: enqueue does not require identity-bound session, only a valid token.
    #[arg(long, env = "QPQ_SEALED_SENDER", default_value_t = false)]
    sealed_sender: bool,

    /// Storage backend: "file" (bincode) or "sql" (SQLCipher-encrypted).
    #[arg(long, default_value = DEFAULT_STORE_BACKEND, env = "QPQ_STORE_BACKEND")]
    store_backend: String,

    /// Path to the SQLCipher database file (only used when --store-backend=sql).
    #[arg(long, default_value = DEFAULT_DB_PATH, env = "QPQ_DB_PATH")]
    db_path: PathBuf,

    /// SQLCipher encryption key. Empty string disables encryption.
    #[arg(long, default_value = "", env = "QPQ_DB_KEY")]
    db_key: String,

    /// Metrics HTTP listen address (e.g. 0.0.0.0:9090). If set and metrics enabled, /metrics is served.
    #[arg(long, env = "QPQ_METRICS_LISTEN")]
    metrics_listen: Option<String>,

    /// Enable metrics server when metrics_listen is set.
    #[arg(long, env = "QPQ_METRICS_ENABLED")]
    metrics_enabled: Option<bool>,

    /// Enable federation (server-to-server message relay).
    #[arg(long, env = "QPQ_FEDERATION_ENABLED", default_value_t = false)]
    federation_enabled: bool,

    /// This server's domain for federation addressing (e.g. "chat.example.com").
    #[arg(long, env = "QPQ_FEDERATION_DOMAIN")]
    federation_domain: Option<String>,

    /// Federation QUIC listen address (default: 0.0.0.0:7001).
    #[arg(long, env = "QPQ_FEDERATION_LISTEN")]
    federation_listen: Option<String>,

    /// Directory containing plugin `.so` / `.dylib` files to load at startup.
    /// Each library must export `extern "C" fn qpq_plugin_init(vtable: *mut HookVTable) -> i32`.
    #[arg(long, env = "QPQ_PLUGIN_DIR")]
    plugin_dir: Option<PathBuf>,

    /// Redact identity key prefixes and payload sizes in audit logs for metadata minimization.
    #[arg(long, env = "QPQ_REDACT_LOGS", default_value_t = false)]
    redact_logs: bool,

    /// WebSocket JSON-RPC bridge listen address (e.g. 0.0.0.0:9000). Enables browser connectivity.
    #[arg(long, env = "QPQ_WS_LISTEN")]
    ws_listen: Option<String>,

    /// WebTransport (HTTP/3) listen address for browser clients (e.g. 0.0.0.0:7443).
    /// Requires --features webtransport.
    #[arg(long, env = "QPQ_WEBTRANSPORT_LISTEN")]
    webtransport_listen: Option<String>,

    /// Graceful shutdown drain timeout in seconds (default: 30). In-flight RPCs get this
    /// long to finish after a shutdown signal before connections are forcefully closed.
    #[arg(long, env = "QPQ_DRAIN_TIMEOUT", default_value_t = config::DEFAULT_DRAIN_TIMEOUT_SECS)]
    drain_timeout: u64,

    /// Default per-RPC timeout in seconds (default: 30). Individual methods may override.
    #[arg(long, env = "QPQ_RPC_TIMEOUT", default_value_t = config::DEFAULT_RPC_TIMEOUT_SECS)]
    rpc_timeout: u64,

    /// Storage/database operation timeout in seconds (default: 10).
    #[arg(long, env = "QPQ_STORAGE_TIMEOUT", default_value_t = config::DEFAULT_STORAGE_TIMEOUT_SECS)]
    storage_timeout: u64,
}

// ── Entry point ───────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let _ = rustls::crypto::ring::default_provider().install_default();

    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    let args = Args::parse();
    let file_cfg = load_config(args.config.as_deref())?;
    let effective = merge_config(&args, &file_cfg);

    let production = std::env::var("QPQ_PRODUCTION")
        .map(|v| matches!(v.to_lowercase().as_str(), "1" | "true" | "yes"))
        .unwrap_or(false);
    if production {
        validate_production_config(&effective)?;
    }

    // Optional metrics server: only start when metrics_enabled and metrics_listen are set.
    if effective.metrics_enabled {
        if let Some(addr_str) = &effective.metrics_listen {
            let addr: std::net::SocketAddr = addr_str
                .parse()
                .context("metrics_listen must be host:port (e.g. 0.0.0.0:9090)")?;
            metrics_exporter_prometheus::PrometheusBuilder::new()
                .with_http_listener(addr)
                .install()
                .context("failed to install Prometheus metrics exporter")?;
            tracing::info!(addr = %addr_str, "metrics server listening on /metrics");
        }
    }

    // In non-production, require an explicit opt-out before running without a static token.
    if !production
        && effective
            .auth_token
            .as_deref()
            .map(|s| s.is_empty())
            .unwrap_or(true)
        && !effective.allow_insecure_auth
    {
        anyhow::bail!(
            "missing QPQ_AUTH_TOKEN; set one or pass --allow-insecure-auth for development"
        );
    }

    if effective.allow_insecure_auth
        && effective
            .auth_token
            .as_deref()
            .map(|s| s.is_empty())
            .unwrap_or(true)
    {
        tracing::warn!("running without QPQ_AUTH_TOKEN (allow-insecure-auth enabled); development only");
    }

    let listen: SocketAddr = effective
        .listen
        .parse()
        .context("--listen must be host:port")?;

    let mut server_config = build_server_config(&effective.tls_cert, &effective.tls_key, production)
        .context("failed to build TLS/QUIC server config")?;

    // Harden QUIC transport: idle timeout, limit stream concurrency.
    let mut transport = quinn::TransportConfig::default();
    transport.max_idle_timeout(Some(
        std::time::Duration::from_secs(300)
            .try_into()
            .map_err(|e| anyhow::anyhow!("idle timeout: {e}"))?,
    ));
    transport.max_concurrent_bidi_streams(1u32.into());
    transport.max_concurrent_uni_streams(0u32.into());
    server_config.transport_config(Arc::new(transport));

    // Shared storage — persisted to disk for restart safety.
    let store: Arc<dyn Store> = match effective.store_backend.as_str() {
        "sql" => {
            if let Some(parent) = effective.db_path.parent() {
                std::fs::create_dir_all(parent).context("create db dir")?;
            }
            tracing::info!(
                path = %effective.db_path.display(),
                encrypted = !effective.db_key.is_empty(),
                "opening SQLCipher store"
            );
            if effective.db_key.is_empty() {
                tracing::warn!("db_key is empty; SQL store will be plaintext (development only)");
            }
            Arc::new(SqlStore::open(&effective.db_path, &effective.db_key)?)
        }
        _ => {
            tracing::info!(dir = %effective.data_dir, "opening file-backed store");
            Arc::new(FileBackedStore::open(&effective.data_dir)?)
        }
    };

    // Ensure blobs directory exists for file transfer support.
    std::fs::create_dir_all(PathBuf::from(&effective.data_dir).join("blobs"))
        .context("create blobs directory")?;

    let auth_cfg = Arc::new(AuthConfig::new(
        effective.auth_token.clone(),
        effective.allow_insecure_auth,
    ));
    let waiters: Arc<DashMap<Vec<u8>, Arc<Notify>>> = Arc::new(DashMap::new());

    // OPAQUE ServerSetup: load from storage or generate fresh.
    let opaque_setup: Arc<ServerSetup<OpaqueSuite>> = match store.get_server_setup() {
        Ok(Some(bytes)) => {
            let setup = ServerSetup::<OpaqueSuite>::deserialize(&bytes)
                .map_err(|e| anyhow::anyhow!("corrupt OPAQUE server setup: {e}"))?;
            tracing::info!("loaded persisted OPAQUE ServerSetup");
            Arc::new(setup)
        }
        Ok(None) => {
            let setup = ServerSetup::<OpaqueSuite>::new(&mut OsRng);
            let bytes = setup.serialize().to_vec();
            store
                .store_server_setup(bytes)
                .context("persist OPAQUE ServerSetup")?;
            tracing::info!("generated and persisted new OPAQUE ServerSetup");
            Arc::new(setup)
        }
        Err(e) => return Err(anyhow::anyhow!("load OPAQUE server setup: {e}")),
    };

    // Server Ed25519 signing key for delivery proofs: load from storage or generate fresh.
    let signing_key: Arc<quicproquo_core::IdentityKeypair> = match store.get_signing_key_seed() {
        Ok(Some(seed_bytes)) => {
            let seed: [u8; 32] = seed_bytes
                .as_slice()
                .try_into()
                .context("signing key seed must be 32 bytes")?;
            tracing::info!("loaded persisted server signing key");
            Arc::new(quicproquo_core::IdentityKeypair::from_seed(seed))
        }
        Ok(None) => {
            let kp = quicproquo_core::IdentityKeypair::generate();
            store
                .store_signing_key_seed(kp.seed_bytes().to_vec())
                .context("persist server signing key")?;
            tracing::info!("generated and persisted new server signing key");
            Arc::new(kp)
        }
        Err(e) => return Err(anyhow::anyhow!("load server signing key: {e}")),
    };

    // Key Transparency Merkle log: load from storage or start fresh.
    let kt_log: Arc<std::sync::Mutex<MerkleLog>> = match store.load_kt_log() {
        Ok(Some(bytes)) => {
            match MerkleLog::from_bytes(&bytes) {
                Ok(log) => {
                    tracing::info!(entries = log.len(), "loaded persisted KT Merkle log");
                    Arc::new(std::sync::Mutex::new(log))
                }
                Err(e) => {
                    tracing::warn!(error = %e, "KT log deserialise failed; starting fresh");
                    Arc::new(std::sync::Mutex::new(MerkleLog::new()))
                }
            }
        }
        Ok(None) => {
            tracing::info!("no KT log found; starting fresh");
            Arc::new(std::sync::Mutex::new(MerkleLog::new()))
        }
        Err(e) => return Err(anyhow::anyhow!("load KT log: {e}")),
    };

    // ── Plugin hooks ──────────────────────────────────────────────────────────
    let hooks: Arc<dyn hooks::ServerHooks> = if let Some(dir) = &effective.plugin_dir {
        let plugins = plugin_loader::load_plugins_from_dir(dir);
        if plugins.is_empty() {
            tracing::info!(dir = %dir.display(), "plugin_dir set but no plugins loaded");
            Arc::new(hooks::NoopHooks)
        } else {
            tracing::info!(count = plugins.len(), "plugins loaded");
            let boxed: Vec<Box<dyn hooks::ServerHooks>> = plugins
                .into_iter()
                .map(|p| Box::new(p) as Box<dyn hooks::ServerHooks>)
                .collect();
            Arc::new(plugin_loader::ChainedHooks::new(boxed))
        }
    } else {
        Arc::new(hooks::NoopHooks)
    };

    let pending_logins: Arc<DashMap<String, PendingLogin>> = Arc::new(DashMap::new());
    let sessions: Arc<DashMap<Vec<u8>, SessionInfo>> = Arc::new(DashMap::new());
    let rate_limits: Arc<DashMap<Vec<u8>, RateEntry>> = Arc::new(DashMap::new());
    let conn_rate_limits: Arc<DashMap<IpAddr, RateEntry>> = Arc::new(DashMap::new());

    // Background cleanup task (expire sessions, pending logins, rate limits, and stale messages).
    spawn_cleanup_task(
        Arc::clone(&sessions),
        Arc::clone(&pending_logins),
        Arc::clone(&rate_limits),
        Arc::clone(&conn_rate_limits),
        Arc::clone(&store),
        Arc::clone(&waiters),
    );

    // ── WebSocket JSON-RPC bridge ──────────────────────────────────────────
    if let Some(ws_addr_str) = &effective.ws_listen {
        let ws_addr: SocketAddr = ws_addr_str
            .parse()
            .context("--ws-listen must be host:port (e.g. 0.0.0.0:9000)")?;
        let ws_state = Arc::new(ws_bridge::WsBridgeState {
            store: Arc::clone(&store),
            waiters: Arc::clone(&waiters),
            auth_cfg: Arc::clone(&auth_cfg),
            sessions: Arc::clone(&sessions),
            rate_limits: Arc::clone(&rate_limits),
            sealed_sender: effective.sealed_sender,
            allow_insecure_auth: effective.allow_insecure_auth,
        });
        ws_bridge::spawn_ws_bridge(ws_addr, ws_state);
    }

    // ── WebTransport (HTTP/3) endpoint ─────────────────────────────────────
    #[cfg(feature = "webtransport")]
    if let Some(wt_addr_str) = &effective.webtransport_listen {
        let wt_addr: SocketAddr = wt_addr_str
            .parse()
            .context("--webtransport-listen must be host:port (e.g. 0.0.0.0:7443)")?;

        let wt_server_config = webtransport::build_webtransport_server_config(
            &effective.tls_cert,
            &effective.tls_key,
        )
        .context("build WebTransport server config")?;

        let wt_state = Arc::new(v2_handlers::ServerState {
            store: Arc::clone(&store),
            waiters: Arc::clone(&waiters),
            auth_cfg: Arc::clone(&auth_cfg),
            opaque_setup: Arc::clone(&opaque_setup),
            pending_logins: Arc::clone(&pending_logins),
            sessions: Arc::clone(&sessions),
            rate_limits: Arc::clone(&rate_limits),
            sealed_sender: effective.sealed_sender,
            hooks: Arc::clone(&hooks),
            signing_key: Arc::clone(&signing_key),
            kt_log: Arc::clone(&kt_log),
            revocation_log: Arc::new(std::sync::Mutex::new(
                quicproquo_kt::RevocationLog::new(),
            )),
            data_dir: PathBuf::from(&effective.data_dir),
            redact_logs: effective.redact_logs,
            audit_logger: Arc::new(audit::NoopAuditLogger),
            draining: Arc::new(std::sync::atomic::AtomicBool::new(false)),
            seen_message_ids: Arc::new(DashMap::new()),
            banned_users: Arc::new(DashMap::new()),
            moderation_reports: Arc::new(std::sync::Mutex::new(Vec::new())),
            node_id: format!("wt-{}", hex::encode(&signing_key.public_key_bytes()[..4])),
            start_time: std::time::Instant::now(),
            storage_backend: effective.store_backend.clone(),
            federation_client: None,
            local_domain: effective.federation.as_ref().map(|f| f.domain.clone()).unwrap_or_default(),
        });

        let wt_registry = Arc::new(v2_handlers::build_registry(
            std::time::Duration::from_secs(effective.rpc_timeout_secs),
        ));

        webtransport::spawn_webtransport_listener(wt_addr, wt_server_config, wt_state, wt_registry)
            .context("spawn WebTransport listener")?;
    }

    let endpoint = Endpoint::server(server_config, listen)?;

    tracing::info!(
        addr = %effective.listen,
        "accepting QUIC connections"
    );

    // ── Federation setup ─────────────────────────────────────────────────────
    let federation_client: Option<Arc<federation::FederationClient>> =
        if let Some(fed_cfg) = &effective.federation {
            tracing::info!(
                domain = %fed_cfg.domain,
                listen = %fed_cfg.listen,
                peers = fed_cfg.peers.len(),
                "federation enabled"
            );

            // Build a client endpoint for outbound federation connections.
            // For now we create a simple endpoint; full mTLS is used when certs are provided.
            let client_config = if fed_cfg.federation_cert.exists()
                && fed_cfg.federation_key.exists()
                && fed_cfg.federation_ca.exists()
            {
                let tls_cfg = federation::tls::build_federation_client_config(
                    &fed_cfg.federation_cert,
                    &fed_cfg.federation_key,
                    &fed_cfg.federation_ca,
                )
                .context("build federation client TLS config")?;

                let crypto = quinn::crypto::rustls::QuicClientConfig::try_from(tls_cfg)
                    .map_err(|e| anyhow::anyhow!("invalid federation client QUIC config: {e}"))?;
                let mut qc = quinn::ClientConfig::new(Arc::new(crypto));
                let mut transport = quinn::TransportConfig::default();
                transport.max_idle_timeout(Some(
                    std::time::Duration::from_secs(120)
                        .try_into()
                        .map_err(|e| anyhow::anyhow!("idle timeout: {e}"))?,
                ));
                qc.transport_config(Arc::new(transport));
                Some(qc)
            } else {
                tracing::warn!("federation cert/key/CA not found; outbound federation connections will fail");
                None
            };

            let fed_bind: SocketAddr = SocketAddr::from(([0, 0, 0, 0], 0));
            let mut fed_endpoint = Endpoint::client(fed_bind)
                .context("create federation client endpoint")?;
            if let Some(cc) = client_config {
                fed_endpoint.set_default_client_config(cc);
            }

            let client = federation::FederationClient::new(fed_cfg, fed_endpoint)
                .context("create federation client")?;

            // Register configured peers in storage.
            for peer in &fed_cfg.peers {
                if let Err(e) = store.upsert_federation_peer(&peer.domain, true) {
                    tracing::warn!(domain = %peer.domain, error = %e, "failed to register federation peer");
                }
            }

            Some(Arc::new(client))
        } else {
            None
        };

    let local_domain: Option<String> = effective.federation.as_ref().map(|f| f.domain.clone());

    // ── Federation listener ──────────────────────────────────────────────────
    let federation_endpoint: Option<Endpoint> =
        if let Some(fed_cfg) = &effective.federation {
            if fed_cfg.federation_cert.exists()
                && fed_cfg.federation_key.exists()
                && fed_cfg.federation_ca.exists()
            {
                let fed_server_config = federation::tls::build_federation_server_config(
                    &fed_cfg.federation_cert,
                    &fed_cfg.federation_key,
                    &fed_cfg.federation_ca,
                )
                .context("build federation server TLS config")?;

                let fed_listen: SocketAddr = fed_cfg
                    .listen
                    .parse()
                    .context("federation listen must be host:port")?;

                let ep = Endpoint::server(fed_server_config, fed_listen)
                    .context("bind federation QUIC endpoint")?;

                tracing::info!(addr = %fed_cfg.listen, "federation endpoint listening");
                Some(ep)
            } else {
                tracing::warn!("federation certs not found; federation listener not started");
                None
            }
        } else {
            None
        };

    // ── mDNS local mesh discovery ─────────────────────────────────────────────
    // Announce this server on the local network so mesh-mode clients (and other
    // Freifunk nodes) can discover it automatically without manual configuration.
    // Non-critical: failures are logged as warnings; the server starts regardless.
    let _mdns_daemon = {
        let listen_port: u16 = listen.port();
        // Use the federation domain as the mDNS instance name when available.
        let mdns_instance = effective
            .federation
            .as_ref()
            .map(|f| f.domain.clone())
            .unwrap_or_else(|| "qpq-server".to_string());
        // mDNS host names must end with a dot.
        let mdns_host = if mdns_instance.ends_with('.') {
            mdns_instance.clone()
        } else {
            format!("{mdns_instance}.local.")
        };

        match mdns_sd::ServiceDaemon::new() {
            Ok(daemon) => {
                let mut props = std::collections::HashMap::new();
                props.insert("ver".to_string(), "1".to_string());
                props.insert("server".to_string(), effective.listen.clone());
                props.insert("domain".to_string(), mdns_instance.clone());

                match mdns_sd::ServiceInfo::new(
                    "_quicproquo._udp.local.",
                    &mdns_instance,
                    &mdns_host,
                    &[] as &[std::net::IpAddr],
                    listen_port,
                    Some(props),
                ) {
                    Ok(info) => match daemon.register(info) {
                        Ok(()) => {
                            tracing::info!(
                                instance = %mdns_instance,
                                port = listen_port,
                                "mDNS: announced qpq server on local network (_quicproquo._udp.local.)"
                            );
                        }
                        Err(e) => {
                            tracing::warn!(error = %e, "mDNS: service registration failed; mesh discovery disabled");
                        }
                    },
                    Err(e) => {
                        tracing::warn!(error = %e, "mDNS: failed to build service info; mesh discovery disabled");
                    }
                }
                Some(daemon)
            }
            Err(e) => {
                tracing::warn!(error = %e, "mDNS: daemon start failed; mesh discovery disabled");
                None
            }
        }
    };

    // capnp-rpc is !Send (Rc internals), so all RPC tasks must stay on a LocalSet.
    let local = LocalSet::new();
    local
        .run_until(async move {
            // Spawn federation acceptor if enabled.
            if let Some(fed_ep) = federation_endpoint {
                let fed_store = Arc::clone(&store);
                let fed_waiters = Arc::clone(&waiters);
                let fed_domain = local_domain.clone().unwrap_or_default();
                tokio::task::spawn_local(async move {
                    loop {
                        let incoming = match fed_ep.accept().await {
                            Some(i) => i,
                            None => break,
                        };
                        let connecting = match incoming.accept() {
                            Ok(c) => c,
                            Err(e) => {
                                tracing::warn!(error = %e, "federation: accept error");
                                continue;
                            }
                        };
                        let store = Arc::clone(&fed_store);
                        let waiters = Arc::clone(&fed_waiters);
                        let domain = fed_domain.clone();
                        tokio::task::spawn_local(async move {
                            match connecting.await {
                                Ok(conn) => {
                                    tracing::info!(
                                        peer = %conn.remote_address(),
                                        "federation: peer connected"
                                    );
                                    let (send, recv) = match conn.accept_bi().await {
                                        Ok(s) => s,
                                        Err(e) => {
                                            tracing::warn!(error = %e, "federation: accept bi error");
                                            return;
                                        }
                                    };
                                    let reader = tokio_util::compat::TokioAsyncReadCompatExt::compat(recv);
                                    let writer = tokio_util::compat::TokioAsyncWriteCompatExt::compat_write(send);

                                    let network = capnp_rpc::twoparty::VatNetwork::new(
                                        reader,
                                        writer,
                                        capnp_rpc::rpc_twoparty_capnp::Side::Server,
                                        Default::default(),
                                    );

                                    let verified_peer_domain =
                                        federation::service::extract_peer_domain(&conn);
                                    if let Some(ref peer) = verified_peer_domain {
                                        tracing::info!(peer_domain = %peer, "federation: mTLS peer authenticated");
                                    } else {
                                        tracing::warn!(peer = %conn.remote_address(), "federation: could not extract peer domain from mTLS cert");
                                    }

                                    let service_impl = federation::service::FederationServiceImpl {
                                        store,
                                        waiters,
                                        local_domain: domain,
                                        verified_peer_domain,
                                        rate_limits: Arc::new(dashmap::DashMap::new()),
                                    };
                                    let client: quicproquo_proto::federation_capnp::federation_service::Client =
                                        capnp_rpc::new_client(service_impl);

                                    if let Err(e) = capnp_rpc::RpcSystem::new(
                                        Box::new(network),
                                        Some(client.client),
                                    ).await {
                                        tracing::warn!(error = %e, "federation: RPC error");
                                    }
                                }
                                Err(e) => {
                                    tracing::warn!(error = %e, "federation: connection error");
                                }
                            }
                        });
                    }
                });
            }

            loop {
                tokio::select! {
                    biased;

                    incoming = endpoint.accept() => {
                        let incoming = match incoming {
                            Some(i) => i,
                            None => break,
                        };

                        // Per-IP connection rate limiting.
                        let remote_ip = incoming.remote_address().ip();
                        if !auth::check_conn_rate_limit(&conn_rate_limits, remote_ip) {
                            tracing::warn!(ip = %remote_ip, "connection rate limit exceeded, dropping");
                            incoming.refuse();
                            continue;
                        }

                        let connecting = match incoming.accept() {
                            Ok(c) => c,
                            Err(e) => {
                                tracing::warn!(error = %e, "failed to accept incoming connection");
                                continue;
                            }
                        };

                        let store = Arc::clone(&store);
                        let waiters = Arc::clone(&waiters);
                        let auth_cfg = Arc::clone(&auth_cfg);
                        let opaque_setup = Arc::clone(&opaque_setup);
                        let pending_logins = Arc::clone(&pending_logins);
                        let sessions = Arc::clone(&sessions);
                        let rate_limits = Arc::clone(&rate_limits);
                        let sealed_sender = effective.sealed_sender;
                        let fed_client = federation_client.clone();
                        let local_dom = local_domain.clone();
                        let sk = Arc::clone(&signing_key);
                        let conn_hooks = Arc::clone(&hooks);
                        let conn_kt_log = Arc::clone(&kt_log);
                        let conn_data_dir = PathBuf::from(&effective.data_dir);
                        let conn_redact_logs = effective.redact_logs;

                        tokio::task::spawn_local(async move {
                            if let Err(e) = handle_node_connection(
                                connecting,
                                store,
                                waiters,
                                auth_cfg,
                                opaque_setup,
                                pending_logins,
                                sessions,
                                rate_limits,
                                sealed_sender,
                                fed_client,
                                local_dom,
                                sk,
                                conn_hooks,
                                conn_kt_log,
                                conn_data_dir,
                                conn_redact_logs,
                            )
                            .await
                            {
                                tracing::warn!(error = %e, "connection error");
                            }
                        });
                    }

                    _ = shutdown_signal() => {
                        tracing::info!("shutdown signal received, draining QUIC connections");
                        // Stop accepting new connections immediately.
                        endpoint.close(0u32.into(), b"server shutdown");
                        break;
                    }
                }
            }

            // Grace period: let in-flight RPC tasks on the LocalSet finish.
            let drain_secs = effective.drain_timeout_secs;
            tracing::info!(drain_timeout_secs = drain_secs, "waiting for in-flight RPCs to complete");
            tokio::time::sleep(std::time::Duration::from_secs(drain_secs)).await;

            Ok::<(), anyhow::Error>(())
        })
        .await?;

    Ok(())
}

/// Wait for either SIGINT (Ctrl-C) or SIGTERM (Unix only).
///
/// Load balancers typically send SIGTERM during rolling deploys. The server
/// should stop accepting new connections, return "draining" from the health
/// endpoint, and wait for in-flight RPCs to finish (up to the drain timeout).
async fn shutdown_signal() {
    let ctrl_c = tokio::signal::ctrl_c();

    #[cfg(unix)]
    {
        let mut sigterm =
            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
                .expect("failed to install SIGTERM handler");
        tokio::select! {
            _ = ctrl_c => {},
            _ = sigterm.recv() => {},
        }
    }

    #[cfg(not(unix))]
    {
        ctrl_c.await.ok();
    }
}
