use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use bytes::Bytes;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tracing;

use crate::auth::{AuthSession, Authenticator};
use crate::codec::Frame;
use crate::codec::frame::{
    CAP_BATCH_ACK, CAP_HEADERS, CAP_NACK, CAP_WILDCARD, Capabilities, ErrorCode, FetchedMessage,
    ReconnectStatus, SERVER_CAPABILITIES, SubscriptionResult, SubscriptionStatus,
};
use crate::config::SharedConfig;
use crate::namespace::{NackResult, NamespaceManager};
use crate::types::{
    ClientId, GroupName, MessageEnvelope, MessageId, Permission, PublishedMessage, SubscribeStart,
    TopicName, TopicOffset, now_millis,
};

/// Maximum number of inflight (delivered but un-acked) messages tracked per
/// connection. Beyond this limit, new deliveries are still sent but their
/// message IDs are not tracked — acks for those messages will return an error,
/// and the namespace manager's inflight timeout will reclaim them.
const MAX_INFLIGHT: usize = 65_536;

/// Reap dead forwarder tasks every N frames to avoid scanning every time.
const REAP_INTERVAL: u64 = 64;

/// Outcome of processing one inbound frame.
#[derive(Debug)]
pub(crate) enum FrameOutcome {
    /// Send this frame back to the client.
    Reply(Frame),
    /// Send multiple frames back to the client in order.
    ReplyMany(Vec<Frame>),
    /// Send this frame and then close the connection.
    ReplyAndClose(Frame),
    /// No response needed.
    None,
}

/// Per-connection state before authentication.
pub(crate) struct PreAuthState {
    client_id: ClientId,
    message_id_counter: Arc<AtomicU64>,
    max_frame_size: u32,
    shared_config: SharedConfig,
}

/// Per-connection state after successful authentication.
pub(crate) struct AuthenticatedState {
    client_id: ClientId,
    auth: AuthSession,
    /// Negotiated protocol capabilities for this session.
    pub(crate) capabilities: Capabilities,
    /// Max frame size for this connection (used to compute fetch response budget).
    max_frame_size: u32,
    /// Maps (topic, group) → forwarder task handle.
    subscriptions: HashMap<(TopicName, Option<GroupName>), JoinHandle<()>>,
    /// Maps delivered MessageId → TopicName (needed because Frame::Ack only
    /// carries MessageId, but manager.ack() requires the topic).
    inflight_topics: HashMap<MessageId, TopicName>,
    message_id_counter: Arc<AtomicU64>,
    /// Counter for periodic dead-forwarder reaping.
    frame_counter: u64,
    /// Fire-once flag for 80% inflight warning; resets when inflight drops below threshold.
    inflight_warning_fired: bool,
    /// Shared runtime config — read per-publish for hot-reloadable TTL limits.
    shared_config: SharedConfig,
}

/// Per-connection session state machine.
///
/// Uses an enum to make the authentication transition a type-level guarantee:
/// post-auth operations cannot be called on a pre-auth session, and there are
/// no `Option::unwrap()` calls on the auth state.
pub(crate) enum Session {
    PreAuth(PreAuthState),
    Authenticated(AuthenticatedState),
    /// Server is draining. Publish/Ack/Unsubscribe/Ping/Pong allowed;
    /// Subscribe/SubscribeAt/SubscribeFrom rejected with ServerDraining.
    Draining(AuthenticatedState),
}

impl Session {
    pub fn new(
        client_id: ClientId,
        message_id_counter: Arc<AtomicU64>,
        max_frame_size: u32,
        shared_config: SharedConfig,
    ) -> Self {
        Self::PreAuth(PreAuthState {
            client_id,
            message_id_counter,
            max_frame_size,
            shared_config,
        })
    }

    pub fn client_id(&self) -> ClientId {
        match self {
            Self::PreAuth(s) => s.client_id,
            Self::Authenticated(s) | Self::Draining(s) => s.client_id,
        }
    }

    /// Returns negotiated capabilities for this session (0 if not yet authenticated).
    pub fn capabilities(&self) -> Capabilities {
        match self {
            Self::PreAuth(_) => 0,
            Self::Authenticated(s) | Self::Draining(s) => s.capabilities,
        }
    }

    /// Transition from Authenticated to Draining. No-op if not authenticated.
    pub fn transition_to_draining(&mut self) {
        // Use take() pattern: swap out Authenticated, wrap as Draining.
        let taken = std::mem::replace(
            self,
            Session::PreAuth(PreAuthState {
                client_id: ClientId(0),
                message_id_counter: Arc::new(AtomicU64::new(0)),
                max_frame_size: 0,
                shared_config: crate::config::shared_config(&crate::config::ServerConfig::default()),
            }),
        );
        match taken {
            Self::Authenticated(state) => *self = Self::Draining(state),
            other => *self = other, // was not Authenticated — restore original
        }
    }

    /// Process one inbound frame. Dispatches to auth/manager as needed.
    pub async fn process_frame<M: NamespaceManager>(
        &mut self,
        frame: Frame,
        auth: &dyn Authenticator,
        manager: &M,
        delivery_tx: &mpsc::Sender<PublishedMessage>,
    ) -> FrameOutcome {
        // 1. Reject server-only frames from clients.
        match &frame {
            Frame::ConnectAck { .. }
            | Frame::SubscribeAck { .. }
            | Frame::UnsubscribeAck { .. }
            | Frame::PublishAck { .. }
            | Frame::Deliver { .. }
            | Frame::GoingAway { .. }
            | Frame::ReconnectAck { .. }
            | Frame::FetchResponse { .. }
            | Frame::NackAck { .. }
            | Frame::BatchAckResponse { .. } => {
                return FrameOutcome::ReplyAndClose(Frame::Error {
                    code: ErrorCode::InvalidFrame,
                    message: "server-only frame received from client".into(),
                });
            }
            _ => {}
        }

        // 2. Ping is always allowed, even before auth.
        if matches!(&frame, Frame::Ping) {
            return FrameOutcome::Reply(Frame::Pong);
        }

        // 3. Pong — no response needed.
        if matches!(&frame, Frame::Pong) {
            return FrameOutcome::None;
        }

        match self {
            Session::PreAuth(pre) => {
                // Reconnect is handled here (not in process_pre_auth) because
                // it needs the manager for batch subscribes and delivery_tx for
                // spawning forwarder tasks.
                if let Frame::Reconnect {
                    version,
                    namespace,
                    api_key,
                    capabilities,
                    subscriptions,
                } = frame
                {
                    return Self::process_reconnect(
                        self,
                        version,
                        namespace,
                        api_key,
                        capabilities,
                        subscriptions,
                        auth,
                        manager,
                        delivery_tx,
                    )
                    .await;
                }

                let (outcome, new_state) = Self::process_pre_auth(pre, frame, auth);
                if let Some(authed) = new_state {
                    let ns = &authed.auth.namespace;
                    match manager.connect(ns, authed.client_id).await {
                        Ok(()) => {
                            *self = Session::Authenticated(authed);
                            outcome
                        }
                        Err(e) => {
                            // manager.connect() already emits fafb_connections_rejected_total
                            FrameOutcome::ReplyAndClose(Frame::ConnectAck {
                                success: false,
                                message: e.to_string(),
                                server_capabilities: 0,
                            })
                        }
                    }
                } else {
                    // Auth failed — emit rejection metric if it was a connect attempt
                    if matches!(
                        &outcome,
                        FrameOutcome::ReplyAndClose(Frame::ConnectAck { success: false, .. })
                    ) {
                        metrics::counter!("fafb_connections_rejected_total", "reason" => "auth_failed").increment(1);
                    }
                    outcome
                }
            }
            Session::Authenticated(state) => {
                state
                    .process_authenticated(frame, manager, delivery_tx)
                    .await
            }
            Session::Draining(state) => {
                // Reject new subscriptions during drain.
                if matches!(
                    frame,
                    Frame::Subscribe { .. }
                        | Frame::SubscribeAt { .. }
                        | Frame::SubscribeFrom { .. }
                ) {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::ServerDraining,
                        message: "server is draining, new subscriptions rejected".into(),
                    });
                }
                // All other authenticated frames (Publish, Ack, Unsubscribe,
                // Fetch, FetchFrom, etc.) are handled normally so inflight
                // work can complete. Fetch is explicitly allowed during drain
                // because it is stateless — no subscriptions, delivery tasks,
                // or server-side state are created.
                state
                    .process_authenticated(frame, manager, delivery_tx)
                    .await
            }
        }
    }

    /// Handle frames in the pre-auth state. Only Connect is allowed.
    /// Returns `(outcome, Some(state))` on successful auth to transition.
    fn process_pre_auth(
        pre: &PreAuthState,
        frame: Frame,
        auth: &dyn Authenticator,
    ) -> (FrameOutcome, Option<AuthenticatedState>) {
        match frame {
            Frame::Connect {
                version,
                namespace: _,
                api_key,
                capabilities,
            } => {
                if version != 1 {
                    return (
                        FrameOutcome::ReplyAndClose(Frame::Error {
                            code: ErrorCode::UnsupportedVersion,
                            message: format!("unsupported protocol version: {version}"),
                        }),
                        None,
                    );
                }
                let agreed = capabilities & SERVER_CAPABILITIES;
                match auth.authenticate(&api_key) {
                    Ok(session) => {
                        let authed = AuthenticatedState {
                            client_id: pre.client_id,
                            auth: session,
                            capabilities: agreed,
                            max_frame_size: pre.max_frame_size,
                            subscriptions: HashMap::new(),
                            inflight_topics: HashMap::new(),
                            message_id_counter: Arc::clone(&pre.message_id_counter),
                            frame_counter: 0,
                            inflight_warning_fired: false,
                            shared_config: Arc::clone(&pre.shared_config),
                        };
                        (
                            FrameOutcome::Reply(Frame::ConnectAck {
                                success: true,
                                message: String::new(),
                                server_capabilities: agreed,
                            }),
                            Some(authed),
                        )
                    }
                    Err(_) => (
                        FrameOutcome::ReplyAndClose(Frame::ConnectAck {
                            success: false,
                            message: "authentication failed".into(),
                            server_capabilities: 0,
                        }),
                        None,
                    ),
                }
            }
            _ => (
                FrameOutcome::ReplyAndClose(Frame::Error {
                    code: ErrorCode::NotAuthenticated,
                    message: "must send Connect before any other frame".into(),
                }),
                None,
            ),
        }
    }

    /// Maximum subscriptions allowed in a single Reconnect frame.
    const MAX_RECONNECT_SUBSCRIPTIONS: usize = 10_000;

    /// Handle a Reconnect frame: auth + capability negotiation + batch subscribe.
    #[allow(clippy::too_many_arguments)]
    async fn process_reconnect<M: NamespaceManager>(
        &mut self,
        version: u8,
        namespace: crate::types::NamespaceId,
        api_key: crate::types::ApiKey,
        capabilities: Capabilities,
        subscriptions: Vec<crate::codec::frame::SubscriptionEntry>,
        auth: &dyn Authenticator,
        manager: &M,
        delivery_tx: &mpsc::Sender<PublishedMessage>,
    ) -> FrameOutcome {
        // 1. Validate protocol version.
        if version != 1 {
            return FrameOutcome::ReplyAndClose(Frame::Error {
                code: ErrorCode::UnsupportedVersion,
                message: format!("unsupported protocol version: {version}"),
            });
        }

        // 2. Authenticate.
        let auth_session = match auth.authenticate(&api_key) {
            Ok(session) => session,
            Err(_) => {
                metrics::counter!("fafb_connections_rejected_total", "reason" => "auth_failed")
                    .increment(1);
                metrics::counter!("fafb_reconnect_total", "status" => "auth_failed").increment(1);
                return FrameOutcome::ReplyAndClose(Frame::ReconnectAck {
                    status: ReconnectStatus::AuthFailed,
                    message: "authentication failed".into(),
                    server_capabilities: 0,
                    results: vec![],
                });
            }
        };

        // 2b. Validate namespace matches the one resolved from the API key.
        // Catches client bugs where a reconnect is sent to the wrong server
        // or with stale namespace state.
        if namespace != auth_session.namespace {
            metrics::counter!("fafb_reconnect_total", "status" => "auth_failed").increment(1);
            return FrameOutcome::ReplyAndClose(Frame::ReconnectAck {
                status: ReconnectStatus::AuthFailed,
                message: format!(
                    "namespace mismatch: frame says '{}' but key belongs to '{}'",
                    namespace.0, auth_session.namespace.0
                ),
                server_capabilities: 0,
                results: vec![],
            });
        }

        // 3. Negotiate capabilities.
        let agreed = capabilities & SERVER_CAPABILITIES;

        // 4. Check namespace connection limit.
        let pre = match self {
            Session::PreAuth(p) => p,
            _ => unreachable!("process_reconnect called from non-PreAuth state"),
        };
        let client_id = pre.client_id;
        let message_id_counter = Arc::clone(&pre.message_id_counter);
        let max_frame_size = pre.max_frame_size;
        let shared_config = Arc::clone(&pre.shared_config);

        if let Err(e) = manager.connect(&auth_session.namespace, client_id).await {
            tracing::warn!(
                client = %client_id,
                namespace = %auth_session.namespace,
                error = %e,
                "reconnect rejected by namespace manager"
            );
            return FrameOutcome::ReplyAndClose(Frame::ReconnectAck {
                status: ReconnectStatus::AuthFailed,
                message: format!("connection rejected: {e}"),
                server_capabilities: 0,
                results: vec![],
            });
        }

        // 5. Validate subscription count.
        if subscriptions.len() > Self::MAX_RECONNECT_SUBSCRIPTIONS {
            return FrameOutcome::ReplyAndClose(Frame::Error {
                code: ErrorCode::InvalidFrame,
                message: format!(
                    "too many subscriptions: {} (max {})",
                    subscriptions.len(),
                    Self::MAX_RECONNECT_SUBSCRIPTIONS
                ),
            });
        }

        // 6. Build authenticated state.
        let mut authed = AuthenticatedState {
            client_id,
            auth: auth_session,
            capabilities: agreed,
            max_frame_size,
            subscriptions: HashMap::new(),
            inflight_topics: HashMap::new(),
            message_id_counter,
            frame_counter: 0,
            inflight_warning_fired: false,
            shared_config,
        };

        // 7. Batch subscribe — collect per-subscription results.
        //
        // Atomic semantics: all subscriptions are attempted, and if any fail,
        // the successful ones are rolled back (unsubscribed). The client gets
        // Partial status with per-subscription details but NO active
        // subscriptions — it can then fall back to Connect + SubscribeAt for
        // individual control. This prevents orphaned forwarders for
        // subscriptions the client may not want.
        let ns = authed.auth.namespace.clone();
        let has_subscribe_perm = authed.auth.permissions.check(Permission::Subscribe).is_ok();
        let mut results = Vec::with_capacity(subscriptions.len());
        let mut any_failed = false;
        // Pending receivers from successful subscribes, held until we know
        // whether to commit or rollback.
        let mut pending_subs: Vec<(
            TopicName,
            Option<GroupName>,
            mpsc::Receiver<PublishedMessage>,
        )> = Vec::new();

        for entry in &subscriptions {
            if !has_subscribe_perm {
                results.push(SubscriptionResult {
                    topic: entry.topic.clone(),
                    group: entry.group.clone(),
                    status: SubscriptionStatus::PermissionDenied,
                    actual_start_offset: TopicOffset(0),
                });
                any_failed = true;
                continue;
            }

            // Resume from last_acked_offset + 1.
            let start_from =
                SubscribeStart::Offset(TopicOffset(entry.last_acked_offset.0.saturating_add(1)));

            match manager
                .subscribe(
                    &ns,
                    &entry.topic,
                    client_id,
                    entry.group.clone(),
                    start_from,
                )
                .await
            {
                Ok(sub) => {
                    pending_subs.push((entry.topic.clone(), entry.group.clone(), sub.rx));

                    // Detect offset clamping.
                    let requested_next = TopicOffset(entry.last_acked_offset.0.saturating_add(1));
                    let status = if sub.start_offset > requested_next {
                        SubscriptionStatus::OffsetClamped
                    } else {
                        SubscriptionStatus::Restored
                    };

                    results.push(SubscriptionResult {
                        topic: entry.topic.clone(),
                        group: entry.group.clone(),
                        status,
                        actual_start_offset: sub.start_offset,
                    });
                }
                Err(e) => {
                    any_failed = true;
                    let status = match crate::codec::frame::ErrorCode::from_error(&e) {
                        ErrorCode::TopicNotFound => SubscriptionStatus::TopicNotFound,
                        ErrorCode::ResourceExhausted => SubscriptionStatus::GroupFull,
                        ErrorCode::PermissionDenied => SubscriptionStatus::PermissionDenied,
                        _ => SubscriptionStatus::TopicNotFound, // fallback
                    };
                    results.push(SubscriptionResult {
                        topic: entry.topic.clone(),
                        group: entry.group.clone(),
                        status,
                        actual_start_offset: TopicOffset(0),
                    });
                }
            }
        }

        // OffsetClamped is a warning, not a failure — the subscription was
        // restored, just at a different offset. Only hard failures (TopicNotFound,
        // PermissionDenied, GroupFull) trigger rollback.
        let any_hard_failure = results.iter().any(|r| {
            matches!(
                r.status,
                SubscriptionStatus::TopicNotFound
                    | SubscriptionStatus::PermissionDenied
                    | SubscriptionStatus::GroupFull
            )
        });

        if any_hard_failure && !pending_subs.is_empty() {
            // Rollback: unsubscribe all successful subscriptions so the client
            // doesn't end up with a partial set of active forwarders.
            for (topic, group, _rx) in &pending_subs {
                let _ = manager
                    .unsubscribe(&ns, topic, client_id, group.clone())
                    .await;
            }
            // Drop receivers — they'll signal the manager's delivery tasks to stop.
            drop(pending_subs);
        } else {
            // Commit: spawn forwarders for all successful subscriptions.
            for (topic, group, rx) in pending_subs {
                let handle = spawn_forwarder(rx, delivery_tx.clone());
                authed.subscriptions.insert((topic, group), handle);
            }
        }

        // 8. Transition to Authenticated state.
        *self = Session::Authenticated(authed);

        let overall_status = if subscriptions.is_empty() || !any_failed && !any_hard_failure {
            ReconnectStatus::Success
        } else {
            ReconnectStatus::Partial
        };

        let status_label = match overall_status {
            ReconnectStatus::Success => "success",
            ReconnectStatus::Partial => "partial",
            ReconnectStatus::AuthFailed => "auth_failed",
        };
        metrics::counter!("fafb_reconnect_total", "status" => status_label).increment(1);
        metrics::histogram!("fafb_reconnect_subscriptions").record(subscriptions.len() as f64);

        tracing::info!(
            client = %client_id,
            namespace = %ns,
            subscriptions = subscriptions.len(),
            status = ?overall_status,
            "reconnect completed"
        );

        FrameOutcome::Reply(Frame::ReconnectAck {
            status: overall_status,
            message: String::new(),
            server_capabilities: agreed,
            results,
        })
    }

    /// Record a delivered message's topic and build the Deliver frame.
    pub fn build_deliver_frame(&mut self, msg: &PublishedMessage) -> Frame {
        let state = match self {
            Session::Authenticated(s) | Session::Draining(s) => Some(s),
            _ => None,
        };
        if let Some(state) = state {
            if state.inflight_topics.len() < MAX_INFLIGHT {
                state
                    .inflight_topics
                    .insert(msg.envelope.id, msg.envelope.topic.clone());
                // 80% inflight warning (fire-once, resets when inflight drops below)
                let threshold = MAX_INFLIGHT * 4 / 5;
                if state.inflight_topics.len() >= threshold {
                    if !state.inflight_warning_fired {
                        state.inflight_warning_fired = true;
                        tracing::warn!(
                            client = %state.client_id,
                            inflight = state.inflight_topics.len(),
                            max_inflight = MAX_INFLIGHT,
                            "inflight tracking at 80%"
                        );
                    }
                } else {
                    state.inflight_warning_fired = false;
                }
            } else {
                tracing::warn!(
                    client = %state.client_id,
                    "inflight tracking limit reached ({MAX_INFLIGHT}), delivery will not be ack-tracked"
                );
                metrics::counter!("fafb_inflight_limit_reached_total").increment(1);
            }
        }
        Frame::Deliver {
            message_id: msg.envelope.id,
            offset: msg.offset,
            timestamp_millis: msg.envelope.timestamp,
            ttl_ms: msg.envelope.ttl_ms,
            topic: msg.envelope.topic.clone(),
            headers: msg.envelope.headers.clone(),
            payload: Bytes::from(msg.envelope.payload.clone()),
        }
    }

    /// Abort all forwarder tasks. Call when connection closes.
    pub fn abort_forwarders(&mut self) {
        let state = match self {
            Session::Authenticated(s) | Session::Draining(s) => Some(s),
            _ => None,
        };
        if let Some(state) = state {
            for (_, handle) in state.subscriptions.drain() {
                handle.abort();
            }
        }
    }
}

impl AuthenticatedState {
    /// Process frames that require authentication.
    async fn process_authenticated<M: NamespaceManager>(
        &mut self,
        frame: Frame,
        manager: &M,
        delivery_tx: &mpsc::Sender<PublishedMessage>,
    ) -> FrameOutcome {
        // Periodic dead-forwarder reaping.
        self.frame_counter += 1;
        if self.frame_counter.is_multiple_of(REAP_INTERVAL) {
            self.reap_dead_forwarders();
        }

        match frame {
            Frame::Connect { .. } => {
                // Already authenticated — reject duplicate Connect.
                FrameOutcome::ReplyAndClose(Frame::Error {
                    code: ErrorCode::InvalidFrame,
                    message: "already authenticated".into(),
                })
            }

            Frame::Subscribe { topic, group } => {
                self.handle_subscribe(topic, group, SubscribeStart::Head, manager, delivery_tx)
                    .await
            }

            Frame::Unsubscribe { topic, group } => {
                let ns = &self.auth.namespace;

                match manager
                    .unsubscribe(ns, &topic, self.client_id, group.clone())
                    .await
                {
                    Ok(()) => {
                        if let Some(handle) =
                            self.subscriptions.remove(&(topic.clone(), group.clone()))
                        {
                            handle.abort();
                        }
                        tracing::debug!(
                            namespace = %ns, topic = %topic,
                            group = ?group, client = %self.client_id,
                            "client unsubscribed"
                        );
                        FrameOutcome::Reply(Frame::UnsubscribeAck { topic, group })
                    }
                    Err(e) => {
                        let code = ErrorCode::from_error(&e);
                        FrameOutcome::Reply(Frame::Error {
                            code,
                            message: e.to_string(),
                        })
                    }
                }
            }

            Frame::Publish {
                topic,
                headers,
                ttl_ms,
                payload,
            } => {
                if self.auth.permissions.check(Permission::Publish).is_err() {
                    tracing::warn!(
                        namespace = %self.auth.namespace, topic = %topic,
                        client = %self.client_id,
                        "publish permission denied"
                    );
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::PermissionDenied,
                        message: "publish permission denied".into(),
                    });
                }

                let ns = &self.auth.namespace;

                // Resolve TTL: explicit > namespace default > 0 (no expiry).
                // Read from SharedConfig per-publish so TTL limits are hot-reloadable.
                let resolved_ttl_ms = {
                    let cfg = self.shared_config.load();
                    let quota = cfg.namespace_quotas.get(ns);
                    let default_ttl_ms =
                        quota.map_or(0u64, |q| q.default_ttl_secs.saturating_mul(1000));
                    let max_ttl_ms = quota.map_or(0u64, |q| q.max_ttl_secs.saturating_mul(1000));
                    let min_ttl_ms = quota.map_or(0u64, |q| q.min_ttl_secs.saturating_mul(1000));

                    let effective = if ttl_ms > 0 {
                        ttl_ms as u64
                    } else if default_ttl_ms > 0 {
                        default_ttl_ms
                    } else {
                        0
                    };

                    if effective > 0 {
                        // Check min TTL
                        if min_ttl_ms > 0 && effective < min_ttl_ms {
                            return FrameOutcome::Reply(Frame::Error {
                                code: ErrorCode::InvalidFrame,
                                message: format!(
                                    "TTL {}ms is below minimum {}ms",
                                    effective, min_ttl_ms,
                                ),
                            });
                        }
                        // Clamp to max TTL
                        if max_ttl_ms > 0 && effective > max_ttl_ms {
                            tracing::warn!(
                                namespace = %ns, topic = %topic,
                                ttl_ms = effective, max_ttl_ms = max_ttl_ms,
                                "TTL clamped to namespace maximum"
                            );
                            max_ttl_ms as u32
                        } else {
                            effective as u32
                        }
                    } else {
                        0
                    }
                };

                let message_id =
                    MessageId::new(self.message_id_counter.fetch_add(1, Ordering::Relaxed));
                let payload_len = payload.len();
                let header_count = headers.len();
                let header_bytes: usize = headers.iter().map(|h| h.key.len() + h.value.len()).sum();
                let envelope = MessageEnvelope::new(
                    message_id,
                    ns.clone(),
                    topic.clone(),
                    payload.to_vec(),
                    now_millis(),
                )
                .with_headers(headers)
                .with_ttl(resolved_ttl_ms);

                // Record TTL metric (always, so 0-TTL messages appear in distribution)
                metrics::histogram!("fafb_message_ttl_seconds")
                    .record(resolved_ttl_ms as f64 / 1000.0);

                // Record header metrics (unconditional so 0-header messages appear in the distribution)
                metrics::histogram!("fafb_message_header_count").record(header_count as f64);
                metrics::histogram!("fafb_message_header_bytes").record(header_bytes as f64);

                match manager.publish(ns, &topic, envelope).await {
                    Ok(offset) => {
                        tracing::trace!(
                            namespace = %ns, topic = %topic,
                            message_id = %message_id, offset = %offset,
                            payload_bytes = payload_len,
                            "message published"
                        );
                        FrameOutcome::Reply(Frame::PublishAck { message_id })
                    }
                    Err(e) => {
                        tracing::warn!(
                            namespace = %ns, topic = %topic,
                            client = %self.client_id, error = %e,
                            "publish failed"
                        );
                        let code = ErrorCode::from_error(&e);
                        FrameOutcome::Reply(Frame::Error {
                            code,
                            message: e.to_string(),
                        })
                    }
                }
            }

            Frame::Ack { message_id } => {
                let Some(topic) = self.inflight_topics.remove(&message_id) else {
                    metrics::counter!("fafb_duplicate_ack_total",
                        "namespace" => self.auth.namespace.0.clone()
                    )
                    .increment(1);
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InternalError,
                        message: format!("unknown message id: {message_id}"),
                    });
                };
                let ns = &self.auth.namespace;

                if let Err(e) = manager.ack(ns, &topic, self.client_id, message_id).await {
                    tracing::warn!(
                        client = %self.client_id, message_id = %message_id,
                        error = %e, "ack failed"
                    );
                } else {
                    tracing::trace!(
                        namespace = %ns, topic = %topic,
                        message_id = %message_id,
                        "message acked"
                    );
                }
                FrameOutcome::None
            }

            Frame::BatchAck { message_ids } => {
                if self.capabilities & CAP_BATCH_ACK == 0 {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InvalidFrame,
                        message: "BatchAck requires CAP_BATCH_ACK capability".into(),
                    });
                }
                let ns = &self.auth.namespace;
                let batch_size = message_ids.len() as f64;
                metrics::histogram!("fafb_batch_ack_size").record(batch_size);

                // Resolve topic for each message_id from session inflight map,
                // then delegate to manager.batch_ack() for grouped processing.
                let mut acks = Vec::with_capacity(message_ids.len());
                let mut session_unknown: u32 = 0;
                for message_id in message_ids {
                    if let Some(topic) = self.inflight_topics.remove(&message_id) {
                        acks.push((message_id, topic));
                    } else {
                        session_unknown += 1;
                    }
                }
                let result = manager.batch_ack(ns, self.client_id, acks).await;
                let acked_count = result.acked;
                let unknown_count = session_unknown + result.unknown;
                tracing::trace!(
                    namespace = %ns, client = %self.client_id,
                    acked = acked_count, unknown = unknown_count,
                    "batch ack processed"
                );
                FrameOutcome::Reply(Frame::BatchAckResponse {
                    acked_count,
                    unknown_count,
                })
            }

            Frame::Error { code, message } => {
                tracing::warn!(
                    client = %self.client_id,
                    code = ?code,
                    message = %message,
                    "received error frame from client"
                );
                FrameOutcome::None
            }

            Frame::SubscribeAt {
                topic,
                offset,
                group,
            } => {
                // SubscribeAt with explicit offset is not supported for wildcards.
                if crate::wildcard::contains_wildcard(&topic.0) {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InvalidOperation,
                        message: "SubscribeAt with explicit offset is not supported for \
                                  wildcard patterns; use Subscribe (head) or SubscribeFrom \
                                  (timestamp)"
                            .into(),
                    });
                }
                self.handle_subscribe(
                    topic,
                    group,
                    SubscribeStart::Offset(offset),
                    manager,
                    delivery_tx,
                )
                .await
            }

            Frame::SubscribeFrom {
                topic,
                timestamp_millis,
                group,
            } => {
                if self.auth.permissions.check(Permission::Subscribe).is_err() {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::PermissionDenied,
                        message: "subscribe permission denied".into(),
                    });
                }

                // For wildcard patterns, pass Timestamp to the manager for
                // per-topic resolution instead of resolving here.
                if crate::wildcard::contains_wildcard(&topic.0) {
                    return self
                        .handle_subscribe(
                            topic,
                            group,
                            SubscribeStart::Timestamp(timestamp_millis),
                            manager,
                            delivery_tx,
                        )
                        .await;
                }

                let ns = &self.auth.namespace;
                let start = match manager
                    .offset_at_timestamp(ns, &topic, timestamp_millis)
                    .await
                {
                    Ok(Some(offset)) => SubscribeStart::Offset(offset),
                    Ok(None) => SubscribeStart::Head,
                    Err(e) => {
                        let code = ErrorCode::from_error(&e);
                        return FrameOutcome::Reply(Frame::Error {
                            code,
                            message: e.to_string(),
                        });
                    }
                };
                self.handle_subscribe(topic, group, start, manager, delivery_tx)
                    .await
            }

            // Already handled in process_frame before dispatch.
            Frame::Ping | Frame::Pong => unreachable!(),
            Frame::ConnectAck { .. }
            | Frame::SubscribeAck { .. }
            | Frame::UnsubscribeAck { .. }
            | Frame::PublishAck { .. }
            | Frame::Deliver { .. }
            | Frame::GoingAway { .. }
            | Frame::ReconnectAck { .. }
            | Frame::FetchResponse { .. }
            | Frame::NackAck { .. }
            | Frame::BatchAckResponse { .. }
            | Frame::NotLeader { .. }
            | Frame::ClusterUnavailable { .. } => unreachable!(),
            Frame::Fetch {
                topic,
                offset,
                max_count,
            } => {
                if self.auth.permissions.check(Permission::Fetch).is_err() {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::PermissionDenied,
                        message: "fetch permission denied".into(),
                    });
                }
                if crate::wildcard::contains_wildcard(&topic.0) {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InvalidTopic,
                        message: "wildcard patterns not supported for Fetch".into(),
                    });
                }
                let frame_overhead = 4 + 1 + 2 + topic.0.len() + 8 + 8 + 4;
                let max_response_bytes =
                    (self.max_frame_size as usize).saturating_sub(frame_overhead);
                self.handle_fetch_at_offset(topic, offset, max_count, max_response_bytes, manager)
                    .await
            }

            Frame::FetchFrom {
                topic,
                timestamp_millis,
                max_count,
            } => {
                if self.auth.permissions.check(Permission::Fetch).is_err() {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::PermissionDenied,
                        message: "fetch permission denied".into(),
                    });
                }
                if crate::wildcard::contains_wildcard(&topic.0) {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InvalidTopic,
                        message: "wildcard patterns not supported for Fetch".into(),
                    });
                }
                let frame_overhead = 4 + 1 + 2 + topic.0.len() + 8 + 8 + 4;
                let max_response_bytes =
                    (self.max_frame_size as usize).saturating_sub(frame_overhead);
                let ns = &self.auth.namespace;
                let resolved_offset = match manager
                    .offset_at_timestamp(ns, &topic, timestamp_millis)
                    .await
                {
                    Ok(Some(offset)) => offset,
                    Ok(None) => {
                        let head = manager
                            .head_offset(ns, &topic)
                            .await
                            .unwrap_or(TopicOffset::ZERO);
                        return FrameOutcome::Reply(Frame::FetchResponse {
                            topic,
                            next_offset: head,
                            head_offset: head,
                            messages: vec![],
                        });
                    }
                    Err(e) => {
                        let code = ErrorCode::from_error(&e);
                        return FrameOutcome::Reply(Frame::Error {
                            code,
                            message: e.to_string(),
                        });
                    }
                };
                self.handle_fetch_at_offset(
                    topic,
                    resolved_offset,
                    max_count,
                    max_response_bytes,
                    manager,
                )
                .await
            }
            Frame::Nack { message_id, reason } => {
                if self.capabilities & CAP_NACK == 0 {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InvalidFrame,
                        message: "Nack requires CAP_NACK capability".into(),
                    });
                }
                let Some(topic) = self.inflight_topics.get(&message_id).cloned() else {
                    return FrameOutcome::Reply(Frame::Error {
                        code: ErrorCode::InternalError,
                        message: format!("unknown message id: {message_id}"),
                    });
                };
                let ns = &self.auth.namespace;

                let status = match manager
                    .nack(ns, &topic, self.client_id, message_id, reason)
                    .await
                {
                    Ok(NackResult::DeadLettered) => {
                        self.inflight_topics.remove(&message_id);
                        tracing::debug!(
                            namespace = %ns, topic = %topic,
                            message_id = %message_id,
                            "message nacked → dead-lettered"
                        );
                        0x00
                    }
                    Ok(NackResult::Dropped) => {
                        self.inflight_topics.remove(&message_id);
                        tracing::debug!(
                            namespace = %ns, topic = %topic,
                            message_id = %message_id,
                            "message nacked → dropped (DLQ disabled)"
                        );
                        0x01
                    }
                    Err(e) => {
                        let err_msg = e.to_string();
                        // "not in inflight" means the server-side entry was
                        // already cleaned up (concurrent ack/sweep). Remove
                        // from session tracking and report as dropped rather
                        // than misleading the client with "DLQ write failed".
                        if err_msg.contains("not in inflight") {
                            self.inflight_topics.remove(&message_id);
                            tracing::debug!(
                                client = %self.client_id, message_id = %message_id,
                                "nack: message already resolved server-side"
                            );
                            0x01
                        } else {
                            tracing::warn!(
                                client = %self.client_id, message_id = %message_id,
                                error = %e, "nack DLQ write failed, message stays inflight"
                            );
                            0x02
                        }
                    }
                };
                FrameOutcome::Reply(Frame::NackAck { message_id, status })
            }

            // Reconnect is only valid in PreAuth state.
            Frame::Reconnect { .. } => FrameOutcome::ReplyAndClose(Frame::Error {
                code: ErrorCode::InvalidFrame,
                message: "already authenticated, Reconnect only valid before auth".into(),
            }),
        }
    }

    /// Shared subscribe logic for Subscribe, SubscribeAt, and SubscribeFrom.
    async fn handle_subscribe<M: NamespaceManager>(
        &mut self,
        topic: TopicName,
        group: Option<GroupName>,
        start_from: SubscribeStart,
        manager: &M,
        delivery_tx: &mpsc::Sender<PublishedMessage>,
    ) -> FrameOutcome {
        if self.auth.permissions.check(Permission::Subscribe).is_err() {
            tracing::warn!(
                namespace = %self.auth.namespace, topic = %topic,
                client = %self.client_id,
                "subscribe permission denied"
            );
            return FrameOutcome::Reply(Frame::Error {
                code: ErrorCode::PermissionDenied,
                message: "subscribe permission denied".into(),
            });
        }

        let is_wildcard = crate::wildcard::contains_wildcard(&topic.0);

        // Gate wildcard subscriptions on CAP_WILDCARD capability.
        if is_wildcard && (self.capabilities & CAP_WILDCARD == 0) {
            return FrameOutcome::Reply(Frame::Error {
                code: ErrorCode::InvalidTopic,
                message: "wildcard subscriptions require CAP_WILDCARD capability".into(),
            });
        }

        // Remember the requested offset for clamping detection
        let requested_offset = match start_from {
            SubscribeStart::Offset(o) => Some(o),
            SubscribeStart::Head | SubscribeStart::Timestamp(_) => None,
        };

        let ns = &self.auth.namespace;
        match manager
            .subscribe(ns, &topic, self.client_id, group.clone(), start_from)
            .await
        {
            Ok(result) => {
                let handle = spawn_forwarder(result.rx, delivery_tx.clone());
                self.subscriptions
                    .insert((topic.clone(), group.clone()), handle);

                let start_offset = result.start_offset;
                let matched_count = result.matched_count;

                tracing::debug!(
                    namespace = %ns, topic = %topic,
                    group = ?group, client = %self.client_id,
                    start_offset = %start_offset,
                    "client subscribed"
                );

                let ack = Frame::SubscribeAck {
                    topic,
                    group,
                    start_offset,
                    matched_count,
                };

                // If the actual start offset differs from the requested offset,
                // send an informational OffsetNotAvailable warning before the ack.
                if let Some(requested) = requested_offset
                    && start_offset > requested
                {
                    return FrameOutcome::ReplyMany(vec![
                        Frame::Error {
                            code: ErrorCode::OffsetNotAvailable,
                            message: format!(
                                "requested offset {requested} not available, starting from {start_offset}"
                            ),
                        },
                        ack,
                    ]);
                }

                FrameOutcome::Reply(ack)
            }
            Err(e) => {
                let code = ErrorCode::from_error(&e);
                FrameOutcome::Reply(Frame::Error {
                    code,
                    message: e.to_string(),
                })
            }
        }
    }

    /// Shared fetch logic for Fetch and FetchFrom (after offset resolution).
    async fn handle_fetch_at_offset<M: NamespaceManager>(
        &self,
        topic: TopicName,
        offset: TopicOffset,
        max_count: u32,
        max_response_bytes: usize,
        manager: &M,
    ) -> FrameOutcome {
        let ns = &self.auth.namespace;
        let headers_enabled = self.capabilities & CAP_HEADERS != 0;
        match manager
            .fetch(
                ns,
                &topic,
                offset,
                max_count,
                max_response_bytes,
                headers_enabled,
            )
            .await
        {
            Ok(result) => {
                let response = Frame::FetchResponse {
                    topic: topic.clone(),
                    next_offset: result.next_offset,
                    head_offset: result.head_offset,
                    messages: result
                        .messages
                        .into_iter()
                        .map(|m| FetchedMessage {
                            offset: m.offset,
                            timestamp_millis: m.envelope.timestamp,
                            ttl_ms: m.envelope.ttl_ms,
                            headers: m.envelope.headers.clone(),
                            payload: Bytes::from(m.envelope.payload.clone()),
                        })
                        .collect(),
                };
                if let Some(original) = result.clamped_from {
                    FrameOutcome::ReplyMany(vec![
                        Frame::Error {
                            code: ErrorCode::OffsetNotAvailable,
                            message: format!(
                                "requested offset {} not available, clamped",
                                original
                            ),
                        },
                        response,
                    ])
                } else {
                    FrameOutcome::Reply(response)
                }
            }
            Err(e) if e.is_message_too_large() => FrameOutcome::Reply(Frame::Error {
                code: ErrorCode::MessageTooLarge,
                message: e.to_string(),
            }),
            Err(e) => {
                let code = ErrorCode::from_error(&e);
                FrameOutcome::Reply(Frame::Error {
                    code,
                    message: e.to_string(),
                })
            }
        }
    }

    /// Remove subscription entries whose forwarder task has exited.
    fn reap_dead_forwarders(&mut self) {
        self.subscriptions.retain(|key, handle| {
            if handle.is_finished() {
                tracing::debug!(
                    client = %self.client_id,
                    topic = %key.0,
                    group = ?key.1,
                    "reaped dead forwarder task"
                );
                false
            } else {
                true
            }
        });
    }
}

/// Bridges the per-subscription receiver from namespace manager into the
/// single per-connection aggregation channel.
fn spawn_forwarder(
    mut rx: mpsc::Receiver<PublishedMessage>,
    delivery_tx: mpsc::Sender<PublishedMessage>,
) -> JoinHandle<()> {
    tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if delivery_tx.send(msg).await.is_err() {
                break; // Connection gone.
            }
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::auth::{MockAuthenticator, PermissionSet};
    use crate::namespace::NamespaceManager;
    use crate::types::{
        ApiKey, ClientId, GroupName, MessageEnvelope, MessageId, NamespaceId, Permission,
        PublishedMessage, TopicName, TopicOffset,
    };
    use std::sync::Arc;
    use std::sync::atomic::AtomicU64;
    use tokio::sync::mpsc;

    // -- Stub namespace manager for unit testing --

    struct StubNamespaceManager {
        connect_tx: mpsc::Sender<(NamespaceId, ClientId)>,
        connect_result: std::sync::Mutex<Option<crate::error::Error>>,
        subscribe_tx: mpsc::Sender<(NamespaceId, TopicName, ClientId, Option<GroupName>)>,
        unsubscribe_tx: mpsc::Sender<(NamespaceId, TopicName, ClientId, Option<GroupName>)>,
        publish_tx: mpsc::Sender<(NamespaceId, TopicName, MessageEnvelope)>,
        ack_tx: mpsc::Sender<(NamespaceId, TopicName, ClientId, MessageId)>,
        disconnect_tx: mpsc::Sender<ClientId>,
    }

    impl StubNamespaceManager {
        fn new() -> (
            Self,
            mpsc::Receiver<(NamespaceId, ClientId)>,
            mpsc::Receiver<(NamespaceId, TopicName, ClientId, Option<GroupName>)>,
            mpsc::Receiver<(NamespaceId, TopicName, ClientId, Option<GroupName>)>,
            mpsc::Receiver<(NamespaceId, TopicName, MessageEnvelope)>,
            mpsc::Receiver<(NamespaceId, TopicName, ClientId, MessageId)>,
            mpsc::Receiver<ClientId>,
        ) {
            let (conn_tx, conn_rx) = mpsc::channel(16);
            let (sub_tx, sub_rx) = mpsc::channel(16);
            let (unsub_tx, unsub_rx) = mpsc::channel(16);
            let (pub_tx, pub_rx) = mpsc::channel(16);
            let (ack_tx, ack_rx) = mpsc::channel(16);
            let (disc_tx, disc_rx) = mpsc::channel(16);
            (
                Self {
                    connect_tx: conn_tx,
                    connect_result: std::sync::Mutex::new(None),
                    subscribe_tx: sub_tx,
                    unsubscribe_tx: unsub_tx,
                    publish_tx: pub_tx,
                    ack_tx: ack_tx,
                    disconnect_tx: disc_tx,
                },
                conn_rx,
                sub_rx,
                unsub_rx,
                pub_rx,
                ack_rx,
                disc_rx,
            )
        }
    }

    impl NamespaceManager for StubNamespaceManager {
        async fn connect(&self, ns: &NamespaceId, client: ClientId) -> crate::error::Result<()> {
            let _ = self.connect_tx.send((ns.clone(), client)).await;
            if let Some(err) = self.connect_result.lock().unwrap().take() {
                return Err(err);
            }
            Ok(())
        }

        async fn subscribe(
            &self,
            ns: &NamespaceId,
            topic: &TopicName,
            client: ClientId,
            group: Option<GroupName>,
            _start_from: SubscribeStart,
        ) -> crate::error::Result<crate::types::SubscribeResult> {
            let _ = self
                .subscribe_tx
                .send((ns.clone(), topic.clone(), client, group))
                .await;
            let (_tx, rx) = mpsc::channel(16);
            Ok(crate::types::SubscribeResult {
                rx,
                start_offset: TopicOffset(0),
                matched_count: 0,
            })
        }

        async fn unsubscribe(
            &self,
            ns: &NamespaceId,
            topic: &TopicName,
            client: ClientId,
            group: Option<GroupName>,
        ) -> crate::error::Result<()> {
            let _ = self
                .unsubscribe_tx
                .send((ns.clone(), topic.clone(), client, group))
                .await;
            Ok(())
        }

        async fn publish(
            &self,
            ns: &NamespaceId,
            topic: &TopicName,
            msg: MessageEnvelope,
        ) -> crate::error::Result<TopicOffset> {
            let _ = self.publish_tx.send((ns.clone(), topic.clone(), msg)).await;
            Ok(TopicOffset(0))
        }

        async fn ack(
            &self,
            ns: &NamespaceId,
            topic: &TopicName,
            client: ClientId,
            message_id: MessageId,
        ) -> crate::error::Result<()> {
            let _ = self
                .ack_tx
                .send((ns.clone(), topic.clone(), client, message_id))
                .await;
            Ok(())
        }

        async fn nack(
            &self,
            ns: &NamespaceId,
            topic: &TopicName,
            client: ClientId,
            message_id: MessageId,
            _reason: String,
        ) -> crate::error::Result<crate::namespace::NackResult> {
            let _ = self
                .ack_tx
                .send((ns.clone(), topic.clone(), client, message_id))
                .await;
            Ok(crate::namespace::NackResult::DeadLettered)
        }

        async fn disconnect(&self, client: &ClientId) -> crate::error::Result<()> {
            let _ = self.disconnect_tx.send(*client).await;
            Ok(())
        }

        async fn list_topics(&self, _ns: &NamespaceId) -> crate::error::Result<Vec<TopicName>> {
            Ok(vec![])
        }

        async fn topic_info(
            &self,
            _ns: &NamespaceId,
            _topic: &TopicName,
        ) -> crate::error::Result<crate::types::TopicInfo> {
            Err(crate::error::Error::TopicNotFound(TopicName("stub".into())))
        }

        async fn head_offset(
            &self,
            _ns: &NamespaceId,
            _topic: &TopicName,
        ) -> crate::error::Result<TopicOffset> {
            Ok(TopicOffset::ZERO)
        }

        async fn min_live_offset(
            &self,
            _ns: &NamespaceId,
            _topic: &TopicName,
        ) -> Option<TopicOffset> {
            None
        }

        async fn offset_at_timestamp(
            &self,
            _ns: &NamespaceId,
            _topic: &TopicName,
            _timestamp_millis: u64,
        ) -> crate::error::Result<Option<TopicOffset>> {
            Ok(None)
        }

        async fn fetch(
            &self,
            _ns: &NamespaceId,
            _topic: &TopicName,
            _offset: TopicOffset,
            _max_count: u32,
            _max_response_bytes: usize,
            _headers_enabled: bool,
        ) -> crate::error::Result<crate::types::FetchResult> {
            Ok(crate::types::FetchResult {
                messages: Vec::new(),
                next_offset: TopicOffset::ZERO,
                head_offset: TopicOffset::ZERO,
                clamped_from: None,
            })
        }

        async fn shutdown(&self) -> crate::error::Result<()> {
            Ok(())
        }

        async fn topic_last_subscribe(
            &self,
            _ns: &crate::types::NamespaceId,
            _topic: &crate::types::TopicName,
        ) -> Option<u64> {
            None
        }

        async fn has_active_subscribers(
            &self,
            _ns: &crate::types::NamespaceId,
            _topic: &crate::types::TopicName,
        ) -> bool {
            false
        }

        async fn has_inflight_messages(
            &self,
            _ns: &crate::types::NamespaceId,
            _topic: &crate::types::TopicName,
        ) -> bool {
            false
        }

        async fn teardown_topic(
            &self,
            _ns: &crate::types::NamespaceId,
            _topic: &crate::types::TopicName,
        ) -> crate::error::Result<()> {
            Ok(())
        }
    }

    fn make_auth(perms: &[Permission]) -> MockAuthenticator {
        MockAuthenticator {
            namespace: NamespaceId("test-ns".into()),
            key: ApiKey("test-key".into()),
            permissions: PermissionSet::from_slice(perms),
        }
    }

    fn make_all_perms_auth() -> MockAuthenticator {
        make_auth(&[
            Permission::Subscribe,
            Permission::Publish,
            Permission::Query,
        ])
    }

    fn test_shared_config() -> SharedConfig {
        crate::config::shared_config(&crate::config::ServerConfig::default())
    }

    fn make_session() -> (Session, Arc<AtomicU64>) {
        let counter = Arc::new(AtomicU64::new(1));
        let session = Session::new(
            ClientId(1),
            Arc::clone(&counter),
            1_048_576,
            test_shared_config(),
        );
        (session, counter)
    }

    fn connect_frame() -> Frame {
        connect_frame_with_caps(0)
    }

    fn connect_frame_with_caps(caps: Capabilities) -> Frame {
        Frame::Connect {
            version: 1,
            namespace: NamespaceId("test-ns".into()),
            api_key: ApiKey("test-key".into()),
            capabilities: caps,
        }
    }

    async fn authenticate_session(
        session: &mut Session,
        auth: &dyn Authenticator,
        manager: &StubNamespaceManager,
        delivery_tx: &mpsc::Sender<PublishedMessage>,
    ) {
        authenticate_session_with_caps(session, auth, manager, delivery_tx, 0).await;
    }

    async fn authenticate_session_with_caps(
        session: &mut Session,
        auth: &dyn Authenticator,
        manager: &StubNamespaceManager,
        delivery_tx: &mpsc::Sender<PublishedMessage>,
        caps: Capabilities,
    ) {
        let outcome = session
            .process_frame(connect_frame_with_caps(caps), auth, manager, delivery_tx)
            .await;
        assert!(matches!(
            outcome,
            FrameOutcome::Reply(Frame::ConnectAck { success: true, .. })
        ));
        assert!(matches!(session, Session::Authenticated(_)));
    }

    // -- Tests --

    #[tokio::test]
    async fn server_only_frame_rejected() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let server_frames = vec![
            Frame::ConnectAck {
                success: true,
                message: String::new(),
                server_capabilities: 0,
            },
            Frame::SubscribeAck {
                topic: TopicName("t".into()),
                group: None,
                start_offset: TopicOffset(0),
                matched_count: 0,
            },
            Frame::UnsubscribeAck {
                topic: TopicName("t".into()),
                group: None,
            },
            Frame::PublishAck {
                message_id: MessageId(1),
            },
            Frame::Deliver {
                message_id: MessageId(1),
                offset: TopicOffset(0),
                timestamp_millis: 0,
                ttl_ms: 0,
                topic: TopicName("t".into()),
                headers: vec![],
                payload: Bytes::new(),
            },
            Frame::GoingAway {
                reason: crate::codec::frame::GoingAwayReason::UPGRADE,
                deadline_ms: 30000,
                message: "test".into(),
            },
            Frame::ReconnectAck {
                status: crate::codec::frame::ReconnectStatus::Success,
                message: String::new(),
                server_capabilities: 0,
                results: vec![],
            },
        ];

        for frame in server_frames {
            let outcome = session
                .process_frame(frame, &auth, &manager, &delivery_tx)
                .await;
            match outcome {
                FrameOutcome::ReplyAndClose(Frame::Error { code, .. }) => {
                    assert_eq!(code, ErrorCode::InvalidFrame);
                }
                _ => panic!("expected ReplyAndClose with InvalidFrame"),
            }
        }
    }

    #[tokio::test]
    async fn pre_auth_gate() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Subscribe {
                    topic: TopicName("t".into()),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::ReplyAndClose(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::NotAuthenticated);
            }
            _ => panic!("expected ReplyAndClose with NotAuthenticated"),
        }
    }

    #[tokio::test]
    async fn ping_before_auth() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(Frame::Ping, &auth, &manager, &delivery_tx)
            .await;
        assert!(matches!(outcome, FrameOutcome::Reply(Frame::Pong)));
    }

    #[tokio::test]
    async fn connect_success() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        assert!(matches!(session, Session::PreAuth(_)));

        let outcome = session
            .process_frame(connect_frame(), &auth, &manager, &delivery_tx)
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::ConnectAck { success, .. }) => {
                assert!(success);
            }
            _ => panic!("expected Reply with ConnectAck success"),
        }

        // Session should now be in Authenticated state.
        assert!(matches!(session, Session::Authenticated(_)));
    }

    #[tokio::test]
    async fn connect_bad_key() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Connect {
                    version: 1,
                    namespace: NamespaceId("test-ns".into()),
                    api_key: ApiKey("wrong-key".into()),
                    capabilities: 0,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::ReplyAndClose(Frame::ConnectAck { success, .. }) => {
                assert!(!success);
            }
            _ => panic!("expected ReplyAndClose with ConnectAck failure"),
        }

        // Should remain in PreAuth state.
        assert!(matches!(session, Session::PreAuth(_)));
    }

    #[tokio::test]
    async fn connect_wrong_version() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Connect {
                    version: 2,
                    namespace: NamespaceId("test-ns".into()),
                    api_key: ApiKey("test-key".into()),
                    capabilities: 0,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::ReplyAndClose(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::UnsupportedVersion);
            }
            _ => panic!("expected ReplyAndClose with UnsupportedVersion"),
        }
    }

    #[tokio::test]
    async fn connect_already_authenticated() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(connect_frame(), &auth, &manager, &delivery_tx)
            .await;
        match outcome {
            FrameOutcome::ReplyAndClose(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::InvalidFrame);
            }
            _ => panic!("expected ReplyAndClose with InvalidFrame"),
        }
    }

    #[tokio::test]
    async fn subscribe_permission_denied() {
        let (mut session, _) = make_session();
        let auth = make_auth(&[Permission::Publish]);
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::Subscribe {
                    topic: TopicName("t".into()),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::PermissionDenied);
            }
            _ => panic!("expected Reply with PermissionDenied"),
        }
    }

    #[tokio::test]
    async fn publish_assigns_message_id() {
        let (mut session, counter) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::Publish {
                    topic: TopicName("t".into()),
                    headers: vec![],
                    ttl_ms: 0,
                    payload: Bytes::from_static(b"hello"),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match &outcome {
            FrameOutcome::Reply(Frame::PublishAck { message_id }) => {
                assert_eq!(*message_id, MessageId(1));
            }
            _ => panic!("expected Reply with PublishAck"),
        }

        let outcome = session
            .process_frame(
                Frame::Publish {
                    topic: TopicName("t".into()),
                    headers: vec![],
                    ttl_ms: 0,
                    payload: Bytes::from_static(b"world"),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match &outcome {
            FrameOutcome::Reply(Frame::PublishAck { message_id }) => {
                assert_eq!(*message_id, MessageId(2));
            }
            _ => panic!("expected Reply with PublishAck"),
        }

        assert_eq!(counter.load(Ordering::Relaxed), 3);
    }

    #[tokio::test]
    async fn ack_unknown_message_id() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::Ack {
                    message_id: MessageId(999),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::InternalError);
            }
            _ => panic!("expected Reply with InternalError"),
        }
    }

    #[tokio::test]
    async fn build_deliver_frame_populates_inflight() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);
        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let msg = PublishedMessage {
            envelope: Arc::new(MessageEnvelope::new(
                MessageId(42),
                NamespaceId("ns".into()),
                TopicName("events".into()),
                vec![1, 2, 3],
                now_millis(),
            )),
            offset: TopicOffset(10),
        };

        let frame = session.build_deliver_frame(&msg);
        assert!(matches!(
            &frame,
            Frame::Deliver { message_id, offset, topic, payload, .. }
            if *message_id == MessageId(42)
                && *offset == TopicOffset(10)
                && *topic == TopicName("events".into())
                && payload.as_ref() == &[1, 2, 3]
        ));

        match &session {
            Session::Authenticated(state) => {
                assert_eq!(
                    state.inflight_topics.get(&MessageId(42)),
                    Some(&TopicName("events".into()))
                );
            }
            _ => panic!("expected Authenticated state"),
        }
    }

    #[tokio::test]
    async fn inflight_cap_prevents_unbounded_growth() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);
        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        for i in 0..MAX_INFLIGHT {
            let msg = PublishedMessage {
                envelope: Arc::new(MessageEnvelope::new(
                    MessageId(i as u64 + 1),
                    NamespaceId("ns".into()),
                    TopicName("t".into()),
                    vec![],
                    now_millis(),
                )),
                offset: TopicOffset(i as u64),
            };
            session.build_deliver_frame(&msg);
        }

        let extra = PublishedMessage {
            envelope: Arc::new(MessageEnvelope::new(
                MessageId(MAX_INFLIGHT as u64 + 1),
                NamespaceId("ns".into()),
                TopicName("t".into()),
                vec![],
                now_millis(),
            )),
            offset: TopicOffset(MAX_INFLIGHT as u64),
        };
        session.build_deliver_frame(&extra);

        match &session {
            Session::Authenticated(state) => {
                assert_eq!(state.inflight_topics.len(), MAX_INFLIGHT);
                assert!(
                    !state
                        .inflight_topics
                        .contains_key(&MessageId(MAX_INFLIGHT as u64 + 1))
                );
            }
            _ => panic!("expected Authenticated state"),
        }
    }

    #[tokio::test]
    async fn connect_called_on_successful_auth() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, mut conn_rx, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let (ns, client) = conn_rx.try_recv().unwrap();
        assert_eq!(ns, NamespaceId("test-ns".into()));
        assert_eq!(client, ClientId(1));
    }

    #[tokio::test]
    async fn connect_rejected_returns_failure_ack() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, _, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        // Pre-set connect to fail
        *manager.connect_result.lock().unwrap() = Some(crate::error::Error::ResourceExhausted(
            "namespace connection limit reached".to_owned(),
        ));

        let outcome = session
            .process_frame(connect_frame(), &auth, &manager, &delivery_tx)
            .await;
        match outcome {
            FrameOutcome::ReplyAndClose(Frame::ConnectAck {
                success, message, ..
            }) => {
                assert!(!success);
                assert!(message.contains("connection limit"));
            }
            _ => panic!("expected ReplyAndClose with ConnectAck failure"),
        }

        // Session should remain PreAuth
        assert!(matches!(session, Session::PreAuth(_)));
    }

    #[tokio::test]
    async fn subscribe_at_returns_ack_with_offset() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::SubscribeAt {
                    topic: TopicName("t".into()),
                    offset: TopicOffset(42),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::SubscribeAck {
                start_offset,
                topic,
                ..
            }) => {
                assert_eq!(topic, TopicName("t".into()));
                // StubNamespaceManager returns TopicOffset(0)
                assert_eq!(start_offset, TopicOffset(0));
            }
            _ => panic!("expected Reply with SubscribeAck, got: {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn subscribe_from_returns_ack() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::SubscribeFrom {
                    topic: TopicName("t".into()),
                    timestamp_millis: 1_700_000_000_000,
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        // StubNamespaceManager.offset_at_timestamp returns None, so falls back to Head
        match outcome {
            FrameOutcome::Reply(Frame::SubscribeAck { start_offset, .. }) => {
                assert_eq!(start_offset, TopicOffset(0));
            }
            _ => panic!("expected Reply with SubscribeAck, got: {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn subscribe_at_permission_denied() {
        let (mut session, _) = make_session();
        let auth = make_auth(&[Permission::Publish]);
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::SubscribeAt {
                    topic: TopicName("t".into()),
                    offset: TopicOffset(0),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::PermissionDenied);
            }
            _ => panic!("expected PermissionDenied error"),
        }
    }

    #[tokio::test]
    async fn subscribe_from_permission_denied() {
        let (mut session, _) = make_session();
        let auth = make_auth(&[Permission::Publish]);
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::SubscribeFrom {
                    topic: TopicName("t".into()),
                    timestamp_millis: 1_700_000_000_000,
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::PermissionDenied);
            }
            _ => panic!("expected PermissionDenied error"),
        }
    }

    // -- Draining session state tests --

    #[tokio::test]
    async fn draining_rejects_subscribe() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;
        session.transition_to_draining();
        assert!(matches!(session, Session::Draining(_)));

        let outcome = session
            .process_frame(
                Frame::Subscribe {
                    topic: TopicName("t".into()),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::ServerDraining);
            }
            _ => panic!("expected ServerDraining error, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn draining_rejects_subscribe_at() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;
        session.transition_to_draining();

        let outcome = session
            .process_frame(
                Frame::SubscribeAt {
                    topic: TopicName("t".into()),
                    offset: TopicOffset(0),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::ServerDraining);
            }
            _ => panic!("expected ServerDraining error"),
        }
    }

    #[tokio::test]
    async fn draining_rejects_subscribe_from() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;
        session.transition_to_draining();

        let outcome = session
            .process_frame(
                Frame::SubscribeFrom {
                    topic: TopicName("t".into()),
                    timestamp_millis: 1_700_000_000_000,
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::ServerDraining);
            }
            _ => panic!("expected ServerDraining error"),
        }
    }

    #[tokio::test]
    async fn draining_allows_publish() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;
        session.transition_to_draining();

        let outcome = session
            .process_frame(
                Frame::Publish {
                    topic: TopicName("t".into()),
                    headers: vec![],
                    ttl_ms: 0,
                    payload: Bytes::from_static(b"hello"),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        assert!(
            matches!(outcome, FrameOutcome::Reply(Frame::PublishAck { .. })),
            "expected PublishAck during drain, got {outcome:?}"
        );
    }

    #[tokio::test]
    async fn draining_allows_unsubscribe() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        // Subscribe first
        let _ = session
            .process_frame(
                Frame::Subscribe {
                    topic: TopicName("t".into()),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;

        session.transition_to_draining();

        let outcome = session
            .process_frame(
                Frame::Unsubscribe {
                    topic: TopicName("t".into()),
                    group: None,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        assert!(
            matches!(outcome, FrameOutcome::Reply(Frame::UnsubscribeAck { .. })),
            "expected UnsubscribeAck during drain, got {outcome:?}"
        );
    }

    #[tokio::test]
    async fn transition_to_draining_noop_when_not_authenticated() {
        let (mut session, _) = make_session();
        session.transition_to_draining();
        // Should remain PreAuth
        assert!(matches!(session, Session::PreAuth(_)));
    }

    #[tokio::test]
    async fn capabilities_in_draining_state() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        // Connect with capabilities
        let outcome = session
            .process_frame(
                Frame::Connect {
                    version: 1,
                    namespace: NamespaceId("test-ns".into()),
                    api_key: ApiKey("test-key".into()),
                    capabilities: crate::codec::frame::CAP_GOING_AWAY,
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        assert!(matches!(
            outcome,
            FrameOutcome::Reply(Frame::ConnectAck { success: true, .. })
        ));

        assert_eq!(session.capabilities(), crate::codec::frame::CAP_GOING_AWAY);

        session.transition_to_draining();
        // Capabilities preserved in Draining state
        assert_eq!(session.capabilities(), crate::codec::frame::CAP_GOING_AWAY);
    }

    // -- Reconnect tests --

    fn reconnect_frame(subs: Vec<crate::codec::frame::SubscriptionEntry>) -> Frame {
        Frame::Reconnect {
            version: 1,
            namespace: NamespaceId("test-ns".into()),
            api_key: ApiKey("test-key".into()),
            capabilities: crate::codec::frame::CAP_GOING_AWAY | crate::codec::frame::CAP_RECONNECT,
            subscriptions: subs,
        }
    }

    #[tokio::test]
    async fn reconnect_successful_no_subscriptions() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, mut conn_rx, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(reconnect_frame(vec![]), &auth, &manager, &delivery_tx)
            .await;

        match outcome {
            FrameOutcome::Reply(Frame::ReconnectAck {
                status,
                server_capabilities,
                results,
                ..
            }) => {
                assert_eq!(status, crate::codec::frame::ReconnectStatus::Success);
                assert!(results.is_empty());
                // Capabilities negotiated
                assert_ne!(server_capabilities, 0);
            }
            _ => panic!("expected ReconnectAck, got {outcome:?}"),
        }

        // Should have called manager.connect()
        let (ns, client) = conn_rx.try_recv().unwrap();
        assert_eq!(ns, NamespaceId("test-ns".into()));
        assert_eq!(client, ClientId(1));

        // Session should be Authenticated
        assert!(matches!(session, Session::Authenticated(_)));
    }

    #[tokio::test]
    async fn reconnect_with_subscriptions() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, _, mut sub_rx, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let subs = vec![
            crate::codec::frame::SubscriptionEntry {
                topic: TopicName("events".into()),
                group: None,
                last_acked_offset: TopicOffset(499),
            },
            crate::codec::frame::SubscriptionEntry {
                topic: TopicName("logs".into()),
                group: Some(GroupName("workers".into())),
                last_acked_offset: TopicOffset(1234),
            },
        ];

        let outcome = session
            .process_frame(reconnect_frame(subs), &auth, &manager, &delivery_tx)
            .await;

        match outcome {
            FrameOutcome::Reply(Frame::ReconnectAck {
                status, results, ..
            }) => {
                assert_eq!(status, crate::codec::frame::ReconnectStatus::Success);
                assert_eq!(results.len(), 2);
                // StubNamespaceManager returns TopicOffset(0) for all subscribes,
                // which is < requested next offset, so status should be Restored.
                // Actually offset 0 < 500, so this would be Restored (0 is not > 500).
                assert_eq!(
                    results[0].status,
                    crate::codec::frame::SubscriptionStatus::Restored
                );
                assert_eq!(
                    results[1].status,
                    crate::codec::frame::SubscriptionStatus::Restored
                );
            }
            _ => panic!("expected ReconnectAck, got {outcome:?}"),
        }

        // Verify subscribe was called twice
        let (ns1, topic1, _, _) = sub_rx.try_recv().unwrap();
        assert_eq!(ns1, NamespaceId("test-ns".into()));
        assert_eq!(topic1, TopicName("events".into()));

        let (_, topic2, _, group2) = sub_rx.try_recv().unwrap();
        assert_eq!(topic2, TopicName("logs".into()));
        assert_eq!(group2, Some(GroupName("workers".into())));
    }

    #[tokio::test]
    async fn reconnect_auth_failed() {
        let (mut session, _) = make_session();
        // Auth that expects a different key
        let auth = MockAuthenticator {
            namespace: NamespaceId("test-ns".into()),
            key: ApiKey("different-key".into()),
            permissions: PermissionSet::from_slice(&[Permission::Subscribe]),
        };
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(reconnect_frame(vec![]), &auth, &manager, &delivery_tx)
            .await;

        match outcome {
            FrameOutcome::ReplyAndClose(Frame::ReconnectAck { status, .. }) => {
                assert_eq!(status, crate::codec::frame::ReconnectStatus::AuthFailed);
            }
            _ => panic!("expected ReplyAndClose with AuthFailed, got {outcome:?}"),
        }

        // Session should remain PreAuth
        assert!(matches!(session, Session::PreAuth(_)));
    }

    #[tokio::test]
    async fn reconnect_bad_version() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Reconnect {
                    version: 2,
                    namespace: NamespaceId("test-ns".into()),
                    api_key: ApiKey("test-key".into()),
                    capabilities: 0,
                    subscriptions: vec![],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;

        match outcome {
            FrameOutcome::ReplyAndClose(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::UnsupportedVersion);
            }
            _ => panic!("expected UnsupportedVersion error, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn reconnect_no_subscribe_permission() {
        let (mut session, _) = make_session();
        let auth = make_auth(&[Permission::Publish, Permission::Query]);
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let subs = vec![crate::codec::frame::SubscriptionEntry {
            topic: TopicName("events".into()),
            group: None,
            last_acked_offset: TopicOffset(0),
        }];

        let outcome = session
            .process_frame(reconnect_frame(subs), &auth, &manager, &delivery_tx)
            .await;

        match outcome {
            FrameOutcome::Reply(Frame::ReconnectAck {
                status, results, ..
            }) => {
                assert_eq!(status, crate::codec::frame::ReconnectStatus::Partial);
                assert_eq!(results.len(), 1);
                assert_eq!(
                    results[0].status,
                    crate::codec::frame::SubscriptionStatus::PermissionDenied
                );
            }
            _ => panic!("expected Partial ReconnectAck, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn reconnect_rejected_when_already_authenticated() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(reconnect_frame(vec![]), &auth, &manager, &delivery_tx)
            .await;

        match outcome {
            FrameOutcome::ReplyAndClose(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::InvalidFrame);
            }
            _ => panic!("expected InvalidFrame error, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn reconnect_capability_negotiation() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Reconnect {
                    version: 1,
                    namespace: NamespaceId("test-ns".into()),
                    api_key: ApiKey("test-key".into()),
                    capabilities: crate::codec::frame::CAP_GOING_AWAY,
                    subscriptions: vec![],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;

        match outcome {
            FrameOutcome::Reply(Frame::ReconnectAck {
                server_capabilities,
                ..
            }) => {
                // Should agree only on CAP_GOING_AWAY (not CAP_RECONNECT since
                // the client didn't request it)
                assert_eq!(server_capabilities, crate::codec::frame::CAP_GOING_AWAY);
            }
            _ => panic!("expected ReconnectAck"),
        }

        // Session capabilities should match
        assert_eq!(session.capabilities(), crate::codec::frame::CAP_GOING_AWAY);
    }

    #[tokio::test]
    async fn reconnect_negotiates_cap_nack() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Reconnect {
                    version: 1,
                    namespace: NamespaceId("test-ns".into()),
                    api_key: ApiKey("test-key".into()),
                    capabilities: CAP_NACK,
                    subscriptions: vec![],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;

        match outcome {
            FrameOutcome::Reply(Frame::ReconnectAck {
                server_capabilities,
                ..
            }) => {
                assert_eq!(server_capabilities, CAP_NACK);
            }
            _ => panic!("expected ReconnectAck"),
        }

        assert_eq!(session.capabilities(), CAP_NACK);
    }

    #[tokio::test]
    async fn reconnect_namespace_mismatch() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let outcome = session
            .process_frame(
                Frame::Reconnect {
                    version: 1,
                    namespace: NamespaceId("wrong-ns".into()),
                    api_key: ApiKey("test-key".into()),
                    capabilities: 0,
                    subscriptions: vec![],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;

        match outcome {
            FrameOutcome::ReplyAndClose(Frame::ReconnectAck {
                status, message, ..
            }) => {
                assert_eq!(status, crate::codec::frame::ReconnectStatus::AuthFailed);
                assert!(
                    message.contains("namespace mismatch"),
                    "expected namespace mismatch message, got: {message}"
                );
            }
            _ => panic!("expected ReplyAndClose with AuthFailed, got {outcome:?}"),
        }

        // Session should remain PreAuth
        assert!(matches!(session, Session::PreAuth(_)));
    }

    #[tokio::test]
    async fn reconnect_partial_failure_rolls_back_subscriptions() {
        let (mut session, _) = make_session();
        // Only Publish permission — no Subscribe. This means all subscription
        // entries will fail with PermissionDenied and get rolled back.
        let auth = make_auth(&[Permission::Publish, Permission::Query]);
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        let subs = vec![
            crate::codec::frame::SubscriptionEntry {
                topic: TopicName("events".into()),
                group: None,
                last_acked_offset: TopicOffset(0),
            },
            crate::codec::frame::SubscriptionEntry {
                topic: TopicName("logs".into()),
                group: None,
                last_acked_offset: TopicOffset(0),
            },
        ];

        let outcome = session
            .process_frame(reconnect_frame(subs), &auth, &manager, &delivery_tx)
            .await;

        match &outcome {
            FrameOutcome::Reply(Frame::ReconnectAck {
                status, results, ..
            }) => {
                assert_eq!(*status, crate::codec::frame::ReconnectStatus::Partial);
                assert_eq!(results.len(), 2);
                // All should be PermissionDenied
                for r in results {
                    assert_eq!(
                        r.status,
                        crate::codec::frame::SubscriptionStatus::PermissionDenied
                    );
                }
            }
            _ => panic!("expected Partial ReconnectAck, got {outcome:?}"),
        }

        // Session should be Authenticated (connection is established) but with
        // NO active subscriptions due to rollback.
        match &session {
            Session::Authenticated(state) => {
                assert!(
                    state.subscriptions.is_empty(),
                    "expected no active subscriptions after rollback, got {}",
                    state.subscriptions.len()
                );
            }
            _ => panic!("expected Authenticated state"),
        }
    }

    #[tokio::test]
    async fn nack_without_cap_nack_rejected() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        // Authenticate WITHOUT CAP_NACK (legacy client).
        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::Nack {
                    message_id: MessageId(1),
                    reason: "test".into(),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, message }) => {
                assert_eq!(code, ErrorCode::InvalidFrame);
                assert!(message.contains("CAP_NACK"));
            }
            _ => panic!("expected InvalidFrame error, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn nack_unknown_message_id() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session_with_caps(&mut session, &auth, &manager, &delivery_tx, CAP_NACK).await;

        let outcome = session
            .process_frame(
                Frame::Nack {
                    message_id: MessageId(999),
                    reason: "bad data".into(),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, .. }) => {
                assert_eq!(code, ErrorCode::InternalError);
            }
            _ => panic!("expected Reply with InternalError, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn nack_known_message_returns_nack_ack() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session_with_caps(&mut session, &auth, &manager, &delivery_tx, CAP_NACK).await;

        // Deliver a message to populate inflight_topics
        let msg = PublishedMessage {
            envelope: Arc::new(MessageEnvelope::new(
                MessageId(42),
                NamespaceId("test-ns".into()),
                TopicName("events".into()),
                vec![1, 2, 3],
                now_millis(),
            )),
            offset: TopicOffset(10),
        };
        session.build_deliver_frame(&msg);

        // Nack the delivered message
        let outcome = session
            .process_frame(
                Frame::Nack {
                    message_id: MessageId(42),
                    reason: "invalid payload".into(),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::NackAck { message_id, status }) => {
                assert_eq!(message_id, MessageId(42));
                assert_eq!(status, 0x00); // DeadLettered
            }
            _ => panic!("expected NackAck reply, got {outcome:?}"),
        }

        // Verify message_id removed from inflight_topics
        match &session {
            Session::Authenticated(state) => {
                assert!(
                    !state.inflight_topics.contains_key(&MessageId(42)),
                    "message should be removed from inflight after nack"
                );
            }
            _ => panic!("expected Authenticated state"),
        }
    }

    #[tokio::test]
    async fn nack_allowed_during_drain() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session_with_caps(&mut session, &auth, &manager, &delivery_tx, CAP_NACK).await;

        // Deliver a message to populate inflight
        let msg = PublishedMessage {
            envelope: Arc::new(MessageEnvelope::new(
                MessageId(55),
                NamespaceId("test-ns".into()),
                TopicName("events".into()),
                vec![4, 5, 6],
                now_millis(),
            )),
            offset: TopicOffset(20),
        };
        session.build_deliver_frame(&msg);

        // Transition to draining
        session.transition_to_draining();
        assert!(matches!(session, Session::Draining(_)));

        // Nack should still be processed during drain
        let outcome = session
            .process_frame(
                Frame::Nack {
                    message_id: MessageId(55),
                    reason: "drain nack".into(),
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::NackAck { message_id, status }) => {
                assert_eq!(message_id, MessageId(55));
                assert_eq!(status, 0x00);
            }
            _ => panic!("expected NackAck reply during drain, got {outcome:?}"),
        }
    }

    // --- BatchAck tests ---

    #[tokio::test]
    async fn batch_ack_without_cap_rejected() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        // Authenticate WITHOUT CAP_BATCH_ACK.
        authenticate_session(&mut session, &auth, &manager, &delivery_tx).await;

        let outcome = session
            .process_frame(
                Frame::BatchAck {
                    message_ids: vec![MessageId(1)],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::Error { code, message }) => {
                assert_eq!(code, ErrorCode::InvalidFrame);
                assert!(message.contains("CAP_BATCH_ACK"));
            }
            _ => panic!("expected InvalidFrame error, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn batch_ack_all_unknown_returns_counts() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session_with_caps(&mut session, &auth, &manager, &delivery_tx, CAP_BATCH_ACK)
            .await;

        // BatchAck with IDs not in inflight_topics → all unknown
        let outcome = session
            .process_frame(
                Frame::BatchAck {
                    message_ids: vec![MessageId(100), MessageId(200), MessageId(300)],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::BatchAckResponse {
                acked_count,
                unknown_count,
            }) => {
                assert_eq!(acked_count, 0);
                assert_eq!(unknown_count, 3);
            }
            _ => panic!("expected BatchAckResponse, got {outcome:?}"),
        }
    }

    #[tokio::test]
    async fn batch_ack_mixed_known_and_unknown() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session_with_caps(&mut session, &auth, &manager, &delivery_tx, CAP_BATCH_ACK)
            .await;

        // Deliver messages to populate inflight_topics
        let topic = TopicName("test-topic".into());
        for id in [10, 20, 30] {
            let msg = PublishedMessage {
                envelope: Arc::new(MessageEnvelope::new(
                    MessageId(id),
                    NamespaceId("test-ns".into()),
                    topic.clone(),
                    b"data".to_vec(),
                    0,
                )),
                offset: TopicOffset(id),
            };
            session.build_deliver_frame(&msg);
        }

        // BatchAck: 10 and 30 are known, 99 is unknown
        let outcome = session
            .process_frame(
                Frame::BatchAck {
                    message_ids: vec![MessageId(10), MessageId(99), MessageId(30)],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::BatchAckResponse {
                acked_count,
                unknown_count,
            }) => {
                assert_eq!(acked_count, 2);
                assert_eq!(unknown_count, 1);
            }
            _ => panic!("expected BatchAckResponse, got {outcome:?}"),
        }

        // Verify 10 and 30 removed from inflight_topics, 20 still there
        match &session {
            Session::Authenticated(state) => {
                assert!(!state.inflight_topics.contains_key(&MessageId(10)));
                assert!(!state.inflight_topics.contains_key(&MessageId(30)));
                assert!(state.inflight_topics.contains_key(&MessageId(20)));
            }
            _ => panic!("expected Authenticated state"),
        }
    }

    #[tokio::test]
    async fn batch_ack_empty_returns_zero_counts() {
        let (mut session, _) = make_session();
        let auth = make_all_perms_auth();
        let (manager, ..) = StubNamespaceManager::new();
        let (delivery_tx, _delivery_rx) = mpsc::channel(16);

        authenticate_session_with_caps(&mut session, &auth, &manager, &delivery_tx, CAP_BATCH_ACK)
            .await;

        let outcome = session
            .process_frame(
                Frame::BatchAck {
                    message_ids: vec![],
                },
                &auth,
                &manager,
                &delivery_tx,
            )
            .await;
        match outcome {
            FrameOutcome::Reply(Frame::BatchAckResponse {
                acked_count,
                unknown_count,
            }) => {
                assert_eq!(acked_count, 0);
                assert_eq!(unknown_count, 0);
            }
            _ => panic!("expected BatchAckResponse, got {outcome:?}"),
        }
    }
}
