//! TCP server implementation using mio for non-blocking I/O.

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

use kimberlite::Kimberlite;
use kimberlite_compliance::audit::{Actor, ComplianceAuditAction, ComponentName};
use kimberlite_types::TenantId;
use kimberlite_wire::{
    ApiKeyInfo, ApiKeyListResponse, ApiKeyRegisterResponse, ApiKeyRevokeResponse,
    ApiKeyRotateResponse, ClusterMode, ColumnInfo, ConsentCheckResponse, ConsentGrantResponse,
    ConsentListResponse, ConsentPurpose as WireConsentPurpose, ConsentRecord as WireConsentRecord,
    ConsentScope as WireConsentScope, ConsentWithdrawResponse, DescribeTableResponse,
    ErasureAuditInfo, ErasureCompleteResponse, ErasureExemptResponse,
    ErasureExemptionBasis as WireExemptionBasis, ErasureListResponse, ErasureMarkProgressResponse,
    ErasureMarkStreamErasedResponse, ErasureRequestInfo, ErasureRequestResponse,
    ErasureStatusResponse, ErasureStatusTag, ErrorCode, ListIndexesResponse, ListTablesResponse,
    Push, PushPayload, Request, RequestPayload, Response, ResponsePayload, ServerInfoResponse,
    SubscribeResponse, SubscriptionAckResponse, SubscriptionCloseReason, TableInfo,
    TenantCreateResponse, TenantDeleteResponse, TenantGetResponse, TenantInfo, TenantListResponse,
};
use mio::net::TcpListener;
use mio::{Events, Interest, Poll, Token};
use tracing::{debug, error, info, trace, warn};

#[cfg(unix)]
use signal_hook::consts::signal::{SIGINT, SIGTERM};
#[cfg(unix)]
use signal_hook_mio::v1_0::Signals;

use crate::auth::AuthService;
use crate::chaos::ChaosHandle;
use crate::config::ServerConfig;
use crate::connection::Connection;
use crate::error::{ServerError, ServerResult};
use crate::handler::RequestHandler;
use crate::health::HealthChecker;
use crate::http::HttpSidecar;
use crate::metrics;
use crate::replication::CommandSubmitter;
use crate::tenant_registry::{RegistryError, TenantRegistry};

/// Token for the listener socket.
const LISTENER_TOKEN: Token = Token(0);

/// Token for signal handling.
const SIGNAL_TOKEN: Token = Token(1);

/// Maximum events to process per poll iteration.
const MAX_EVENTS: usize = 1024;

/// Default shutdown drain timeout.
const SHUTDOWN_DRAIN_TIMEOUT: Duration = Duration::from_secs(30);

/// Maximum time the main event loop will block between subscription pumps.
///
/// Keeps tail-of-log push frames flowing even when the subscribed client is
/// otherwise idle and no I/O events fire.
const SUBSCRIPTION_POLL_INTERVAL: Duration = Duration::from_millis(50);

/// TCP server for `Kimberlite`.
///
/// Uses mio's poll-based event loop for handling multiple connections
/// without async runtimes.
pub struct Server {
    config: ServerConfig,
    poll: Poll,
    listener: TcpListener,
    connections: HashMap<Token, Connection>,
    handler: RequestHandler,
    next_token: usize,
    /// Whether shutdown has been requested.
    shutdown_requested: Arc<AtomicBool>,
    /// Health checker.
    health_checker: HealthChecker,
    /// HTTP sidecar for metrics/health endpoints.
    http_sidecar: Option<HttpSidecar>,
    /// Signal handler for SIGTERM/SIGINT (Unix only).
    #[cfg(unix)]
    signals: Option<Signals>,
    /// In-memory tenant registry backing the admin-API tenant ops.
    tenant_registry: Arc<TenantRegistry>,
    /// Wall-clock instant at which the event loop was constructed, used to
    /// report uptime via `GetServerInfo`.
    started_at: Instant,
}

impl Server {
    /// Creates a new server with the given configuration.
    pub fn new(config: ServerConfig, db: Kimberlite) -> ServerResult<Self> {
        let poll = Poll::new()?;

        // Bind the listener
        let addr = config.bind_addr;
        let mut listener =
            TcpListener::bind(addr).map_err(|e| ServerError::BindFailed { addr, source: e })?;

        // Register the listener with the poll
        poll.registry()
            .register(&mut listener, LISTENER_TOKEN, Interest::READABLE)?;

        // Create auth service
        let auth_service = AuthService::new(config.auth.clone());

        if matches!(config.auth, crate::auth::AuthMode::None) {
            warn!(
                "AuthMode::None is active — all connections accepted without authentication. \
                 Pass --allow-unauthenticated to suppress this warning in development."
            );
        }

        // Create health checker
        let health_checker = HealthChecker::new(&config.data_dir);

        // Create command submitter with replication mode. Wrapped in Arc
        // so the HTTP sidecar's chaos worker can share a handle with the
        // binary-protocol handler without either owning it exclusively.
        let submitter = Arc::new(CommandSubmitter::new(
            &config.replication,
            db,
            &config.data_dir,
        )?);

        if submitter.is_replicated() {
            info!(
                "Server listening on {} with {:?} replication",
                addr,
                submitter.status().mode
            );
        } else {
            info!("Server listening on {}", addr);
        }

        // Chaos HTTP endpoints activate only when the operator opts in via
        // env var. Production builds must NOT expose the probe contract by
        // default; chaos VMs set `KMB_ENABLE_CHAOS_ENDPOINTS=1` in their
        // init script.
        let chaos_enabled = std::env::var("KMB_ENABLE_CHAOS_ENDPOINTS")
            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
            .unwrap_or(false);

        let chaos_handle = if chaos_enabled {
            match submitter.subscribe_applied_commits(1024) {
                Some(rx) => {
                    info!("chaos HTTP endpoints enabled (KMB_ENABLE_CHAOS_ENDPOINTS=1)");
                    Some(ChaosHandle::spawn(
                        Arc::clone(&submitter),
                        rx,
                        config.data_dir.clone(),
                    ))
                }
                None => {
                    warn!(
                        "KMB_ENABLE_CHAOS_ENDPOINTS=1 set but replication is not in cluster \
                         mode — chaos endpoints require VSR commit fanout. Disabling."
                    );
                    None
                }
            }
        } else {
            None
        };

        // Bind HTTP sidecar for metrics/health/chaos if configured.
        let http_sidecar = if let Some(http_addr) = config.metrics_bind_addr {
            match HttpSidecar::bind_with_chaos(http_addr, &poll, chaos_handle) {
                Ok(sidecar) => Some(sidecar),
                Err(e) => {
                    warn!("Failed to bind HTTP sidecar on {http_addr}: {e}");
                    None
                }
            }
        } else {
            None
        };

        Ok(Self {
            config,
            poll,
            listener,
            connections: HashMap::new(),
            handler: RequestHandler::new(submitter, auth_service),
            next_token: 2, // Start at 2 since 0 is LISTENER_TOKEN and 1 is SIGNAL_TOKEN
            shutdown_requested: Arc::new(AtomicBool::new(false)),
            health_checker,
            http_sidecar,
            #[cfg(unix)]
            signals: None,
            tenant_registry: Arc::new(TenantRegistry::new()),
            started_at: Instant::now(),
        })
    }

    /// Creates a new server with signal handling enabled.
    ///
    /// On Unix: handles SIGTERM and SIGINT.
    /// On Windows: handles Ctrl+C and Ctrl+Break.
    pub fn with_signal_handling(config: ServerConfig, db: Kimberlite) -> ServerResult<Self> {
        let mut server = Self::new(config, db)?;

        #[cfg(unix)]
        {
            // Set up signal handling for SIGTERM and SIGINT
            let mut signals = Signals::new([SIGTERM, SIGINT]).map_err(ServerError::Io)?;

            // Register signals with the poll
            server
                .poll
                .registry()
                .register(&mut signals, SIGNAL_TOKEN, Interest::READABLE)?;

            server.signals = Some(signals);
            info!("Signal handling enabled (SIGTERM/SIGINT)");
        }

        #[cfg(windows)]
        {
            // Set up Ctrl+C handler for Windows
            let shutdown_flag = Arc::clone(&server.shutdown_requested);
            ctrlc::set_handler(move || {
                info!("Received Ctrl+C, initiating graceful shutdown");
                shutdown_flag.store(true, Ordering::SeqCst);
            })
            .map_err(|e| ServerError::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))?;

            info!("Signal handling enabled (Ctrl+C)");
        }

        Ok(server)
    }

    /// Returns the address the server is listening on.
    pub fn local_addr(&self) -> ServerResult<SocketAddr> {
        Ok(self.listener.local_addr()?)
    }

    /// Runs the server event loop.
    ///
    /// This method blocks until the server is shut down.
    pub fn run(&mut self) -> ServerResult<()> {
        let mut events = Events::with_capacity(MAX_EVENTS);

        info!("Server event loop started");

        loop {
            // Wait for events, but wake up at least every subscription
            // poll interval so active subscriptions get their tail pushed
            // even while the subscribed client is otherwise idle.
            let timeout = self.has_active_subscriptions()
                .then_some(SUBSCRIPTION_POLL_INTERVAL);
            if let Err(e) = self.poll.poll(&mut events, timeout) {
                if e.kind() == std::io::ErrorKind::Interrupted {
                    continue;
                }
                return Err(e.into());
            }

            // Process events
            for event in &events {
                match event.token() {
                    LISTENER_TOKEN => {
                        self.accept_connections()?;
                    }
                    crate::http::HTTP_LISTENER_TOKEN => {
                        if let Some(ref sidecar) = self.http_sidecar {
                            sidecar.handle_accept(&self.health_checker);
                        }
                    }
                    token => {
                        if event.is_readable() {
                            self.handle_readable(token)?;
                        }
                        if event.is_writable() {
                            self.handle_writable(token)?;
                        }
                    }
                }
            }

            // Pump subscriptions on timeout wakeups OR to drain any
            // newly-available events for subscriptions on other connections.
            if self.has_active_subscriptions() {
                self.pump_all_subscriptions();
                self.flush_pending_writes()?;
            }

            // Clean up closed connections
            self.cleanup_closed();
        }
    }

    /// True if any connection has at least one active subscription.
    fn has_active_subscriptions(&self) -> bool {
        self.connections.values().any(|c| !c.subscriptions.is_empty())
    }

    /// After queueing push frames, ensure each affected connection has
    /// WRITABLE interest registered so mio will flush the buffer.
    fn flush_pending_writes(&mut self) -> ServerResult<()> {
        let tokens: Vec<Token> = self
            .connections
            .iter()
            .filter(|(_, c)| !c.write_buf.is_empty())
            .map(|(t, _)| *t)
            .collect();
        for token in tokens {
            self.update_interest(token)?;
        }
        Ok(())
    }

    /// Runs a single iteration of the event loop.
    ///
    /// Useful for testing or custom event loops.
    pub fn poll_once(&mut self, timeout: Option<std::time::Duration>) -> ServerResult<()> {
        let mut events = Events::with_capacity(MAX_EVENTS);

        self.poll.poll(&mut events, timeout)?;

        for event in &events {
            match event.token() {
                LISTENER_TOKEN => {
                    self.accept_connections()?;
                }
                crate::http::HTTP_LISTENER_TOKEN => {
                    if let Some(ref sidecar) = self.http_sidecar {
                        sidecar.handle_accept(&self.health_checker);
                    }
                }
                token => {
                    if event.is_readable() {
                        self.handle_readable(token)?;
                    }
                    if event.is_writable() {
                        self.handle_writable(token)?;
                    }
                }
            }
        }

        self.cleanup_closed();
        Ok(())
    }

    /// Accepts new connections from the listener.
    fn accept_connections(&mut self) -> ServerResult<()> {
        loop {
            match self.listener.accept() {
                Ok((mut stream, addr)) => {
                    // Disable Nagle's algorithm for low-latency client messaging
                    if let Err(e) = stream.set_nodelay(true) {
                        tracing::warn!("Failed to set TCP_NODELAY for {}: {}", addr, e);
                    }

                    // Check connection limit
                    if self.connections.len() >= self.config.max_connections {
                        warn!(
                            "Max connections reached, rejecting connection from {}",
                            addr
                        );
                        // Just drop the stream to reject
                        continue;
                    }

                    // Allocate a token for this connection
                    let token = Token(self.next_token);
                    self.next_token += 1;

                    // Register the stream
                    self.poll
                        .registry()
                        .register(&mut stream, token, Interest::READABLE)?;

                    // Create the connection (with rate limiting if configured)
                    let conn = if let Some(rate_config) = self.config.rate_limit {
                        Connection::with_rate_limit(
                            token,
                            stream,
                            self.config.read_buffer_size,
                            rate_config,
                        )
                    } else {
                        Connection::new(token, stream, self.config.read_buffer_size)
                    };
                    self.connections.insert(token, conn);

                    // Record metrics
                    metrics::record_connection_accepted();

                    debug!("Accepted connection from {} (token {:?})", addr, token);
                }
                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    // No more connections to accept
                    break;
                }
                Err(e) => {
                    error!("Error accepting connection: {}", e);
                    break;
                }
            }
        }
        Ok(())
    }

    /// Intercepts protocol-v2 subscription lifecycle requests so they can
    /// operate on the per-connection `SubscriptionRegistry`.
    ///
    /// Returns `Some(response)` if the request was a subscription op; callers
    /// must not also dispatch through the stateless request handler.
    fn try_handle_subscription_request(
        &mut self,
        token: Token,
        request: &Request,
    ) -> Option<Response> {
        match &request.payload {
            RequestPayload::Subscribe(req) => {
                let conn = self.connections.get_mut(&token)?;
                // Validate the stream exists by reading zero bytes from it.
                if let Err(e) = self
                    .handler
                    .kimberlite()
                    .tenant(request.tenant_id)
                    .read_events(req.stream_id, req.from_offset, 0)
                {
                    return Some(Response::error(
                        request.id,
                        ErrorCode::StreamNotFound,
                        e.to_string(),
                    ));
                }
                let sub_id = conn.subscriptions.register(
                    request.tenant_id,
                    req.stream_id,
                    req.from_offset,
                    req.initial_credits,
                    req.consumer_group.clone(),
                );
                trace!(
                    "Registered subscription {} on stream {:?} for token {:?}",
                    sub_id, req.stream_id, token
                );
                Some(Response::new(
                    request.id,
                    ResponsePayload::Subscribe(SubscribeResponse {
                        subscription_id: sub_id,
                        start_offset: req.from_offset,
                        credits: req.initial_credits,
                    }),
                ))
            }
            RequestPayload::SubscribeCredit(req) => {
                let conn = self.connections.get_mut(&token)?;
                match conn
                    .subscriptions
                    .grant_credits(req.subscription_id, req.additional_credits)
                {
                    Some(new_balance) => Some(Response::new(
                        request.id,
                        ResponsePayload::SubscriptionAck(SubscriptionAckResponse {
                            subscription_id: req.subscription_id,
                            credits_remaining: new_balance,
                        }),
                    )),
                    None => Some(Response::error(
                        request.id,
                        ErrorCode::SubscriptionNotFound,
                        format!("subscription {} not found", req.subscription_id),
                    )),
                }
            }
            RequestPayload::Unsubscribe(req) => {
                let conn = self.connections.get_mut(&token)?;
                if conn.subscriptions.remove(req.subscription_id).is_some() {
                    let _ = conn.queue_push(Push::new(PushPayload::SubscriptionClosed {
                        subscription_id: req.subscription_id,
                        reason: SubscriptionCloseReason::ClientCancelled,
                    }));
                    Some(Response::new(
                        request.id,
                        ResponsePayload::SubscriptionAck(SubscriptionAckResponse {
                            subscription_id: req.subscription_id,
                            credits_remaining: 0,
                        }),
                    ))
                } else {
                    Some(Response::error(
                        request.id,
                        ErrorCode::SubscriptionNotFound,
                        format!("subscription {} not found", req.subscription_id),
                    ))
                }
            }
            _ => None,
        }
    }

    /// Intercepts Phase 4 admin + schema + server-info requests. All are
    /// admin-only (gated on `Role::Admin` or AuthMode::None); non-Admin
    /// callers receive `AuthenticationFailed`.
    fn try_handle_admin_request(
        &mut self,
        token: Token,
        request: &Request,
    ) -> Option<Response> {
        // Every request we match here also implies admin-only. We perform
        // the role check up-front so the later arms can assume admin.
        let is_admin = self.connection_is_admin(token);

        let require_admin = || {
            Response::error(
                request.id,
                ErrorCode::AuthenticationFailed,
                "admin operations require the Admin role".to_string(),
            )
        };

        match &request.payload {
            RequestPayload::ListTables(_) => Some(self.handle_list_tables(request)),
            RequestPayload::DescribeTable(req) => {
                Some(self.handle_describe_table(request, &req.table_name))
            }
            RequestPayload::ListIndexes(req) => {
                Some(self.handle_list_indexes(request, &req.table_name))
            }
            RequestPayload::TenantCreate(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_tenant_create(request, req.tenant_id, req.name.clone()))
            }
            RequestPayload::TenantList(_) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_tenant_list(request))
            }
            RequestPayload::TenantDelete(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_tenant_delete(request, req.tenant_id))
            }
            RequestPayload::TenantGet(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_tenant_get(request, req.tenant_id))
            }
            RequestPayload::ApiKeyRegister(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_api_key_register(
                    request,
                    req.subject.clone(),
                    req.tenant_id,
                    req.roles.clone(),
                    req.expires_at_nanos,
                ))
            }
            RequestPayload::ApiKeyRevoke(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_api_key_revoke(request, &req.key))
            }
            RequestPayload::ApiKeyList(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_api_key_list(request, req.tenant_id))
            }
            RequestPayload::ApiKeyRotate(req) => {
                if !is_admin {
                    return Some(require_admin());
                }
                Some(self.handle_api_key_rotate(request, &req.old_key))
            }
            RequestPayload::GetServerInfo(_) => Some(self.handle_server_info(request)),
            _ => None,
        }
    }

    fn connection_is_admin(&self, token: Token) -> bool {
        let auth_mode_is_none = matches!(
            self.handler.auth_service().mode(),
            crate::auth::AuthMode::None
        );
        if auth_mode_is_none {
            return true;
        }
        self.connections
            .get(&token)
            .and_then(|c| c.authenticated_identity.as_ref())
            .is_some_and(|id| id.roles.iter().any(|r| r.eq_ignore_ascii_case("Admin")))
    }

    /// Resolve the [`Actor`] for a request: prefer SDK-supplied
    /// `request.audit.actor`, fall back to the authenticated identity
    /// (via the first connection with a matching tenant), else attribute
    /// to the compliance-api system component.
    ///
    /// AUDIT-2026-04 S3.9 — this is the single choke point where wire
    /// `AuditMetadata` meets the server's `ComplianceAuditLog` so every
    /// compliance-tracked mutation preserves client attribution.
    fn actor_from_request(&self, request: &Request) -> Actor {
        if let Some(audit) = request.audit.as_ref() {
            if let Some(actor) = audit.actor.as_ref().filter(|s| !s.is_empty()) {
                return Actor::Authenticated(actor.clone());
            }
        }
        for conn in self.connections.values() {
            if let Some(id) = conn.authenticated_identity.as_ref() {
                if id.tenant_id == request.tenant_id {
                    return Actor::Authenticated(id.subject.clone());
                }
            }
        }
        Actor::System(ComponentName::Other("compliance-api".to_string()))
    }

    /// Append a compliance audit event with the SDK-supplied actor, best-effort.
    /// Failure to append is logged but does not fail the underlying request —
    /// the mutation has already committed and the client shouldn't be punished
    /// for a transient audit-log write hiccup.
    fn record_audit_event(&self, request: &Request, action: ComplianceAuditAction) {
        let actor = self.actor_from_request(request);
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        if let Err(e) = tenant.audit_log_append_with_actor(action, actor) {
            tracing::warn!(
                tenant_id = %request.tenant_id,
                request_id = request.id.0,
                error = %e,
                "failed to append compliance audit event"
            );
        }
    }

    fn handle_list_tables(&mut self, request: &Request) -> Response {
        self.tenant_registry.touch(request.tenant_id);
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        match tenant.query("SHOW TABLES", &[]) {
            Ok(result) => {
                let tables = result
                    .rows
                    .into_iter()
                    .filter_map(|row| {
                        let name = match row.first() {
                            Some(kimberlite_query::Value::Text(s)) => s.clone(),
                            _ => return None,
                        };
                        let column_count = match row.get(1) {
                            Some(kimberlite_query::Value::BigInt(n)) => *n as u32,
                            Some(kimberlite_query::Value::Integer(n)) => *n as u32,
                            _ => 0,
                        };
                        Some(TableInfo { name, column_count })
                    })
                    .collect();
                Response::new(
                    request.id,
                    ResponsePayload::ListTables(ListTablesResponse { tables }),
                )
            }
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_describe_table(&mut self, request: &Request, table_name: &str) -> Response {
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let sql = format!("SHOW COLUMNS FROM {table_name}");
        match tenant.query(&sql, &[]) {
            Ok(result) => {
                let columns = result
                    .rows
                    .into_iter()
                    .filter_map(|row| {
                        let name = match row.first() {
                            Some(kimberlite_query::Value::Text(s)) => s.clone(),
                            _ => return None,
                        };
                        let data_type = match row.get(1) {
                            Some(kimberlite_query::Value::Text(s)) => s.clone(),
                            _ => "UNKNOWN".to_string(),
                        };
                        let nullable = match row.get(2) {
                            Some(kimberlite_query::Value::Boolean(b)) => *b,
                            _ => true,
                        };
                        let primary_key = match row.get(3) {
                            Some(kimberlite_query::Value::Boolean(b)) => *b,
                            _ => false,
                        };
                        Some(ColumnInfo {
                            name,
                            data_type,
                            nullable,
                            primary_key,
                        })
                    })
                    .collect();
                Response::new(
                    request.id,
                    ResponsePayload::DescribeTable(DescribeTableResponse {
                        table_name: table_name.to_string(),
                        columns,
                    }),
                )
            }
            Err(e) => {
                let msg = e.to_string();
                let code = if msg.contains("not found") || msg.contains("does not exist") {
                    ErrorCode::TableNotFound
                } else {
                    ErrorCode::InternalError
                };
                Response::error(request.id, code, msg)
            }
        }
    }

    fn handle_list_indexes(&mut self, request: &Request, table_name: &str) -> Response {
        // Validate the table exists via SHOW TABLES. Returning index metadata
        // requires exposing `kernel_state.indexes()` through a `Kimberlite`
        // accessor; that exposure is tracked in ROADMAP v0.6 alongside a
        // proper `SHOW INDEXES` SQL grammar. Until then we return an empty
        // list for known tables so admin UIs can render a "no indexes" state
        // without surfacing a misleading error.
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let Ok(show) = tenant.query("SHOW TABLES", &[]) else {
            return Response::error(
                request.id,
                ErrorCode::InternalError,
                "failed to enumerate tables".to_string(),
            );
        };
        let exists = show.rows.iter().any(|row| {
            matches!(row.first(), Some(kimberlite_query::Value::Text(s)) if s == table_name)
        });
        if !exists {
            return Response::error(
                request.id,
                ErrorCode::TableNotFound,
                format!("table {table_name} not found"),
            );
        }
        Response::new(
            request.id,
            ResponsePayload::ListIndexes(ListIndexesResponse { indexes: Vec::new() }),
        )
    }

    fn handle_tenant_create(
        &mut self,
        request: &Request,
        tenant_id: TenantId,
        name: Option<String>,
    ) -> Response {
        match self.tenant_registry.register(tenant_id, name.clone()) {
            Ok((entry, created)) => {
                let info = TenantInfo {
                    tenant_id,
                    name: entry.name,
                    table_count: self.tenant_table_count(tenant_id),
                    created_at_nanos: Some(entry.created_at_nanos),
                };
                Response::new(
                    request.id,
                    ResponsePayload::TenantCreate(TenantCreateResponse { tenant: info, created }),
                )
            }
            Err(RegistryError::AlreadyExistsDifferentName { .. }) => Response::error(
                request.id,
                ErrorCode::TenantAlreadyExists,
                "tenant already exists with a different name".to_string(),
            ),
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_tenant_list(&mut self, request: &Request) -> Response {
        let tenants: Vec<TenantInfo> = self
            .tenant_registry
            .list()
            .into_iter()
            .map(|(tenant_id, entry)| TenantInfo {
                tenant_id,
                name: entry.name,
                table_count: self.tenant_table_count(tenant_id),
                created_at_nanos: Some(entry.created_at_nanos),
            })
            .collect();
        Response::new(
            request.id,
            ResponsePayload::TenantList(TenantListResponse { tenants }),
        )
    }

    fn handle_tenant_delete(&mut self, request: &Request, tenant_id: TenantId) -> Response {
        let tenant = self.handler.kimberlite().tenant(tenant_id);
        let mut tables_dropped = 0u32;
        if let Ok(result) = tenant.query("SHOW TABLES", &[]) {
            for row in result.rows {
                if let Some(kimberlite_query::Value::Text(name)) = row.first() {
                    let drop_sql = format!("DROP TABLE {name}");
                    if tenant.execute(&drop_sql, &[]).is_ok() {
                        tables_dropped += 1;
                    }
                }
            }
        }
        let deleted = self.tenant_registry.remove(tenant_id);
        Response::new(
            request.id,
            ResponsePayload::TenantDelete(TenantDeleteResponse {
                deleted,
                tables_dropped,
            }),
        )
    }

    fn handle_tenant_get(&mut self, request: &Request, tenant_id: TenantId) -> Response {
        match self.tenant_registry.get(tenant_id) {
            Some(entry) => Response::new(
                request.id,
                ResponsePayload::TenantGet(TenantGetResponse {
                    tenant: TenantInfo {
                        tenant_id,
                        name: entry.name,
                        table_count: self.tenant_table_count(tenant_id),
                        created_at_nanos: Some(entry.created_at_nanos),
                    },
                }),
            ),
            None => Response::error(
                request.id,
                ErrorCode::TenantNotFound,
                format!("tenant {} not registered", u64::from(tenant_id)),
            ),
        }
    }

    fn handle_api_key_register(
        &mut self,
        request: &Request,
        subject: String,
        tenant_id: TenantId,
        roles: Vec<String>,
        expires_at_nanos: Option<u64>,
    ) -> Response {
        let expires_at = expires_at_nanos.map(|ns| {
            std::time::UNIX_EPOCH + std::time::Duration::from_nanos(ns)
        });
        match self
            .handler
            .auth_service()
            .issue_api_key(subject, tenant_id, roles, expires_at)
        {
            Ok((key, listing)) => Response::new(
                request.id,
                ResponsePayload::ApiKeyRegister(ApiKeyRegisterResponse {
                    key,
                    info: ApiKeyInfo {
                        key_id: listing.key_id,
                        subject: listing.subject,
                        tenant_id: listing.tenant_id,
                        roles: listing.roles,
                        expires_at_nanos: listing.expires_at_nanos,
                    },
                }),
            ),
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_api_key_revoke(&mut self, request: &Request, key: &str) -> Response {
        match self.handler.auth_service().revoke_api_key(key) {
            Ok(revoked) => {
                if revoked {
                    Response::new(
                        request.id,
                        ResponsePayload::ApiKeyRevoke(ApiKeyRevokeResponse { revoked: true }),
                    )
                } else {
                    Response::error(
                        request.id,
                        ErrorCode::ApiKeyNotFound,
                        "API key not found".to_string(),
                    )
                }
            }
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_api_key_list(
        &mut self,
        request: &Request,
        tenant_filter: Option<TenantId>,
    ) -> Response {
        match self.handler.auth_service().list_api_keys(tenant_filter) {
            Ok(listings) => {
                let keys = listings
                    .into_iter()
                    .map(|l| ApiKeyInfo {
                        key_id: l.key_id,
                        subject: l.subject,
                        tenant_id: l.tenant_id,
                        roles: l.roles,
                        expires_at_nanos: l.expires_at_nanos,
                    })
                    .collect();
                Response::new(
                    request.id,
                    ResponsePayload::ApiKeyList(ApiKeyListResponse { keys }),
                )
            }
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_api_key_rotate(&mut self, request: &Request, old_key: &str) -> Response {
        match self.handler.auth_service().rotate_api_key(old_key) {
            Ok((new_key, listing)) => Response::new(
                request.id,
                ResponsePayload::ApiKeyRotate(ApiKeyRotateResponse {
                    new_key,
                    info: ApiKeyInfo {
                        key_id: listing.key_id,
                        subject: listing.subject,
                        tenant_id: listing.tenant_id,
                        roles: listing.roles,
                        expires_at_nanos: listing.expires_at_nanos,
                    },
                }),
            ),
            Err(_) => Response::error(
                request.id,
                ErrorCode::ApiKeyNotFound,
                "old API key not found".to_string(),
            ),
        }
    }

    fn handle_server_info(&mut self, request: &Request) -> Response {
        let cluster_mode = if self.handler.kimberlite_submitter_is_replicated() {
            ClusterMode::Clustered
        } else {
            ClusterMode::Standalone
        };
        let capabilities = vec![
            "query".to_string(),
            "append".to_string(),
            "subscribe.v2".to_string(),
            "admin.v1".to_string(),
            "schema.v1".to_string(),
            "server_info.v1".to_string(),
        ];
        Response::new(
            request.id,
            ResponsePayload::ServerInfo(ServerInfoResponse {
                build_version: env!("CARGO_PKG_VERSION").to_string(),
                protocol_version: kimberlite_wire::PROTOCOL_VERSION,
                capabilities,
                uptime_secs: self.started_at.elapsed().as_secs(),
                cluster_mode,
                tenant_count: self.tenant_registry.len() as u32,
            }),
        )
    }

    /// Intercepts Phase 5 compliance requests (consent + erasure).
    ///
    /// These ops are tenant-scoped and do not require Admin — any
    /// authenticated identity for that tenant can call them. Admin-role
    /// enforcement for cross-tenant scans is handled one layer up.
    fn try_handle_compliance_request(&mut self, request: &Request) -> Option<Response> {
        match &request.payload {
            RequestPayload::ConsentGrant(req) => Some(self.handle_consent_grant(
                request,
                req.subject_id.clone(),
                req.purpose,
                req.scope,
            )),
            RequestPayload::ConsentWithdraw(req) => {
                Some(self.handle_consent_withdraw(request, &req.consent_id))
            }
            RequestPayload::ConsentCheck(req) => Some(self.handle_consent_check(
                request,
                &req.subject_id,
                req.purpose,
            )),
            RequestPayload::ConsentList(req) => Some(self.handle_consent_list(
                request,
                &req.subject_id,
                req.valid_only,
            )),
            RequestPayload::ErasureRequest(req) => {
                Some(self.handle_erasure_request(request, &req.subject_id))
            }
            RequestPayload::ErasureMarkProgress(req) => Some(self.handle_erasure_mark_progress(
                request,
                &req.request_id,
                &req.streams,
            )),
            RequestPayload::ErasureMarkStreamErased(req) => {
                Some(self.handle_erasure_mark_stream_erased(
                    request,
                    &req.request_id,
                    req.stream_id,
                    req.records_erased,
                ))
            }
            RequestPayload::ErasureComplete(req) => {
                Some(self.handle_erasure_complete(request, &req.request_id))
            }
            RequestPayload::ErasureExempt(req) => {
                Some(self.handle_erasure_exempt(request, &req.request_id, req.basis))
            }
            RequestPayload::ErasureStatus(req) => {
                Some(self.handle_erasure_status(request, &req.request_id))
            }
            RequestPayload::ErasureList(_) => Some(self.handle_erasure_list(request)),

            // Phase 6 — audit / export / breach
            RequestPayload::AuditQuery(_) => Some(Response::error(
                request.id,
                ErrorCode::InternalError,
                "AuditQuery is wired at the wire-protocol level; server handler lands in v0.5.1"
                    .to_string(),
            )),
            RequestPayload::ExportSubject(_) => Some(Response::error(
                request.id,
                ErrorCode::InternalError,
                "ExportSubject wire surface defined; server handler lands in v0.5.1".to_string(),
            )),
            RequestPayload::VerifyExport(_) => Some(Response::error(
                request.id,
                ErrorCode::InternalError,
                "VerifyExport wire surface defined; server handler lands in v0.5.1".to_string(),
            )),
            RequestPayload::BreachReportIndicator(_)
            | RequestPayload::BreachQueryStatus(_)
            | RequestPayload::BreachConfirm(_)
            | RequestPayload::BreachResolve(_) => Some(Response::error(
                request.id,
                ErrorCode::InternalError,
                "Breach wire surface defined; server handlers land in v0.5.1".to_string(),
            )),
            _ => None,
        }
    }

    fn handle_consent_grant(
        &mut self,
        request: &Request,
        subject_id: String,
        purpose: WireConsentPurpose,
        scope: Option<WireConsentScope>,
    ) -> Response {
        self.tenant_registry.touch(request.tenant_id);
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let native_purpose = wire_to_native_purpose(purpose);
        let scope_label = scope
            .map(|s| format!("{s:?}"))
            .unwrap_or_else(|| "default".to_string());
        let result = tenant.grant_consent(&subject_id, native_purpose);
        drop(tenant);
        match result {
            Ok(consent_id) => {
                self.record_audit_event(
                    request,
                    ComplianceAuditAction::ConsentGranted {
                        subject_id: subject_id.clone(),
                        purpose: format!("{native_purpose:?}"),
                        scope: scope_label,
                    },
                );
                Response::new(
                    request.id,
                    ResponsePayload::ConsentGrant(ConsentGrantResponse {
                        consent_id: consent_id.to_string(),
                        granted_at_nanos: now_nanos_u64(),
                    }),
                )
            }
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_consent_withdraw(&mut self, request: &Request, consent_id: &str) -> Response {
        let uuid = match uuid::Uuid::parse_str(consent_id) {
            Ok(u) => u,
            Err(_) => {
                return Response::error(
                    request.id,
                    ErrorCode::InvalidRequest,
                    format!("invalid consent_id: {consent_id}"),
                );
            }
        };
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let subject_id = tenant
            .get_consent(uuid)
            .ok()
            .flatten()
            .map(|rec| rec.subject_id);
        let result = tenant.withdraw_consent(uuid);
        drop(tenant);
        match result {
            Ok(()) => {
                self.record_audit_event(
                    request,
                    ComplianceAuditAction::ConsentWithdrawn {
                        subject_id: subject_id.unwrap_or_default(),
                        consent_id: uuid,
                    },
                );
                Response::new(
                    request.id,
                    ResponsePayload::ConsentWithdraw(ConsentWithdrawResponse {
                        withdrawn_at_nanos: now_nanos_u64(),
                    }),
                )
            }
            Err(e) => {
                let msg = e.to_string();
                let code = if msg.contains("not found") {
                    ErrorCode::ConsentNotFound
                } else if msg.contains("expired") {
                    ErrorCode::ConsentExpired
                } else {
                    ErrorCode::InternalError
                };
                Response::error(request.id, code, msg)
            }
        }
    }

    fn handle_consent_check(
        &mut self,
        request: &Request,
        subject_id: &str,
        purpose: WireConsentPurpose,
    ) -> Response {
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let native_purpose = wire_to_native_purpose(purpose);
        match tenant.check_consent(subject_id, native_purpose) {
            Ok(is_valid) => Response::new(
                request.id,
                ResponsePayload::ConsentCheck(ConsentCheckResponse { is_valid }),
            ),
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_consent_list(
        &mut self,
        request: &Request,
        subject_id: &str,
        valid_only: bool,
    ) -> Response {
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        match tenant.get_consents_for_subject(subject_id) {
            Ok(records) => {
                let consents: Vec<WireConsentRecord> = records
                    .into_iter()
                    .filter(|r| {
                        !valid_only
                            || (r.withdrawn_at.is_none()
                                && r.expires_at.is_none_or(|t| t > chrono::Utc::now()))
                    })
                    .map(consent_record_to_wire)
                    .collect();
                Response::new(
                    request.id,
                    ResponsePayload::ConsentList(ConsentListResponse { consents }),
                )
            }
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_erasure_request(&mut self, request: &Request, subject_id: &str) -> Response {
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let result = tenant.request_erasure(subject_id);
        drop(tenant);
        match result {
            Ok(req) => {
                self.record_audit_event(
                    request,
                    ComplianceAuditAction::ErasureRequested {
                        subject_id: subject_id.to_string(),
                        request_id: req.request_id,
                    },
                );
                Response::new(
                    request.id,
                    ResponsePayload::ErasureRequest(ErasureRequestResponse {
                        request: erasure_request_to_wire(&req),
                    }),
                )
            }
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn handle_erasure_mark_progress(
        &mut self,
        request: &Request,
        request_id: &str,
        streams: &[kimberlite_types::StreamId],
    ) -> Response {
        let uuid = match uuid::Uuid::parse_str(request_id) {
            Ok(u) => u,
            Err(_) => {
                return Response::error(
                    request.id,
                    ErrorCode::InvalidRequest,
                    "invalid request_id".to_string(),
                );
            }
        };
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        if let Err(e) = tenant.mark_erasure_in_progress(uuid, streams.to_vec()) {
            return Response::error(
                request.id,
                erasure_error_code(&e.to_string()),
                e.to_string(),
            );
        }
        self.respond_with_erasure_request(request, uuid, |r| {
            ResponsePayload::ErasureMarkProgress(ErasureMarkProgressResponse {
                request: erasure_request_to_wire(r),
            })
        })
    }

    fn handle_erasure_mark_stream_erased(
        &mut self,
        request: &Request,
        request_id: &str,
        stream_id: kimberlite_types::StreamId,
        records_erased: u64,
    ) -> Response {
        let uuid = match uuid::Uuid::parse_str(request_id) {
            Ok(u) => u,
            Err(_) => {
                return Response::error(
                    request.id,
                    ErrorCode::InvalidRequest,
                    "invalid request_id".to_string(),
                );
            }
        };
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        if let Err(e) = tenant.mark_stream_erased(uuid, stream_id, records_erased) {
            return Response::error(request.id, erasure_error_code(&e.to_string()), e.to_string());
        }
        self.respond_with_erasure_request(request, uuid, |r| {
            ResponsePayload::ErasureMarkStreamErased(ErasureMarkStreamErasedResponse {
                request: erasure_request_to_wire(r),
            })
        })
    }

    fn handle_erasure_complete(&mut self, request: &Request, request_id: &str) -> Response {
        let uuid = match uuid::Uuid::parse_str(request_id) {
            Ok(u) => u,
            Err(_) => {
                return Response::error(
                    request.id,
                    ErrorCode::InvalidRequest,
                    "invalid request_id".to_string(),
                );
            }
        };
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let result = tenant.complete_erasure(uuid);
        drop(tenant);
        match result {
            Ok(audit) => {
                self.record_audit_event(
                    request,
                    ComplianceAuditAction::ErasureCompleted {
                        subject_id: audit.subject_id.clone(),
                        records_erased: audit.records_erased,
                        request_id: uuid,
                    },
                );
                Response::new(
                    request.id,
                    ResponsePayload::ErasureComplete(ErasureCompleteResponse {
                        audit: erasure_audit_to_wire(&audit),
                    }),
                )
            }
            Err(e) => {
                Response::error(request.id, erasure_error_code(&e.to_string()), e.to_string())
            }
        }
    }

    fn handle_erasure_exempt(
        &mut self,
        request: &Request,
        request_id: &str,
        basis: WireExemptionBasis,
    ) -> Response {
        let uuid = match uuid::Uuid::parse_str(request_id) {
            Ok(u) => u,
            Err(_) => {
                return Response::error(
                    request.id,
                    ErrorCode::InvalidRequest,
                    "invalid request_id".to_string(),
                );
            }
        };
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        let native_basis = wire_to_native_exemption(basis);
        let subject_id = tenant
            .get_erasure_request(uuid)
            .ok()
            .flatten()
            .map(|r| r.subject_id)
            .unwrap_or_default();
        let exempt_result = tenant.exempt_from_erasure(uuid, native_basis);
        drop(tenant);
        if let Err(e) = exempt_result {
            return Response::error(request.id, erasure_error_code(&e.to_string()), e.to_string());
        }
        self.record_audit_event(
            request,
            ComplianceAuditAction::ErasureExempted {
                subject_id,
                request_id: uuid,
                basis: format!("{native_basis:?}"),
            },
        );
        self.respond_with_erasure_request(request, uuid, |r| {
            ResponsePayload::ErasureExempt(ErasureExemptResponse {
                request: erasure_request_to_wire(r),
            })
        })
    }

    fn handle_erasure_status(&mut self, request: &Request, request_id: &str) -> Response {
        let uuid = match uuid::Uuid::parse_str(request_id) {
            Ok(u) => u,
            Err(_) => {
                return Response::error(
                    request.id,
                    ErrorCode::InvalidRequest,
                    "invalid request_id".to_string(),
                );
            }
        };
        self.respond_with_erasure_request(request, uuid, |r| {
            ResponsePayload::ErasureStatus(ErasureStatusResponse {
                request: erasure_request_to_wire(r),
            })
        })
    }

    fn handle_erasure_list(&mut self, request: &Request) -> Response {
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        match tenant.erasure_audit_trail() {
            Ok(records) => Response::new(
                request.id,
                ResponsePayload::ErasureList(ErasureListResponse {
                    audit: records.iter().map(erasure_audit_to_wire).collect(),
                }),
            ),
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    fn respond_with_erasure_request(
        &self,
        request: &Request,
        uuid: uuid::Uuid,
        build: impl FnOnce(&kimberlite_compliance::erasure::ErasureRequest) -> ResponsePayload,
    ) -> Response {
        let tenant = self.handler.kimberlite().tenant(request.tenant_id);
        match tenant.get_erasure_request(uuid) {
            Ok(Some(r)) => Response::new(request.id, build(&r)),
            Ok(None) => Response::error(
                request.id,
                ErrorCode::ErasureNotFound,
                format!("erasure request {uuid} not found"),
            ),
            Err(e) => Response::error(request.id, ErrorCode::InternalError, e.to_string()),
        }
    }

    /// Counts tables owned by a specific tenant. Since `kernel_state` is
    /// global (no tenant dimension), this runs a SQL `SHOW TABLES` on the
    /// tenant handle which the query engine scopes correctly.
    fn tenant_table_count(&self, tenant_id: TenantId) -> u32 {
        self.handler
            .kimberlite()
            .tenant(tenant_id)
            .query("SHOW TABLES", &[])
            .map(|r| r.rows.len() as u32)
            .unwrap_or(0)
    }

    /// Drains ready events for every active subscription on this connection
    /// and queues them as `Push` frames on the connection's write buffer.
    fn pump_connection_subscriptions(&mut self, token: Token) {
        let kb = self.handler.kimberlite();
        let Some(conn) = self.connections.get_mut(&token) else {
            return;
        };
        if conn.subscriptions.is_empty() || conn.closing {
            return;
        }
        let pushes = conn.subscriptions.pump(kb);
        for push in pushes {
            if let Err(e) = conn.queue_push(push) {
                error!("Error encoding push frame for {:?}: {}", token, e);
                conn.closing = true;
                break;
            }
        }
    }

    /// Pumps subscriptions for every connection. Called periodically on poll
    /// timeouts and after each request batch so even quiet clients make
    /// subscription progress.
    fn pump_all_subscriptions(&mut self) {
        let tokens: Vec<Token> = self.connections.keys().copied().collect();
        for token in tokens {
            self.pump_connection_subscriptions(token);
        }
    }

    /// Handles readable events on a connection.
    fn handle_readable(&mut self, token: Token) -> ServerResult<()> {
        let Some(conn) = self.connections.get_mut(&token) else {
            warn!("Readable event for unknown token {:?}", token);
            return Ok(());
        };

        // Update activity timestamp
        conn.touch();

        // Read data from the socket
        match conn.read() {
            Ok(true) => {
                // Connection still open, process requests
                self.process_requests(token);
                // Flush any ready subscription events to the same connection.
                self.pump_connection_subscriptions(token);
            }
            Ok(false) => {
                // Connection closed by peer
                debug!("Connection {:?} closed by peer", token);
                if let Some(c) = self.connections.get_mut(&token) {
                    c.closing = true;
                }
            }
            Err(e) => {
                error!("Error reading from {:?}: {}", token, e);
                if let Some(c) = self.connections.get_mut(&token) {
                    c.closing = true;
                }
            }
        }

        // Update interest if needed
        self.update_interest(token)?;
        Ok(())
    }

    /// Handles writable events on a connection.
    fn handle_writable(&mut self, token: Token) -> ServerResult<()> {
        let Some(conn) = self.connections.get_mut(&token) else {
            warn!("Writable event for unknown token {:?}", token);
            return Ok(());
        };

        match conn.write() {
            Ok(true) => {
                // All data written
                trace!("All data written to {:?}", token);
            }
            Ok(false) => {
                // More data to write
                trace!("More data to write to {:?}", token);
            }
            Err(e) => {
                error!("Error writing to {:?}: {}", token, e);
                conn.closing = true;
            }
        }

        // Update interest
        self.update_interest(token)?;
        Ok(())
    }

    /// Processes pending requests on a connection.
    fn process_requests(&mut self, token: Token) {
        loop {
            let Some(conn) = self.connections.get_mut(&token) else {
                return;
            };

            // Check if there's enough data for a frame
            if !conn.has_pending_data() {
                break;
            }

            // Try to decode a request
            match conn.try_decode_request() {
                Ok(Some(request)) => {
                    trace!("Received request {:?} from {:?}", request.id, token);

                    // Apply per-tenant rate limit if tenant tags are configured
                    // and this connection doesn't have a tenant-specific limiter yet
                    if let Some(ref tag_config) = self.config.tenant_tags {
                        let conn = self.connections.get_mut(&token).expect("just checked");
                        if conn.tenant_priority.is_none() {
                            let tenant_id = u64::from(request.tenant_id);
                            let priority = tag_config.priority_for(tenant_id);
                            let rate_config = tag_config.rate_limit_for(tenant_id);
                            conn.set_tenant_rate_limit(priority, rate_config);
                        }
                    }

                    // Check rate limit before processing
                    let Some(conn) = self.connections.get_mut(&token) else {
                        return;
                    };

                    if !conn.check_rate_limit() {
                        warn!("Rate limit exceeded for {:?}", token);
                        metrics::record_rate_limited();
                        let response = Response::error(
                            request.id,
                            ErrorCode::RateLimited,
                            "rate limit exceeded".to_string(),
                        );
                        if let Err(e) = conn.queue_response(&response) {
                            error!("Error encoding rate limit response: {}", e);
                            conn.closing = true;
                        }
                        continue;
                    }

                    // v2 subscription lifecycle — intercept before the
                    // stateless handler so we have access to the connection's
                    // subscription registry.
                    if let Some(response) = self.try_handle_subscription_request(token, &request) {
                        if let Some(c) = self.connections.get_mut(&token) {
                            if let Err(e) = c.queue_response(&response) {
                                error!("Error encoding subscription response: {}", e);
                                c.closing = true;
                            }
                        }
                        continue;
                    }

                    // Phase 4 admin/schema/server-info requests — intercepted
                    // here so they can touch the tenant registry + auth service
                    // without threading state through the stateless handler.
                    if let Some(response) = self.try_handle_admin_request(token, &request) {
                        if let Some(c) = self.connections.get_mut(&token) {
                            if let Err(e) = c.queue_response(&response) {
                                error!("Error encoding admin response: {}", e);
                                c.closing = true;
                            }
                        }
                        continue;
                    }

                    // Phase 5 compliance requests — consent + erasure.
                    if let Some(response) = self.try_handle_compliance_request(&request) {
                        if let Some(c) = self.connections.get_mut(&token) {
                            if let Err(e) = c.queue_response(&response) {
                                error!("Error encoding compliance response: {}", e);
                                c.closing = true;
                            }
                        }
                        continue;
                    }


                    // Clone the connection's current identity before calling the
                    // handler (we need &mut self.connections later).
                    let conn_identity_opt = self
                        .connections
                        .get(&token)
                        .and_then(|c| c.authenticated_identity.clone());

                    // Handle the request.
                    let (response, new_identity) =
                        self.handler.handle(request, conn_identity_opt.as_ref());

                    // Store the identity returned by a successful Handshake.
                    if let Some(identity) = new_identity {
                        if let Some(conn) = self.connections.get_mut(&token) {
                            conn.authenticated_identity = Some(identity);
                        }
                    }

                    // Queue the response
                    if let Some(c) = self.connections.get_mut(&token) {
                        if let Err(e) = c.queue_response(&response) {
                            error!("Error encoding response: {}", e);
                            c.closing = true;
                        }
                    }
                }
                Ok(None) => {
                    // Need more data
                    break;
                }
                Err(e) => {
                    error!("Error decoding request from {:?}: {}", token, e);
                    if let Some(c) = self.connections.get_mut(&token) {
                        c.closing = true;
                    }
                    break;
                }
            }
        }
    }

    /// Updates the interest flags for a connection.
    fn update_interest(&mut self, token: Token) -> ServerResult<()> {
        let Some(conn) = self.connections.get_mut(&token) else {
            return Ok(());
        };

        let interest = conn.interest();
        self.poll
            .registry()
            .reregister(&mut conn.stream, token, interest)?;

        Ok(())
    }

    /// Cleans up connections that have been marked as closing or are idle.
    fn cleanup_closed(&mut self) {
        let idle_timeout = self.config.idle_timeout;

        let to_close: Vec<Token> = self
            .connections
            .iter()
            .filter(|(_, c)| {
                if c.closing {
                    return true;
                }
                // Check idle timeout
                if let Some(timeout) = idle_timeout {
                    if c.is_idle(timeout) {
                        return true;
                    }
                }
                false
            })
            .map(|(t, _)| *t)
            .collect();

        for token in to_close {
            if let Some(mut conn) = self.connections.remove(&token) {
                if conn.closing {
                    debug!("Closing connection {:?}", token);
                } else {
                    debug!("Closing idle connection {:?}", token);
                }
                let _ = self.poll.registry().deregister(&mut conn.stream);

                // Record metrics
                metrics::record_connection_closed();
            }
        }
    }

    /// Returns the number of active connections.
    pub fn connection_count(&self) -> usize {
        self.connections.len()
    }

    /// Returns a handle that can be used to request shutdown from another thread.
    pub fn shutdown_handle(&self) -> ShutdownHandle {
        ShutdownHandle {
            shutdown_requested: Arc::clone(&self.shutdown_requested),
        }
    }

    /// Requests graceful shutdown.
    ///
    /// The server will stop accepting new connections and drain existing ones.
    pub fn shutdown(&self) {
        info!("Shutdown requested");
        self.shutdown_requested.store(true, Ordering::SeqCst);
    }

    /// Returns true if shutdown has been requested.
    pub fn is_shutdown_requested(&self) -> bool {
        self.shutdown_requested.load(Ordering::SeqCst)
    }

    /// Runs the server with graceful shutdown support.
    ///
    /// This method blocks until shutdown is requested and all connections are drained.
    /// If signal handling was enabled via `with_signal_handling`, the server will
    /// automatically shut down on SIGTERM or SIGINT.
    pub fn run_with_shutdown(&mut self) -> ServerResult<()> {
        let mut events = Events::with_capacity(MAX_EVENTS);

        info!("Server event loop started (with shutdown support)");

        loop {
            // Check if shutdown was requested
            if self.is_shutdown_requested() {
                info!("Shutdown requested, draining connections...");
                return self.drain_connections();
            }

            // Wait for events with a timeout to check shutdown periodically
            let timeout = Some(Duration::from_millis(100));
            if let Err(e) = self.poll.poll(&mut events, timeout) {
                if e.kind() == std::io::ErrorKind::Interrupted {
                    continue;
                }
                return Err(e.into());
            }

            // Process events
            for event in &events {
                match event.token() {
                    LISTENER_TOKEN => {
                        if !self.is_shutdown_requested() {
                            self.accept_connections()?;
                        }
                    }
                    SIGNAL_TOKEN => {
                        // Handle signals
                        self.handle_signals();
                    }
                    crate::http::HTTP_LISTENER_TOKEN => {
                        if let Some(ref sidecar) = self.http_sidecar {
                            sidecar.handle_accept(&self.health_checker);
                        }
                    }
                    token => {
                        if event.is_readable() {
                            self.handle_readable(token)?;
                        }
                        if event.is_writable() {
                            self.handle_writable(token)?;
                        }
                    }
                }
            }

            // Clean up closed connections
            self.cleanup_closed();
        }
    }

    /// Handles incoming signals (SIGTERM/SIGINT on Unix).
    #[cfg(unix)]
    fn handle_signals(&mut self) {
        if let Some(signals) = &mut self.signals {
            for signal in signals.pending() {
                match signal {
                    SIGTERM => {
                        info!("Received SIGTERM, initiating graceful shutdown");
                        self.shutdown();
                    }
                    SIGINT => {
                        info!("Received SIGINT, initiating graceful shutdown");
                        self.shutdown();
                    }
                    _ => {
                        debug!("Received signal {}, ignoring", signal);
                    }
                }
            }
        }
    }

    /// Stub for Windows - signal handling is done via ctrlc handler.
    #[cfg(windows)]
    fn handle_signals(&mut self) {
        // On Windows, Ctrl+C is handled by the ctrlc handler set in with_signal_handling
        // This method is called when SIGNAL_TOKEN is triggered, but on Windows that won't happen
    }

    /// Drains all active connections gracefully.
    ///
    /// Waits up to `SHUTDOWN_DRAIN_TIMEOUT` for connections to complete.
    fn drain_connections(&mut self) -> ServerResult<()> {
        let deadline = Instant::now() + SHUTDOWN_DRAIN_TIMEOUT;
        let mut events = Events::with_capacity(MAX_EVENTS);

        // Mark all connections as closing (no new requests)
        for conn in self.connections.values_mut() {
            conn.closing = true;
        }

        // Continue processing until all connections are drained or timeout
        while !self.connections.is_empty() && Instant::now() < deadline {
            let remaining = deadline.saturating_duration_since(Instant::now());
            let timeout = Some(remaining.min(Duration::from_millis(100)));

            if let Err(e) = self.poll.poll(&mut events, timeout) {
                if e.kind() == std::io::ErrorKind::Interrupted {
                    continue;
                }
                return Err(e.into());
            }

            for event in &events {
                let token = event.token();
                if token == LISTENER_TOKEN {
                    continue; // Don't accept new connections
                }
                if event.is_readable() {
                    let _ = self.handle_readable(token);
                }
                if event.is_writable() {
                    let _ = self.handle_writable(token);
                }
            }

            self.cleanup_closed();
        }

        let remaining = self.connections.len();
        if remaining > 0 {
            warn!(
                "Shutdown timeout reached with {} connections still active",
                remaining
            );
        } else {
            info!("All connections drained successfully");
        }

        Ok(())
    }

    /// Returns the health checker.
    pub fn health_checker(&self) -> &HealthChecker {
        &self.health_checker
    }

    /// Returns the authentication service.
    pub fn auth_service(&self) -> &AuthService {
        self.handler.auth_service()
    }

    /// Returns the Prometheus metrics.
    pub fn metrics(&self) -> String {
        metrics::Metrics::global().render()
    }
}

/// A handle that can be used to request shutdown from another thread.
#[derive(Clone)]
pub struct ShutdownHandle {
    shutdown_requested: Arc<AtomicBool>,
}

impl ShutdownHandle {
    /// Requests graceful shutdown.
    pub fn shutdown(&self) {
        self.shutdown_requested.store(true, Ordering::SeqCst);
    }

    /// Returns true if shutdown has been requested.
    pub fn is_shutdown_requested(&self) -> bool {
        self.shutdown_requested.load(Ordering::SeqCst)
    }
}

// ============================================================================
// Phase 5 — helpers for wire ↔ compliance-crate type conversions.
// ============================================================================

fn now_nanos_u64() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .ok()
        .and_then(|d| u64::try_from(d.as_nanos()).ok())
        .unwrap_or(0)
}

fn wire_to_native_purpose(
    purpose: WireConsentPurpose,
) -> kimberlite_compliance::purpose::Purpose {
    use kimberlite_compliance::purpose::Purpose;
    match purpose {
        WireConsentPurpose::Marketing => Purpose::Marketing,
        WireConsentPurpose::Analytics => Purpose::Analytics,
        WireConsentPurpose::Contractual => Purpose::Contractual,
        WireConsentPurpose::LegalObligation => Purpose::LegalObligation,
        WireConsentPurpose::VitalInterests => Purpose::VitalInterests,
        WireConsentPurpose::PublicTask => Purpose::PublicTask,
        WireConsentPurpose::Research => Purpose::Research,
        WireConsentPurpose::Security => Purpose::Security,
    }
}

fn native_to_wire_purpose(
    p: kimberlite_compliance::purpose::Purpose,
) -> WireConsentPurpose {
    use kimberlite_compliance::purpose::Purpose;
    match p {
        Purpose::Marketing => WireConsentPurpose::Marketing,
        Purpose::Analytics => WireConsentPurpose::Analytics,
        Purpose::Contractual => WireConsentPurpose::Contractual,
        Purpose::LegalObligation => WireConsentPurpose::LegalObligation,
        Purpose::VitalInterests => WireConsentPurpose::VitalInterests,
        Purpose::PublicTask => WireConsentPurpose::PublicTask,
        Purpose::Research => WireConsentPurpose::Research,
        Purpose::Security => WireConsentPurpose::Security,
    }
}

fn native_to_wire_scope(
    s: kimberlite_compliance::consent::ConsentScope,
) -> WireConsentScope {
    use kimberlite_compliance::consent::ConsentScope as N;
    match s {
        N::AllData => WireConsentScope::AllData,
        N::ContactInfo => WireConsentScope::ContactInfo,
        N::AnalyticsOnly => WireConsentScope::AnalyticsOnly,
        N::ContractualNecessity => WireConsentScope::ContractualNecessity,
    }
}

fn wire_to_native_exemption(
    basis: WireExemptionBasis,
) -> kimberlite_compliance::erasure::ExemptionBasis {
    use kimberlite_compliance::erasure::ExemptionBasis as E;
    match basis {
        WireExemptionBasis::LegalObligation => E::LegalObligation,
        WireExemptionBasis::PublicHealth => E::PublicHealth,
        WireExemptionBasis::Archiving => E::Archiving,
        WireExemptionBasis::LegalClaims => E::LegalClaims,
    }
}

fn native_to_wire_exemption(
    basis: kimberlite_compliance::erasure::ExemptionBasis,
) -> WireExemptionBasis {
    use kimberlite_compliance::erasure::ExemptionBasis as E;
    match basis {
        E::LegalObligation => WireExemptionBasis::LegalObligation,
        E::PublicHealth => WireExemptionBasis::PublicHealth,
        E::Archiving => WireExemptionBasis::Archiving,
        E::LegalClaims => WireExemptionBasis::LegalClaims,
    }
}

fn consent_record_to_wire(
    r: kimberlite_compliance::consent::ConsentRecord,
) -> WireConsentRecord {
    WireConsentRecord {
        consent_id: r.consent_id.to_string(),
        subject_id: r.subject_id,
        purpose: native_to_wire_purpose(r.purpose),
        scope: native_to_wire_scope(r.scope),
        granted_at_nanos: datetime_to_nanos(r.granted_at),
        withdrawn_at_nanos: r.withdrawn_at.map(datetime_to_nanos),
        expires_at_nanos: r.expires_at.map(datetime_to_nanos),
        notes: r.notes,
    }
}

fn erasure_request_to_wire(
    r: &kimberlite_compliance::erasure::ErasureRequest,
) -> ErasureRequestInfo {
    ErasureRequestInfo {
        request_id: r.request_id.to_string(),
        subject_id: r.subject_id.clone(),
        requested_at_nanos: datetime_to_nanos(r.requested_at),
        deadline_nanos: datetime_to_nanos(r.deadline),
        status: erasure_status_to_wire(&r.status),
        records_erased: r.records_erased,
        streams_affected: r.affected_streams.clone(),
    }
}

fn erasure_status_to_wire(
    s: &kimberlite_compliance::erasure::ErasureStatus,
) -> ErasureStatusTag {
    use kimberlite_compliance::erasure::ErasureStatus;
    match s {
        ErasureStatus::Pending => ErasureStatusTag::Pending,
        ErasureStatus::InProgress { streams_remaining } => ErasureStatusTag::InProgress {
            streams_remaining: *streams_remaining as u32,
        },
        ErasureStatus::Complete { erased_at, total_records } => ErasureStatusTag::Complete {
            erased_at_nanos: datetime_to_nanos(*erased_at),
            total_records: *total_records,
        },
        ErasureStatus::Failed { reason, retry_at } => ErasureStatusTag::Failed {
            reason: reason.clone(),
            retry_at_nanos: datetime_to_nanos(*retry_at),
        },
        ErasureStatus::Exempt { basis } => ErasureStatusTag::Exempt {
            basis: native_to_wire_exemption(*basis),
        },
    }
}

fn erasure_audit_to_wire(
    r: &kimberlite_compliance::erasure::ErasureAuditRecord,
) -> ErasureAuditInfo {
    ErasureAuditInfo {
        request_id: r.request_id.to_string(),
        subject_id: r.subject_id.clone(),
        requested_at_nanos: datetime_to_nanos(r.requested_at),
        completed_at_nanos: r
            .completed_at
            .map(datetime_to_nanos)
            .unwrap_or(0),
        records_erased: r.records_erased,
        streams_affected: r.streams_affected.clone(),
        erasure_proof_hex: r.erasure_proof.as_ref().map(hex_encode_hash),
    }
}

fn datetime_to_nanos(dt: chrono::DateTime<chrono::Utc>) -> u64 {
    let secs = dt.timestamp();
    let sub_nanos = u64::from(dt.timestamp_subsec_nanos());
    if secs < 0 {
        return 0;
    }
    let secs = secs as u64;
    secs.saturating_mul(1_000_000_000).saturating_add(sub_nanos)
}

fn hex_encode_hash(h: &kimberlite_types::Hash) -> String {
    h.as_bytes()
        .iter()
        .map(|b| format!("{b:02x}"))
        .collect::<String>()
}

fn erasure_error_code(msg: &str) -> ErrorCode {
    if msg.contains("not found") {
        ErrorCode::ErasureNotFound
    } else if msg.contains("already complete") || msg.contains("already Complete") {
        ErrorCode::ErasureAlreadyComplete
    } else if msg.contains("Exempt") || msg.contains("exempt") {
        ErrorCode::ErasureExempt
    } else {
        ErrorCode::InternalError
    }
}
