rivven_client/

client.rs

use crate::{Error, MessageData, Request, Response, Result};
use bytes::Bytes;
use rivven_core::PasswordHash;
use rivven_protocol::SyncGroupAssignments;
use sha2::{Digest, Sha256};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpStream;
use tracing::{debug, info};

#[cfg(feature = "tls")]
use rivven_core::tls::{TlsClientStream, TlsConfig, TlsConnector};

// Default maximum response size (100 MB) - prevents malicious server from exhausting client memory
const DEFAULT_MAX_RESPONSE_SIZE: usize = 100 * 1024 * 1024;

// Maximum request size — aligned with server `max_request_size` and
// `rivven_protocol::MAX_MESSAGE_SIZE` (10 MiB). Requests exceeding this
// are rejected client-side before touching the wire, preventing TCP-level
// deadlocks when the server must drain an oversized body.
const DEFAULT_MAX_REQUEST_SIZE: usize = rivven_protocol::MAX_MESSAGE_SIZE;

/// Default connection timeout (10 seconds) — prevents hanging on
/// unreachable hosts instead of relying on OS TCP timeout (~75-120s).
const DEFAULT_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);

/// Default request timeout (30 seconds) — prevents callers from
/// blocking forever when a server stalls or stops responding.
const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);

// ============================================================================
// Stream Wrapper
// ============================================================================

/// Wrapper for either plaintext or TLS streams
/// Note: TLS variant is significantly larger due to TLS state, but boxing
/// would add indirection overhead for every I/O operation
#[allow(clippy::large_enum_variant)]
pub(crate) enum ClientStream {
    Plaintext(TcpStream),
    #[cfg(feature = "tls")]
    Tls(TlsClientStream<TcpStream>),
}

impl AsyncRead for ClientStream {
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            ClientStream::Plaintext(s) => std::pin::Pin::new(s).poll_read(cx, buf),
            #[cfg(feature = "tls")]
            ClientStream::Tls(s) => std::pin::Pin::new(s).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for ClientStream {
    fn poll_write(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        match self.get_mut() {
            ClientStream::Plaintext(s) => std::pin::Pin::new(s).poll_write(cx, buf),
            #[cfg(feature = "tls")]
            ClientStream::Tls(s) => std::pin::Pin::new(s).poll_write(cx, buf),
        }
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            ClientStream::Plaintext(s) => std::pin::Pin::new(s).poll_flush(cx),
            #[cfg(feature = "tls")]
            ClientStream::Tls(s) => std::pin::Pin::new(s).poll_flush(cx),
        }
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            ClientStream::Plaintext(s) => std::pin::Pin::new(s).poll_shutdown(cx),
            #[cfg(feature = "tls")]
            ClientStream::Tls(s) => std::pin::Pin::new(s).poll_shutdown(cx),
        }
    }
}

// ============================================================================
// Client
// ============================================================================

/// Rivven client for connecting to a Rivven server
pub struct Client {
    stream: ClientStream,
    next_correlation_id: u32,
    /// Per-request timeout for send_request() I/O.
    request_timeout: Duration,
    /// Set to true when the stream is desynchronized (e.g. correlation ID
    /// mismatch). All subsequent requests immediately return an error,
    /// forcing the caller to reconnect.
    poisoned: bool,
}

impl Client {
    /// Connect to a Rivven server (plaintext)
    ///
    /// Automatically sends a protocol handshake after connecting.
    /// The handshake validates protocol version compatibility before any
    /// application requests are sent.
    ///
    /// Uses a default connection timeout of 10 seconds. For a custom
    /// timeout, use [`connect_with_timeout`](Self::connect_with_timeout).
    pub async fn connect(addr: &str) -> Result<Self> {
        Self::connect_with_timeout(addr, DEFAULT_CONNECTION_TIMEOUT).await
    }

    /// Connect to a Rivven server with a custom connection timeout.
    ///
    /// Wraps the TCP connect + handshake in `tokio::time::timeout` so
    /// callers never block longer than `timeout` on an unreachable host.
    pub async fn connect_with_timeout(addr: &str, timeout: Duration) -> Result<Self> {
        info!("Connecting to Rivven server at {}", addr);
        let stream = tokio::time::timeout(timeout, TcpStream::connect(addr))
            .await
            .map_err(|_| Error::TimeoutWithMessage(format!("Connection to {} timed out", addr)))?
            .map_err(|e| Error::ConnectionError(e.to_string()))?;

        // Disable Nagle's algorithm for lower latency on small request/response RPC.
        let _ = stream.set_nodelay(true);

        let mut client = Self {
            stream: ClientStream::Plaintext(stream),
            next_correlation_id: 0,
            request_timeout: DEFAULT_REQUEST_TIMEOUT,
            poisoned: false,
        };

        // auto-handshake on connect
        client.handshake("rivven-client").await?;

        Ok(client)
    }

    /// Connect to a Rivven server with TLS
    #[cfg(feature = "tls")]
    pub async fn connect_tls(
        addr: &str,
        tls_config: &TlsConfig,
        server_name: &str,
    ) -> Result<Self> {
        Self::connect_tls_with_timeout(addr, tls_config, server_name, DEFAULT_CONNECTION_TIMEOUT)
            .await
    }

    /// Connect to a Rivven server with TLS and a custom connection timeout.
    #[cfg(feature = "tls")]
    pub async fn connect_tls_with_timeout(
        addr: &str,
        tls_config: &TlsConfig,
        server_name: &str,
        timeout: Duration,
    ) -> Result<Self> {
        info!("Connecting to Rivven server at {} with TLS", addr);

        // Resolve address — supports both IP:port and DNS hostname:port,
        // unlike the old SocketAddr::parse() which rejected DNS names.
        let tcp_stream = tokio::time::timeout(timeout, TcpStream::connect(addr))
            .await
            .map_err(|_| {
                Error::TimeoutWithMessage(format!("TLS connection to {} timed out", addr))
            })?
            .map_err(|e| Error::ConnectionError(format!("TCP connection error: {}", e)))?;

        // Disable Nagle for low-latency RPC (matches plaintext path)
        tcp_stream
            .set_nodelay(true)
            .map_err(|e| Error::ConnectionError(format!("Failed to set TCP_NODELAY: {}", e)))?;

        // Create TLS connector
        let connector = TlsConnector::new(tls_config)
            .map_err(|e| Error::ConnectionError(format!("TLS config error: {}", e)))?;

        // Wrap the TCP stream in TLS
        let tls_stream = connector
            .connect(tcp_stream, server_name)
            .await
            .map_err(|e| Error::ConnectionError(format!("TLS handshake error: {}", e)))?;

        info!("TLS connection established to {} ({})", addr, server_name);

        let mut client = Self {
            stream: ClientStream::Tls(tls_stream),
            next_correlation_id: 0,
            request_timeout: DEFAULT_REQUEST_TIMEOUT,
            poisoned: false,
        };

        // auto-handshake on TLS connect
        client.handshake("rivven-client").await?;

        Ok(client)
    }

    /// Connect with mTLS (mutual TLS) using client certificate
    #[cfg(feature = "tls")]
    pub async fn connect_mtls(
        addr: &str,
        cert_path: impl Into<std::path::PathBuf>,
        key_path: impl Into<std::path::PathBuf>,
        ca_path: impl Into<std::path::PathBuf> + Clone,
        server_name: &str,
    ) -> Result<Self> {
        let tls_config = TlsConfig::mtls_from_pem_files(cert_path, key_path, ca_path);
        Self::connect_tls(addr, &tls_config, server_name).await
    }

    // ========================================================================
    // Handshake
    // ========================================================================

    /// Send a protocol version handshake to the server.
    ///
    /// Validates that the server speaks a compatible protocol version.
    /// Called automatically by `connect()` and `connect_tls()`.
    pub async fn handshake(&mut self, client_id: &str) -> Result<()> {
        let request = Request::Handshake {
            protocol_version: rivven_protocol::PROTOCOL_VERSION,
            client_id: client_id.to_string(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::HandshakeResult {
                compatible,
                message: _,
                server_version,
            } => {
                if compatible {
                    info!(
                        "Handshake OK (client v{}, server v{})",
                        rivven_protocol::PROTOCOL_VERSION,
                        server_version
                    );
                    Ok(())
                } else {
                    Err(Error::ProtocolError(
                        rivven_protocol::ProtocolError::VersionMismatch {
                            expected: rivven_protocol::PROTOCOL_VERSION,
                            actual: server_version,
                        },
                    ))
                }
            }
            Response::Error { message } => {
                // Server doesn't support handshake — log warning but proceed
                // for backward compatibility with older servers
                tracing::warn!(
                    "Server returned error on handshake: {}, proceeding anyway",
                    message
                );
                Ok(())
            }
            _ => {
                // Unknown response — old server, proceed anyway
                tracing::warn!(
                    "Server did not return HandshakeResult, proceeding without version check"
                );
                Ok(())
            }
        }
    }

    /// Consume the client and return the underlying TCP stream.
    ///
    /// This is used by the producer to hand off the authenticated +
    /// handshaked connection to its background sender task.
    /// Only works for plaintext connections; for TLS use `into_client_stream()`.
    pub fn into_stream(self) -> Result<TcpStream> {
        match self.stream {
            ClientStream::Plaintext(s) => Ok(s),
            #[cfg(feature = "tls")]
            ClientStream::Tls(_) => Err(Error::ConnectionError(
                "Cannot extract TcpStream from TLS connection. Use into_client_stream() instead."
                    .to_string(),
            )),
        }
    }

    /// Consume the client and return the underlying `ClientStream`.
    ///
    /// Works for both plaintext and TLS connections. Used by the
    /// producer to hand off the authenticated + handshaked connection
    /// to its background sender task.
    pub(crate) fn into_client_stream(self) -> ClientStream {
        self.stream
    }

    /// Set the per-request timeout for `send_request()` I/O.
    pub fn set_request_timeout(&mut self, timeout: Duration) {
        self.request_timeout = timeout;
    }

    /// Returns `true` when the underlying transport is TLS-encrypted.
    pub fn is_tls(&self) -> bool {
        match &self.stream {
            ClientStream::Plaintext(_) => false,
            #[cfg(feature = "tls")]
            ClientStream::Tls(_) => true,
        }
    }

    // ========================================================================
    // Authentication Methods
    // ========================================================================

    /// Authenticate with simple username/password
    ///
    /// # Security — plaintext credentials
    ///
    /// This uses SASL/PLAIN which sends the password in **cleartext** on the
    /// wire.  The client automatically sets `require_tls = true` when the
    /// underlying connection is TLS-encrypted so the server will reject the
    /// request if it somehow arrives over a non-TLS channel.  For untrusted
    /// networks, prefer `authenticate_scram()` which never sends the password.
    #[allow(deprecated)]
    pub async fn authenticate(&mut self, username: &str, password: &str) -> Result<AuthSession> {
        // Client-side guard: SASL/PLAIN sends the password in cleartext.
        // Without TLS, a network observer captures credentials immediately.
        if !self.is_tls() {
            return Err(Error::AuthenticationFailed(
                "SASL/PLAIN requires a TLS connection — use authenticate_scram() for plaintext channels".to_string(),
            ));
        }

        // Tell the server whether we are on a TLS connection so it can
        // reject the request when TLS is not active.
        let require_tls = true; // always true — we checked above
        let request = Request::Authenticate {
            username: username.to_string(),
            password: password.to_string(),
            require_tls,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Authenticated {
                session_id,
                expires_in,
            } => {
                info!("Authenticated as '{}'", username);
                Ok(AuthSession {
                    session_id,
                    expires_in,
                })
            }
            Response::Error { message } => Err(Error::AuthenticationFailed(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Authenticate using SCRAM-SHA-256 (secure challenge-response)
    ///
    /// SCRAM-SHA-256 (RFC 5802/7677) provides:
    /// - Password never sent over the wire
    /// - Mutual authentication (server proves it knows password too)
    /// - Protection against replay attacks
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// let session = client.authenticate_scram("alice", "password123").await?;
    /// println!("Session: {} (expires in {}s)", session.session_id, session.expires_in);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn authenticate_scram(
        &mut self,
        username: &str,
        password: &str,
    ) -> Result<AuthSession> {
        // Step 1: Generate client nonce and send client-first message
        let client_nonce = generate_nonce();
        let client_first_bare = format!("n={},r={}", escape_username(username), client_nonce);
        let client_first = format!("n,,{}", client_first_bare);

        debug!("SCRAM: Sending client-first");
        let request = Request::ScramClientFirst {
            message: Bytes::from(client_first.clone()),
        };

        let response = self.send_request(request).await?;

        // Step 2: Parse server-first message
        let server_first = match response {
            Response::ScramServerFirst { message } => String::from_utf8(message.to_vec())
                .map_err(|_| Error::AuthenticationFailed("Invalid server-first encoding".into()))?,
            Response::Error { message } => return Err(Error::AuthenticationFailed(message)),
            _ => return Err(Error::InvalidResponse),
        };

        debug!("SCRAM: Received server-first");

        // Parse server-first: r=<nonce>,s=<salt>,i=<iterations>
        let (combined_nonce, salt_b64, iterations) = parse_server_first(&server_first)?;

        // Verify server nonce starts with our client nonce
        if !combined_nonce.starts_with(&client_nonce) {
            return Err(Error::AuthenticationFailed("Server nonce mismatch".into()));
        }

        // Decode salt
        let salt = base64_decode(&salt_b64)
            .map_err(|_| Error::AuthenticationFailed("Invalid salt encoding".into()))?;

        // Step 3: Compute client proof
        let salted_password = pbkdf2_sha256(password.as_bytes(), &salt, iterations);
        let client_key = PasswordHash::hmac_sha256(&salted_password, b"Client Key");
        let stored_key = sha256(&client_key);

        let client_final_without_proof = format!("c=biws,r={}", combined_nonce);
        let auth_message = format!(
            "{},{},{}",
            client_first_bare, server_first, client_final_without_proof
        );

        let client_signature = PasswordHash::hmac_sha256(&stored_key, auth_message.as_bytes());
        let client_proof = xor_bytes(&client_key, &client_signature);
        let client_proof_b64 = base64_encode(&client_proof);

        // Step 4: Send client-final message
        let client_final = format!("{},p={}", client_final_without_proof, client_proof_b64);

        debug!("SCRAM: Sending client-final");
        let request = Request::ScramClientFinal {
            message: Bytes::from(client_final),
        };

        let response = self.send_request(request).await?;

        // Step 5: Verify server-final and get session
        match response {
            Response::ScramServerFinal {
                message,
                session_id,
                expires_in,
            } => {
                let server_final = String::from_utf8(message.to_vec()).map_err(|_| {
                    Error::AuthenticationFailed("Invalid server-final encoding".into())
                })?;

                // Check for error response
                if let Some(error_msg) = server_final.strip_prefix("e=") {
                    return Err(Error::AuthenticationFailed(error_msg.to_string()));
                }

                // Verify server signature (mutual authentication)
                if let Some(verifier_b64) = server_final.strip_prefix("v=") {
                    let server_key = PasswordHash::hmac_sha256(&salted_password, b"Server Key");
                    let expected_server_sig =
                        PasswordHash::hmac_sha256(&server_key, auth_message.as_bytes());
                    let expected_verifier = base64_encode(&expected_server_sig);

                    if verifier_b64 != expected_verifier {
                        return Err(Error::AuthenticationFailed(
                            "Server verification failed".into(),
                        ));
                    }
                }

                let session_id = session_id.ok_or_else(|| {
                    Error::AuthenticationFailed("No session ID in response".into())
                })?;
                let expires_in = expires_in
                    .ok_or_else(|| Error::AuthenticationFailed("No expiry in response".into()))?;

                info!("SCRAM authentication successful for '{}'", username);
                Ok(AuthSession {
                    session_id,
                    expires_in,
                })
            }
            Response::Error { message } => Err(Error::AuthenticationFailed(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    // ========================================================================
    // Request/Response Handling
    // ========================================================================

    /// Send a request and receive a response.
    ///
    /// The entire I/O round-trip (write + read) is wrapped in
    /// `tokio::time::timeout` using the client's `request_timeout`
    /// (default 30 s). A stalled or unresponsive server will therefore
    /// return `Error::Timeout` instead of blocking the caller forever.
    pub(crate) async fn send_request(&mut self, request: Request) -> Result<Response> {
        let timeout_dur = self.request_timeout;
        match tokio::time::timeout(timeout_dur, self.send_request_inner(request)).await {
            Ok(result) => result,
            Err(_elapsed) => {
                // Timeout may have cancelled mid-I/O — stream is potentially
                // desynchronized. Poison so the next call reconnects.
                self.poisoned = true;
                Err(Error::Timeout)
            }
        }
    }

    /// Inner implementation of send_request without timeout wrapper.
    async fn send_request_inner(&mut self, request: Request) -> Result<Response> {
        // Fail immediately if the stream is desynchronized. The caller must
        // reconnect to get a new, clean Client instance.
        if self.poisoned {
            return Err(Error::ConnectionError(
                "Client stream is desynchronized — reconnect required".into(),
            ));
        }

        // Generate sequential correlation ID
        let correlation_id = self.next_correlation_id;
        self.next_correlation_id = self.next_correlation_id.wrapping_add(1);

        // Serialize request with wire format prefix and correlation ID
        let request_bytes =
            request.to_wire(rivven_protocol::WireFormat::Postcard, correlation_id)?;

        // Reject oversized requests client-side before touching the wire.
        // This prevents a TCP-level deadlock where write_all() blocks waiting
        // for the server to read, while the server rejects and tries to respond.
        if request_bytes.len() > DEFAULT_MAX_REQUEST_SIZE {
            return Err(Error::RequestTooLarge(
                request_bytes.len(),
                DEFAULT_MAX_REQUEST_SIZE,
            ));
        }

        // Write length prefix + request.
        // After the first write_all succeeds, bytes may be on the wire.
        // Any subsequent I/O failure desynchronizes the TCP stream, so we
        // must poison the client to prevent silent corruption.
        let len: u32 = request_bytes
            .len()
            .try_into()
            .map_err(|_| Error::RequestTooLarge(request_bytes.len(), u32::MAX as usize))?;
        self.stream
            .write_all(&len.to_be_bytes())
            .await
            .map_err(|e| {
                self.poisoned = true;
                Error::from(e)
            })?;
        self.stream.write_all(&request_bytes).await.map_err(|e| {
            self.poisoned = true;
            Error::from(e)
        })?;
        self.stream.flush().await.map_err(|e| {
            self.poisoned = true;
            Error::from(e)
        })?;

        // Read length prefix — request was sent, so read failure desynchronizes
        let mut len_buf = [0u8; 4];
        self.stream.read_exact(&mut len_buf).await.map_err(|e| {
            self.poisoned = true;
            Error::from(e)
        })?;
        let msg_len = u32::from_be_bytes(len_buf) as usize;

        // Validate response size to prevent memory exhaustion from malicious server
        if msg_len > DEFAULT_MAX_RESPONSE_SIZE {
            self.poisoned = true;
            return Err(Error::ResponseTooLarge(msg_len, DEFAULT_MAX_RESPONSE_SIZE));
        }

        // Read response — partial read desynchronizes
        let mut response_buf = vec![0u8; msg_len];
        self.stream
            .read_exact(&mut response_buf)
            .await
            .map_err(|e| {
                self.poisoned = true;
                Error::from(e)
            })?;

        // Deserialize response (auto-detects wire format).
        // The full response was consumed from the wire, so framing is intact
        // regardless of deserialization outcome — no need to poison here.
        let (response, _format, response_correlation_id) = Response::from_wire(&response_buf)?;

        // Validate that the response correlation ID matches the
        // request we sent. A mismatch indicates stream desynchronization
        // (e.g. partial reads) or a buggy server.
        if response_correlation_id != correlation_id {
            // Mark the client as poisoned — the stream is no longer usable
            // because subsequent reads would parse at wrong byte boundaries.
            self.poisoned = true;
            return Err(Error::ProtocolError(
                rivven_protocol::ProtocolError::InvalidFormat(format!(
                    "Correlation ID mismatch: expected {}, got {}",
                    correlation_id, response_correlation_id
                )),
            ));
        }

        Ok(response)
    }

    /// Consume from multiple partitions using request pipelining.
    ///
    /// Sends all `Consume` requests back-to-back on the wire *before*
    /// reading any response. This eliminates per-partition round-trip
    /// latency and avoids head-of-line blocking when fetching from many
    /// partitions over a single connection.
    ///
    /// The returned `Vec` has one entry per input partition in the same
    /// order. Each entry is `Ok(messages)` or `Err(Error)`.
    pub async fn consume_pipelined(
        &mut self,
        fetches: &[(&str, u32, u64, u32, Option<u8>)],
    ) -> Result<Vec<Result<Vec<MessageData>>>> {
        if fetches.is_empty() {
            return Ok(Vec::new());
        }
        if self.poisoned {
            return Err(Error::ConnectionError(
                "Client stream is desynchronized — reconnect required".into(),
            ));
        }

        let timeout_dur = self.request_timeout;
        match tokio::time::timeout(timeout_dur, self.consume_pipelined_inner(fetches)).await {
            Ok(result) => result,
            Err(_elapsed) => {
                self.poisoned = true;
                Err(Error::Timeout)
            }
        }
    }

    async fn consume_pipelined_inner(
        &mut self,
        fetches: &[(&str, u32, u64, u32, Option<u8>)],
    ) -> Result<Vec<Result<Vec<MessageData>>>> {
        let mut correlation_ids = Vec::with_capacity(fetches.len());
        let mut bytes_sent = false;

        // Phase 1: Send all requests without waiting for responses.
        for &(topic, partition, offset, max_messages, isolation_level) in fetches {
            let correlation_id = self.next_correlation_id;
            self.next_correlation_id = self.next_correlation_id.wrapping_add(1);
            correlation_ids.push(correlation_id);

            let request = Request::Consume {
                topic: topic.to_string(),
                partition,
                offset,
                max_messages,
                isolation_level,
                max_wait_ms: None,
            };

            let request_bytes = request
                .to_wire(rivven_protocol::WireFormat::Postcard, correlation_id)
                .inspect_err(|_| {
                    if bytes_sent {
                        self.poisoned = true;
                    }
                })?;

            if request_bytes.len() > DEFAULT_MAX_REQUEST_SIZE {
                if bytes_sent {
                    self.poisoned = true;
                }
                return Err(Error::RequestTooLarge(
                    request_bytes.len(),
                    DEFAULT_MAX_REQUEST_SIZE,
                ));
            }

            let len: u32 = request_bytes.len().try_into().map_err(|_| {
                if bytes_sent {
                    self.poisoned = true;
                }
                Error::RequestTooLarge(request_bytes.len(), u32::MAX as usize)
            })?;
            self.stream
                .write_all(&len.to_be_bytes())
                .await
                .map_err(|e| {
                    if bytes_sent {
                        self.poisoned = true;
                    }
                    Error::from(e)
                })?;
            self.stream.write_all(&request_bytes).await.map_err(|e| {
                self.poisoned = true; // At least length prefix was sent
                Error::from(e)
            })?;
            bytes_sent = true;
        }
        self.stream.flush().await.map_err(|e| {
            self.poisoned = true;
            Error::from(e)
        })?;

        // Phase 2: Read all responses in-order.
        let mut results = Vec::with_capacity(fetches.len());
        let mut response_buf: Vec<u8> = Vec::with_capacity(4096);
        for &expected_cid in &correlation_ids {
            let mut len_buf = [0u8; 4];
            self.stream.read_exact(&mut len_buf).await.map_err(|e| {
                self.poisoned = true;
                Error::from(e)
            })?;
            let msg_len = u32::from_be_bytes(len_buf) as usize;

            if msg_len > DEFAULT_MAX_RESPONSE_SIZE {
                self.poisoned = true;
                return Err(Error::ResponseTooLarge(msg_len, DEFAULT_MAX_RESPONSE_SIZE));
            }

            response_buf.resize(msg_len, 0);
            self.stream
                .read_exact(&mut response_buf)
                .await
                .map_err(|e| {
                    self.poisoned = true;
                    Error::from(e)
                })?;
            let (response, _format, response_cid) = Response::from_wire(&response_buf)
                .inspect_err(|_| {
                    self.poisoned = true;
                })?;

            if response_cid != expected_cid {
                self.poisoned = true;
                return Err(Error::ProtocolError(
                    rivven_protocol::ProtocolError::InvalidFormat(format!(
                        "Correlation ID mismatch: expected {}, got {}",
                        expected_cid, response_cid
                    )),
                ));
            }

            let result = match response {
                Response::Messages { messages } => Ok(messages),
                Response::Error { message } => Err(Error::ServerError(message)),
                _ => Err(Error::InvalidResponse),
            };
            results.push(result);
        }

        Ok(results)
    }

    /// Publish a message to a topic
    pub async fn publish(
        &mut self,
        topic: impl Into<String>,
        value: impl Into<Bytes>,
    ) -> Result<u64> {
        self.publish_with_key(topic, None::<Bytes>, value).await
    }

    /// Publish a message with a key to a topic
    pub async fn publish_with_key(
        &mut self,
        topic: impl Into<String>,
        key: Option<impl Into<Bytes>>,
        value: impl Into<Bytes>,
    ) -> Result<u64> {
        let request = Request::Publish {
            topic: topic.into(),
            partition: None,
            key: key.map(|k| k.into()),
            value: value.into(),
            leader_epoch: None,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Published { offset, .. } => Ok(offset),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Publish a message to a specific partition
    pub async fn publish_to_partition(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        key: Option<impl Into<Bytes>>,
        value: impl Into<Bytes>,
    ) -> Result<u64> {
        let request = Request::Publish {
            topic: topic.into(),
            partition: Some(partition),
            key: key.map(|k| k.into()),
            value: value.into(),
            leader_epoch: None,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Published { offset, .. } => Ok(offset),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Consume messages from a topic partition
    ///
    /// Uses read_uncommitted isolation level (default).
    /// For transactional consumers that should not see aborted transaction messages,
    /// use [`Self::consume_with_isolation`] with `isolation_level = 1` (read_committed).
    pub async fn consume(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        offset: u64,
        max_messages: u32,
    ) -> Result<Vec<MessageData>> {
        self.consume_with_isolation(topic, partition, offset, max_messages, None)
            .await
    }

    /// Consume messages from a topic partition with specified isolation level
    ///
    /// # Arguments
    /// * `topic` - Topic name
    /// * `partition` - Partition number
    /// * `offset` - Starting offset
    /// * `max_messages` - Maximum messages to return
    /// * `isolation_level` - Transaction isolation level:
    ///   - `None` or `Some(0)` = read_uncommitted (default): Returns all messages
    ///   - `Some(1)` = read_committed: Filters out messages from aborted transactions
    ///
    /// # Read Committed Isolation
    ///
    /// When using `isolation_level = Some(1)` (read_committed), the consumer will:
    /// - Not see messages from transactions that were aborted
    /// - Not see control records (transaction markers)
    /// - Only see committed transactional messages
    ///
    /// This is essential for exactly-once semantics (EOS) consumers.
    pub async fn consume_with_isolation(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        offset: u64,
        max_messages: u32,
        isolation_level: Option<u8>,
    ) -> Result<Vec<MessageData>> {
        let request = Request::Consume {
            topic: topic.into(),
            partition,
            offset,
            max_messages,
            isolation_level,
            max_wait_ms: None,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Messages { messages } => Ok(messages),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Consume messages with long-polling support
    ///
    /// If no data is available immediately, the server will hold the request
    /// for up to `max_wait_ms` milliseconds before returning an empty response.
    /// This avoids tight polling loops and reduces network overhead.
    ///
    /// Capped server-side at 30 000 ms. `0` or `None` = immediate (no waiting).
    pub async fn consume_long_poll(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        offset: u64,
        max_messages: u32,
        isolation_level: Option<u8>,
        max_wait_ms: u64,
    ) -> Result<Vec<MessageData>> {
        let request = Request::Consume {
            topic: topic.into(),
            partition,
            offset,
            max_messages,
            isolation_level,
            max_wait_ms: Some(max_wait_ms),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Messages { messages } => Ok(messages),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Consume messages with read_committed isolation level
    ///
    /// This is a convenience method for transactional consumers that should
    /// only see committed messages. Messages from aborted transactions are filtered out.
    ///
    /// Equivalent to calling [`Self::consume_with_isolation`] with `isolation_level = Some(1)`.
    pub async fn consume_read_committed(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        offset: u64,
        max_messages: u32,
    ) -> Result<Vec<MessageData>> {
        self.consume_with_isolation(topic, partition, offset, max_messages, Some(1))
            .await
    }

    /// Create a new topic
    pub async fn create_topic(
        &mut self,
        name: impl Into<String>,
        partitions: Option<u32>,
    ) -> Result<u32> {
        let name = name.into();
        let request = Request::CreateTopic {
            name: name.clone(),
            partitions,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TopicCreated { partitions, .. } => Ok(partitions),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// List all topics
    pub async fn list_topics(&mut self) -> Result<Vec<String>> {
        let request = Request::ListTopics;
        let response = self.send_request(request).await?;

        match response {
            Response::Topics { topics } => Ok(topics),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Delete a topic
    pub async fn delete_topic(&mut self, name: impl Into<String>) -> Result<()> {
        let request = Request::DeleteTopic { name: name.into() };
        let response = self.send_request(request).await?;

        match response {
            Response::TopicDeleted => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Commit consumer offset
    pub async fn commit_offset(
        &mut self,
        consumer_group: impl Into<String>,
        topic: impl Into<String>,
        partition: u32,
        offset: u64,
    ) -> Result<()> {
        let request = Request::CommitOffset {
            consumer_group: consumer_group.into(),
            topic: topic.into(),
            partition,
            offset,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::OffsetCommitted => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Commit offsets for multiple partitions using request pipelining.
    ///
    /// Sends all `CommitOffset` requests at once, then reads all responses.
    /// Returns per-partition results in the same order as `offsets`.
    pub async fn commit_offsets_pipelined(
        &mut self,
        consumer_group: &str,
        offsets: &[(String, u32, u64)],
    ) -> Result<Vec<Result<()>>> {
        if offsets.is_empty() {
            return Ok(Vec::new());
        }
        if self.poisoned {
            return Err(Error::ConnectionError(
                "Client stream is desynchronized — reconnect required".into(),
            ));
        }

        let timeout_dur = self.request_timeout;
        match tokio::time::timeout(
            timeout_dur,
            self.commit_offsets_pipelined_inner(consumer_group, offsets),
        )
        .await
        {
            Ok(result) => result,
            Err(_elapsed) => {
                self.poisoned = true;
                Err(Error::Timeout)
            }
        }
    }

    async fn commit_offsets_pipelined_inner(
        &mut self,
        consumer_group: &str,
        offsets: &[(String, u32, u64)],
    ) -> Result<Vec<Result<()>>> {
        let mut correlation_ids = Vec::with_capacity(offsets.len());
        let mut bytes_sent = false;

        // Phase 1: Send all commit requests.
        for (topic, partition, offset) in offsets {
            let correlation_id = self.next_correlation_id;
            self.next_correlation_id = self.next_correlation_id.wrapping_add(1);
            correlation_ids.push(correlation_id);

            let request = Request::CommitOffset {
                consumer_group: consumer_group.to_string(),
                topic: topic.clone(),
                partition: *partition,
                offset: *offset,
            };

            let request_bytes = request
                .to_wire(rivven_protocol::WireFormat::Postcard, correlation_id)
                .inspect_err(|_| {
                    if bytes_sent {
                        self.poisoned = true;
                    }
                })?;

            if request_bytes.len() > DEFAULT_MAX_REQUEST_SIZE {
                if bytes_sent {
                    self.poisoned = true;
                }
                return Err(Error::RequestTooLarge(
                    request_bytes.len(),
                    DEFAULT_MAX_REQUEST_SIZE,
                ));
            }

            let len: u32 = request_bytes.len().try_into().map_err(|_| {
                if bytes_sent {
                    self.poisoned = true;
                }
                Error::RequestTooLarge(request_bytes.len(), u32::MAX as usize)
            })?;
            self.stream
                .write_all(&len.to_be_bytes())
                .await
                .map_err(|e| {
                    if bytes_sent {
                        self.poisoned = true;
                    }
                    Error::from(e)
                })?;
            self.stream.write_all(&request_bytes).await.map_err(|e| {
                self.poisoned = true;
                Error::from(e)
            })?;
            bytes_sent = true;
        }
        self.stream.flush().await.map_err(|e| {
            self.poisoned = true;
            Error::from(e)
        })?;

        // Phase 2: Read all responses in-order.
        let mut results = Vec::with_capacity(offsets.len());
        for &expected_cid in &correlation_ids {
            let mut len_buf = [0u8; 4];
            self.stream.read_exact(&mut len_buf).await.map_err(|e| {
                self.poisoned = true;
                Error::from(e)
            })?;
            let msg_len = u32::from_be_bytes(len_buf) as usize;

            if msg_len > DEFAULT_MAX_RESPONSE_SIZE {
                self.poisoned = true;
                return Err(Error::ResponseTooLarge(msg_len, DEFAULT_MAX_RESPONSE_SIZE));
            }

            let mut response_buf = vec![0u8; msg_len];
            self.stream
                .read_exact(&mut response_buf)
                .await
                .map_err(|e| {
                    self.poisoned = true;
                    Error::from(e)
                })?;
            let (response, _format, response_cid) = Response::from_wire(&response_buf)
                .inspect_err(|_| {
                    self.poisoned = true;
                })?;

            if response_cid != expected_cid {
                self.poisoned = true;
                return Err(Error::ProtocolError(
                    rivven_protocol::ProtocolError::InvalidFormat(format!(
                        "Correlation ID mismatch: expected {}, got {}",
                        expected_cid, response_cid
                    )),
                ));
            }

            let result = match response {
                Response::OffsetCommitted => Ok(()),
                Response::Error { message } => Err(Error::ServerError(message)),
                _ => Err(Error::InvalidResponse),
            };
            results.push(result);
        }

        Ok(results)
    }

    /// Returns `true` if the client stream is desynchronized and unusable.
    pub fn is_poisoned(&self) -> bool {
        self.poisoned
    }

    /// Get consumer offset
    pub async fn get_offset(
        &mut self,
        consumer_group: impl Into<String>,
        topic: impl Into<String>,
        partition: u32,
    ) -> Result<Option<u64>> {
        let request = Request::GetOffset {
            consumer_group: consumer_group.into(),
            topic: topic.into(),
            partition,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Offset { offset } => Ok(offset),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Get earliest and latest offsets for a topic partition
    ///
    /// Returns (earliest, latest) where:
    /// - earliest: First available offset (messages before this are deleted/compacted)
    /// - latest: Next offset to be assigned (one past the last message)
    pub async fn get_offset_bounds(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
    ) -> Result<(u64, u64)> {
        let request = Request::GetOffsetBounds {
            topic: topic.into(),
            partition,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::OffsetBounds { earliest, latest } => Ok((earliest, latest)),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Get topic metadata
    pub async fn get_metadata(&mut self, topic: impl Into<String>) -> Result<(String, u32)> {
        let request = Request::GetMetadata {
            topic: topic.into(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::Metadata { name, partitions } => Ok((name, partitions)),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Ping the server
    pub async fn ping(&mut self) -> Result<()> {
        let request = Request::Ping;
        let response = self.send_request(request).await?;

        match response {
            Response::Pong => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Register a schema with the schema registry (via HTTP REST API)
    ///
    /// The schema registry runs as a separate service (`rivven-schema`) with a
    /// Confluent-compatible REST API. This method performs a POST to
    /// `{registry_url}/subjects/{subject}/versions`.
    ///
    /// Supports both `http://` and `https://` registry URLs. HTTPS requires the
    /// `schema-registry` feature which brings in `reqwest` with `rustls-tls`.
    /// Without the feature flag, a minimal inline HTTP/1.1 client is used (HTTP only).
    ///
    /// # Arguments
    /// * `registry_url` - Schema registry base URL (e.g., `http://localhost:8081` or `https://registry.example.com`)
    /// * `subject` - Subject name (typically `{topic}-key` or `{topic}-value`)
    /// * `schema_type` - Schema format: `"AVRO"`, `"PROTOBUF"`, or `"JSON"`
    /// * `schema` - The schema definition string
    ///
    /// # Returns
    /// The global schema ID on success.
    pub async fn register_schema(
        &self,
        registry_url: &str,
        subject: &str,
        schema_type: &str,
        schema: &str,
    ) -> Result<u32> {
        let url = registry_url.trim_end_matches('/');
        let endpoint = format!("{}/subjects/{}/versions", url, subject);

        let body = serde_json::json!({
            "schema": schema,
            "schemaType": schema_type,
        });

        #[cfg(feature = "schema-registry")]
        {
            self.register_schema_reqwest(&endpoint, &body).await
        }

        #[cfg(not(feature = "schema-registry"))]
        {
            self.register_schema_inline(url, &endpoint, &body).await
        }
    }

    /// HTTPS-capable schema registration using reqwest (requires `schema-registry` feature)
    #[cfg(feature = "schema-registry")]
    async fn register_schema_reqwest(
        &self,
        endpoint: &str,
        body: &serde_json::Value,
    ) -> Result<u32> {
        let client = reqwest::Client::new();
        let response = client
            .post(endpoint)
            .header("Content-Type", "application/vnd.schemaregistry.v1+json")
            .json(body)
            .send()
            .await
            .map_err(|e| Error::ConnectionError(format!("schema registry request failed: {e}")))?;

        let status = response.status();
        if !status.is_success() {
            let body_text = response.text().await.unwrap_or_default();
            return Err(Error::ServerError(format!(
                "schema registry returned HTTP {status}: {body_text}"
            )));
        }

        #[derive(serde::Deserialize)]
        struct RegisterResponse {
            id: u32,
        }

        let result: RegisterResponse = response
            .json()
            .await
            .map_err(|e| Error::ConnectionError(format!("failed to parse response: {e}")))?;

        Ok(result.id)
    }

    /// Minimal inline HTTP/1.1 client for schema registration (HTTP only, no external deps).
    ///
    /// Handles both `Content-Length` and `Transfer-Encoding: chunked` responses.
    /// For production deployments behind proxies or with HTTPS requirements,
    /// enable the `schema-registry` feature to use `reqwest` instead.
    #[cfg(not(feature = "schema-registry"))]
    async fn register_schema_inline(
        &self,
        base_url: &str,
        _endpoint: &str,
        body: &serde_json::Value,
    ) -> Result<u32> {
        use tokio::io::AsyncBufReadExt;
        use tokio::io::BufReader;
        use tokio::net::TcpStream as TokioTcpStream;

        let stripped = base_url.strip_prefix("http://").ok_or_else(|| {
            Error::ConnectionError(
                "HTTPS requires the `schema-registry` feature; URL must start with http:// without it".into(),
            )
        })?;
        let (host_port, _) = stripped.split_once('/').unwrap_or((stripped, ""));

        // Extract path from endpoint (skip the base URL part)
        let path = _endpoint.strip_prefix(base_url).unwrap_or(_endpoint);

        let body_bytes = serde_json::to_vec(body)
            .map_err(|e| Error::ConnectionError(format!("failed to serialize schema: {e}")))?;

        let request = format!(
            "POST {} HTTP/1.1\r\nHost: {}\r\nContent-Type: application/vnd.schemaregistry.v1+json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
            path, host_port, body_bytes.len()
        );

        let timeout = tokio::time::Duration::from_secs(30);

        let mut stream = tokio::time::timeout(timeout, TokioTcpStream::connect(host_port))
            .await
            .map_err(|_| Error::ConnectionError("schema registry connect timed out".into()))?
            .map_err(|e| {
                Error::ConnectionError(format!("failed to connect to schema registry: {e}"))
            })?;

        stream
            .write_all(request.as_bytes())
            .await
            .map_err(|e| Error::ConnectionError(format!("failed to send request: {e}")))?;
        stream
            .write_all(&body_bytes)
            .await
            .map_err(|e| Error::ConnectionError(format!("failed to send body: {e}")))?;

        // Read response with timeout
        let response_body =
            tokio::time::timeout(timeout, async {
                let mut reader = BufReader::new(stream);

                // --- Parse status line ---
                let mut status_line = String::new();
                reader.read_line(&mut status_line).await.map_err(|e| {
                    Error::ConnectionError(format!("failed to read status line: {e}"))
                })?;

                let status_code: u16 = status_line
                    .split_whitespace()
                    .nth(1)
                    .and_then(|s| s.parse().ok())
                    .unwrap_or(0);

                if !(200..300).contains(&status_code) {
                    return Err(Error::ServerError(format!(
                        "schema registry returned HTTP {status_code}"
                    )));
                }

                // --- Parse headers ---
                let mut content_length: Option<usize> = None;
                let mut is_chunked = false;
                loop {
                    let mut header_line = String::new();
                    reader.read_line(&mut header_line).await.map_err(|e| {
                        Error::ConnectionError(format!("failed to read header: {e}"))
                    })?;

                    let trimmed = header_line.trim();
                    if trimmed.is_empty() {
                        break; // End of headers
                    }

                    let lower = trimmed.to_ascii_lowercase();
                    if let Some(val) = lower.strip_prefix("content-length:") {
                        content_length = val.trim().parse().ok();
                    } else if lower.starts_with("transfer-encoding:") && lower.contains("chunked") {
                        is_chunked = true;
                    }
                }

                // --- Read body ---
                let body_bytes = if is_chunked {
                    // Chunked transfer-encoding: read chunk-size\r\n, chunk-data\r\n, repeat until 0\r\n
                    const MAX_CHUNK_SIZE: usize = 16 * 1024 * 1024;
                    const MAX_TOTAL_BODY: usize = 16 * 1024 * 1024;
                    let mut body = Vec::new();
                    loop {
                        let mut size_line = String::new();
                        reader.read_line(&mut size_line).await.map_err(|e| {
                            Error::ConnectionError(format!("failed to read chunk size: {e}"))
                        })?;

                        let chunk_size =
                            usize::from_str_radix(size_line.trim(), 16).map_err(|_| {
                                Error::ConnectionError(format!(
                                    "invalid chunk size: {:?}",
                                    size_line.trim()
                                ))
                            })?;
                        if chunk_size == 0 {
                            // Terminal chunk — consume trailing CRLF per RFC 7230 §4.1
                            let mut trailer_buf = [0u8; 2];
                            let _ = reader.read_exact(&mut trailer_buf).await;
                            break;
                        }

                        // Guard against malicious chunk sizes
                        if chunk_size > MAX_CHUNK_SIZE {
                            return Err(Error::ConnectionError(format!(
                                "chunk size {} exceeds maximum {}",
                                chunk_size, MAX_CHUNK_SIZE
                            )));
                        }

                        let mut chunk = vec![0u8; chunk_size];
                        reader.read_exact(&mut chunk).await.map_err(|e| {
                            Error::ConnectionError(format!("failed to read chunk data: {e}"))
                        })?;
                        body.extend_from_slice(&chunk);

                        // Guard against unbounded total body accumulation
                        if body.len() > MAX_TOTAL_BODY {
                            return Err(Error::ConnectionError(format!(
                                "chunked body {} bytes exceeds maximum {}",
                                body.len(),
                                MAX_TOTAL_BODY
                            )));
                        }

                        // Consume trailing \r\n (exactly 2 bytes) after chunk data.
                        // Using read_exact instead of read_line prevents unbounded
                        // memory reads on malformed input.
                        let mut crlf_buf = [0u8; 2];
                        reader.read_exact(&mut crlf_buf).await.map_err(|e| {
                            Error::ConnectionError(format!("failed to read chunk CRLF: {e}"))
                        })?;
                        if crlf_buf != [b'\r', b'\n'] {
                            return Err(Error::ConnectionError(format!(
                                "expected CRLF after chunk data, got {:02x?}",
                                crlf_buf
                            )));
                        }
                    }
                    body
                } else if let Some(len) = content_length {
                    const MAX_CONTENT_LENGTH: usize = 16 * 1024 * 1024;
                    if len > MAX_CONTENT_LENGTH {
                        return Err(Error::ConnectionError(format!(
                            "response Content-Length {} bytes exceeds maximum {}",
                            len, MAX_CONTENT_LENGTH
                        )));
                    }
                    let mut body = vec![0u8; len];
                    reader.read_exact(&mut body).await.map_err(|e| {
                        Error::ConnectionError(format!("failed to read response body: {e}"))
                    })?;
                    body
                } else {
                    // Fallback: Connection: close — read until EOF (capped at 16 MB)
                    const MAX_RESPONSE_SIZE: usize = 16 * 1024 * 1024;
                    let mut body = Vec::with_capacity(4096);
                    reader.read_to_end(&mut body).await.map_err(|e| {
                        Error::ConnectionError(format!("failed to read response: {e}"))
                    })?;
                    if body.len() > MAX_RESPONSE_SIZE {
                        return Err(Error::ConnectionError(format!(
                            "response body {} bytes exceeds maximum {}",
                            body.len(),
                            MAX_RESPONSE_SIZE
                        )));
                    }
                    body
                };

                Ok(body_bytes)
            })
            .await
            .map_err(|_| Error::ConnectionError("schema registry response timed out".into()))??;

        #[derive(serde::Deserialize)]
        struct RegisterResponse {
            id: u32,
        }

        let result: RegisterResponse = serde_json::from_slice(&response_body).map_err(|e| {
            Error::ConnectionError(format!("failed to parse schema registry response: {e}"))
        })?;

        Ok(result.id)
    }

    /// List all consumer groups
    pub async fn list_groups(&mut self) -> Result<Vec<String>> {
        let request = Request::ListGroups;

        let response = self.send_request(request).await?;

        match response {
            Response::Groups { groups } => Ok(groups),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Describe a consumer group (get all committed offsets)
    pub async fn describe_group(
        &mut self,
        consumer_group: impl Into<String>,
    ) -> Result<std::collections::HashMap<String, std::collections::HashMap<u32, u64>>> {
        let request = Request::DescribeGroup {
            consumer_group: consumer_group.into(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::GroupDescription { offsets, .. } => Ok(offsets),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Delete a consumer group
    pub async fn delete_group(&mut self, consumer_group: impl Into<String>) -> Result<()> {
        let request = Request::DeleteGroup {
            consumer_group: consumer_group.into(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::GroupDeleted => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    // =========================================================================
    // Consumer Group Coordination
    // =========================================================================

    /// Join a consumer group.
    ///
    /// The coordinator assigns (or generates) a member ID and returns the
    /// current generation, leader, and the full member list. The leader
    /// uses the member list to compute partition assignments and sends
    /// them via [`sync_group`](Self::sync_group).
    ///
    /// # Returns
    /// `(generation_id, protocol_type, member_id, leader_id, members)`
    /// where `members` is `Vec<(member_id, subscriptions)>`.
    pub async fn join_group(
        &mut self,
        group_id: impl Into<String>,
        member_id: impl Into<String>,
        session_timeout_ms: u32,
        rebalance_timeout_ms: u32,
        protocol_type: impl Into<String>,
        subscriptions: Vec<String>,
    ) -> Result<(u32, String, String, String, Vec<(String, Vec<String>)>)> {
        let request = Request::JoinGroup {
            group_id: group_id.into(),
            member_id: member_id.into(),
            session_timeout_ms,
            rebalance_timeout_ms,
            protocol_type: protocol_type.into(),
            subscriptions,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::JoinGroupResult {
                generation_id,
                protocol_type,
                member_id,
                leader_id,
                members,
            } => Ok((generation_id, protocol_type, member_id, leader_id, members)),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Synchronize consumer group assignments.
    ///
    /// The group leader sends partition assignments for all members;
    /// followers send an empty assignment list. Every member receives
    /// their own assignment in the response.
    ///
    /// # Arguments
    /// * `assignments` — `Vec<(member_id, Vec<(topic, Vec<partition>)>)>`.
    ///   Only the leader should provide non-empty assignments.
    ///
    /// # Returns
    /// This member's partition assignment: `Vec<(topic, Vec<partition>)>`.
    pub async fn sync_group(
        &mut self,
        group_id: impl Into<String>,
        generation_id: u32,
        member_id: impl Into<String>,
        assignments: SyncGroupAssignments,
    ) -> Result<Vec<(String, Vec<u32>)>> {
        let request = Request::SyncGroup {
            group_id: group_id.into(),
            generation_id,
            member_id: member_id.into(),
            assignments,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::SyncGroupResult { assignments } => Ok(assignments),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Send a heartbeat to the consumer group coordinator.
    ///
    /// # Returns
    /// * `0` — OK, member is in sync
    /// * `27` — REBALANCE_IN_PROGRESS, member should rejoin
    pub async fn heartbeat(
        &mut self,
        group_id: impl Into<String>,
        generation_id: u32,
        member_id: impl Into<String>,
    ) -> Result<i32> {
        let request = Request::Heartbeat {
            group_id: group_id.into(),
            generation_id,
            member_id: member_id.into(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::HeartbeatResult { error_code } => Ok(error_code),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Leave a consumer group.
    ///
    /// Gracefully removes this member, triggering a rebalance for
    /// remaining group members. Call this during consumer shutdown.
    pub async fn leave_group(
        &mut self,
        group_id: impl Into<String>,
        member_id: impl Into<String>,
    ) -> Result<()> {
        let request = Request::LeaveGroup {
            group_id: group_id.into(),
            member_id: member_id.into(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::LeaveGroupResult => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Get the first offset with timestamp >= the given timestamp
    ///
    /// # Arguments
    /// * `topic` - The topic name
    /// * `partition` - The partition number
    /// * `timestamp_ms` - Timestamp in milliseconds since Unix epoch
    ///
    /// # Returns
    /// * `Some(offset)` - The first offset with message timestamp >= timestamp_ms
    /// * `None` - No messages found with timestamp >= timestamp_ms
    pub async fn get_offset_for_timestamp(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        timestamp_ms: i64,
    ) -> Result<Option<u64>> {
        let request = Request::GetOffsetForTimestamp {
            topic: topic.into(),
            partition,
            timestamp_ms,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::OffsetForTimestamp { offset } => Ok(offset),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    // ========================================================================
    // Admin API
    // ========================================================================

    /// Describe topic configurations
    ///
    /// Returns the current configuration for the specified topics.
    ///
    /// # Arguments
    /// * `topics` - Topics to describe (empty slice = all topics)
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// let configs = client.describe_topic_configs(&["orders", "events"]).await?;
    /// for (topic, config) in configs {
    ///     println!("{}: {:?}", topic, config);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn describe_topic_configs(
        &mut self,
        topics: &[&str],
    ) -> Result<std::collections::HashMap<String, std::collections::HashMap<String, String>>> {
        let request = Request::DescribeTopicConfigs {
            topics: topics.iter().map(|s| s.to_string()).collect(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TopicConfigsDescribed { configs } => {
                let mut result = std::collections::HashMap::new();
                for desc in configs {
                    let mut topic_configs = std::collections::HashMap::new();
                    for (key, value) in desc.configs {
                        topic_configs.insert(key, value.value);
                    }
                    result.insert(desc.topic, topic_configs);
                }
                Ok(result)
            }
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Alter topic configuration
    ///
    /// Modifies configuration for an existing topic. Pass `None` as value to reset
    /// a configuration key to its default.
    ///
    /// # Arguments
    /// * `topic` - Topic name
    /// * `configs` - Configuration changes: (key, value) pairs. Use `None` to reset to default.
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// let result = client.alter_topic_config("orders", &[
    ///     ("retention.ms", Some("86400000")),  // 1 day retention
    ///     ("cleanup.policy", Some("compact")), // Enable compaction
    /// ]).await?;
    /// println!("Changed {} configs", result.changed_count);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn alter_topic_config(
        &mut self,
        topic: impl Into<String>,
        configs: &[(&str, Option<&str>)],
    ) -> Result<AlterTopicConfigResult> {
        use rivven_protocol::TopicConfigEntry;

        let request = Request::AlterTopicConfig {
            topic: topic.into(),
            configs: configs
                .iter()
                .map(|(k, v)| TopicConfigEntry {
                    key: k.to_string(),
                    value: v.map(|s| s.to_string()),
                })
                .collect(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TopicConfigAltered {
                topic,
                changed_count,
            } => Ok(AlterTopicConfigResult {
                topic,
                changed_count,
            }),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Create additional partitions for an existing topic
    ///
    /// Increases the partition count for a topic. The new partition count
    /// must be greater than the current count (you cannot reduce partitions).
    ///
    /// # Arguments
    /// * `topic` - Topic name
    /// * `new_partition_count` - New total partition count
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// // Increase from 3 to 6 partitions
    /// let new_count = client.create_partitions("orders", 6).await?;
    /// println!("Topic now has {} partitions", new_count);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn create_partitions(
        &mut self,
        topic: impl Into<String>,
        new_partition_count: u32,
    ) -> Result<u32> {
        let request = Request::CreatePartitions {
            topic: topic.into(),
            new_partition_count,
            assignments: vec![], // Let broker auto-assign
        };

        let response = self.send_request(request).await?;

        match response {
            Response::PartitionsCreated {
                new_partition_count,
                ..
            } => Ok(new_partition_count),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Delete records before a given offset (log truncation)
    ///
    /// Removes all records with offsets less than the specified offset for each
    /// partition. This is useful for freeing up disk space or removing old data.
    ///
    /// # Arguments
    /// * `topic` - Topic name
    /// * `partition_offsets` - List of (partition, before_offset) pairs
    ///
    /// # Returns
    /// A list of results indicating the new low watermark for each partition.
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// // Delete records before offset 1000 on partitions 0, 1, 2
    /// let results = client.delete_records("orders", &[
    ///     (0, 1000),
    ///     (1, 1000),
    ///     (2, 1000),
    /// ]).await?;
    /// for r in results {
    ///     println!("Partition {}: low watermark now {}", r.partition, r.low_watermark);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn delete_records(
        &mut self,
        topic: impl Into<String>,
        partition_offsets: &[(u32, u64)],
    ) -> Result<Vec<DeleteRecordsResult>> {
        let request = Request::DeleteRecords {
            topic: topic.into(),
            partition_offsets: partition_offsets.to_vec(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::RecordsDeleted { results, .. } => Ok(results),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    // =========================================================================
    // Idempotent Producer API
    // =========================================================================

    /// Initialize an idempotent producer
    ///
    /// Returns a producer ID and epoch that should be used for all subsequent
    /// idempotent publish operations. The broker uses these to detect and
    /// deduplicate messages in case of retries.
    ///
    /// # Arguments
    /// * `previous_producer_id` - If reconnecting, pass the previous producer_id
    ///   to bump the epoch (prevents zombie producers)
    ///
    /// # Returns
    /// `ProducerState` containing the producer_id and producer_epoch
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// let producer = client.init_producer_id(None).await?;
    /// println!("Producer ID: {}, Epoch: {}", producer.producer_id, producer.producer_epoch);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn init_producer_id(
        &mut self,
        previous_producer_id: Option<u64>,
    ) -> Result<ProducerState> {
        let request = Request::InitProducerId {
            producer_id: previous_producer_id,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::ProducerIdInitialized {
                producer_id,
                producer_epoch,
            } => Ok(ProducerState {
                producer_id,
                producer_epoch,
                partition_sequences: std::collections::HashMap::new(),
                next_sequence: 0,
            }),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Publish a message with idempotent semantics
    ///
    /// Uses producer_id/epoch/sequence for exactly-once delivery. The broker
    /// deduplicates messages based on these values, making retries safe.
    ///
    /// # Arguments
    /// * `topic` - Topic to publish to
    /// * `key` - Optional message key (used for partitioning)
    /// * `value` - Message payload
    /// * `producer` - Producer state from `init_producer_id`
    ///
    /// # Returns
    /// Tuple of (offset, partition, was_duplicate)
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// let mut producer = client.init_producer_id(None).await?;
    ///
    /// let (offset, partition, duplicate) = client
    ///     .publish_idempotent("orders", None::<Vec<u8>>, b"order-1".to_vec(), &mut producer)
    ///     .await?;
    ///
    /// println!("Published to partition {} at offset {}", partition, offset);
    /// if duplicate {
    ///     println!("(This was a retry - message already existed)");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn publish_idempotent(
        &mut self,
        topic: impl Into<String>,
        key: Option<impl Into<Bytes>>,
        value: impl Into<Bytes>,
        producer: &mut ProducerState,
    ) -> Result<(u64, u32, bool)> {
        let topic_str = topic.into();
        // Use per-partition sequence when a partition is known; fall back to
        // the global counter for server-assigned partitions. The broker
        // deduplicates per (producer_id, partition, sequence).
        let sequence = producer.next_sequence;
        producer.next_sequence = producer.next_sequence.wrapping_add(1);
        // After wrapping past i32::MAX, reset to 1 (not 0)
        // because sequence 0 was used for the first message. Reusing 0
        // could collide with the broker's dedup window.
        if producer.next_sequence <= 0 {
            producer.next_sequence = 1;
        }

        let request = Request::IdempotentPublish {
            topic: topic_str,
            partition: None,
            key: key.map(|k| k.into()),
            value: value.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
            sequence,
            leader_epoch: None,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::IdempotentPublished {
                offset,
                partition,
                duplicate,
            } => Ok((offset, partition, duplicate)),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Publish a message idempotently to a **specific** partition.
    ///
    /// Uses the per-(topic, partition) sequence counter from `ProducerState`
    /// instead of the global counter, giving the broker correct dedup tracking
    /// per partition.
    pub async fn publish_idempotent_to_partition(
        &mut self,
        topic: impl Into<String>,
        partition: u32,
        key: Option<impl Into<Bytes>>,
        value: impl Into<Bytes>,
        producer: &mut ProducerState,
    ) -> Result<(u64, u32, bool)> {
        let topic_str = topic.into();
        let sequence = producer.next_sequence_for(&topic_str, partition);

        let request = Request::IdempotentPublish {
            topic: topic_str,
            partition: Some(partition),
            key: key.map(|k| k.into()),
            value: value.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
            sequence,
            leader_epoch: None,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::IdempotentPublished {
                offset,
                partition: resp_partition,
                duplicate,
            } => Ok((offset, resp_partition, duplicate)),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    // =========================================================================
    // Transaction API
    // =========================================================================

    /// Begin a new transaction
    ///
    /// Starts a transaction that can span multiple topics and partitions.
    /// All writes within the transaction are atomic - they either all succeed
    /// or all fail together.
    ///
    /// # Arguments
    /// * `txn_id` - Unique transaction identifier (should be stable per producer)
    /// * `producer` - Producer state from `init_producer_id`
    /// * `timeout_ms` - Optional transaction timeout (defaults to 60s)
    ///
    /// # Example
    /// ```no_run
    /// # use rivven_client::Client;
    /// # async fn example() -> rivven_client::Result<()> {
    /// let mut client = Client::connect("127.0.0.1:9092").await?;
    /// let producer = client.init_producer_id(None).await?;
    ///
    /// // Start a transaction
    /// client.begin_transaction("txn-1", &producer, None).await?;
    /// // ... publish messages ...
    /// client.commit_transaction("txn-1", &producer).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn begin_transaction(
        &mut self,
        txn_id: impl Into<String>,
        producer: &ProducerState,
        timeout_ms: Option<u64>,
    ) -> Result<()> {
        let request = Request::BeginTransaction {
            txn_id: txn_id.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
            timeout_ms,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TransactionStarted { .. } => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Add partitions to an active transaction
    ///
    /// Registers partitions that will be written to within the transaction.
    /// This must be called before publishing to a new partition.
    ///
    /// # Arguments
    /// * `txn_id` - Transaction identifier
    /// * `producer` - Producer state
    /// * `partitions` - List of (topic, partition) pairs to add
    pub async fn add_partitions_to_txn(
        &mut self,
        txn_id: impl Into<String>,
        producer: &ProducerState,
        partitions: &[(&str, u32)],
    ) -> Result<usize> {
        let request = Request::AddPartitionsToTxn {
            txn_id: txn_id.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
            partitions: partitions
                .iter()
                .map(|(t, p)| (t.to_string(), *p))
                .collect(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::PartitionsAddedToTxn {
                partition_count, ..
            } => Ok(partition_count),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Publish a message within a transaction
    ///
    /// Like `publish_idempotent`, but the message is only visible to consumers
    /// after the transaction is committed.
    ///
    /// # Arguments
    /// * `txn_id` - Transaction identifier
    /// * `topic` - Topic to publish to
    /// * `key` - Optional message key
    /// * `value` - Message payload
    /// * `producer` - Producer state with sequence tracking
    ///
    /// # Returns
    /// Tuple of (offset, partition, sequence) - offset is pending until commit
    pub async fn publish_transactional(
        &mut self,
        txn_id: impl Into<String>,
        topic: impl Into<String>,
        key: Option<impl Into<Bytes>>,
        value: impl Into<Bytes>,
        producer: &mut ProducerState,
    ) -> Result<(u64, u32, i32)> {
        let sequence = producer.next_sequence;
        producer.next_sequence = producer.next_sequence.wrapping_add(1);
        // Avoid reusing sequence 0 after wrap
        if producer.next_sequence <= 0 {
            producer.next_sequence = 1;
        }

        let request = Request::TransactionalPublish {
            txn_id: txn_id.into(),
            topic: topic.into(),
            partition: None,
            key: key.map(|k| k.into()),
            value: value.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
            sequence,
            leader_epoch: None,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TransactionalPublished {
                offset,
                partition,
                sequence,
            } => Ok((offset, partition, sequence)),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Add consumer offsets to a transaction
    ///
    /// For exactly-once consume-transform-produce patterns: commits consumer
    /// offsets atomically with the produced messages.
    ///
    /// # Arguments
    /// * `txn_id` - Transaction identifier
    /// * `producer` - Producer state
    /// * `group_id` - Consumer group ID
    /// * `offsets` - List of (topic, partition, offset) to commit
    pub async fn add_offsets_to_txn(
        &mut self,
        txn_id: impl Into<String>,
        producer: &ProducerState,
        group_id: impl Into<String>,
        offsets: &[(&str, u32, i64)],
    ) -> Result<()> {
        let request = Request::AddOffsetsToTxn {
            txn_id: txn_id.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
            group_id: group_id.into(),
            offsets: offsets
                .iter()
                .map(|(t, p, o)| (t.to_string(), *p, *o))
                .collect(),
        };

        let response = self.send_request(request).await?;

        match response {
            Response::OffsetsAddedToTxn { .. } => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Commit a transaction
    ///
    /// Makes all writes in the transaction visible to consumers atomically.
    /// If this fails, the transaction should be aborted.
    ///
    /// # Arguments
    /// * `txn_id` - Transaction identifier
    /// * `producer` - Producer state
    pub async fn commit_transaction(
        &mut self,
        txn_id: impl Into<String>,
        producer: &ProducerState,
    ) -> Result<()> {
        let request = Request::CommitTransaction {
            txn_id: txn_id.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TransactionCommitted { .. } => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }

    /// Abort a transaction
    ///
    /// Discards all writes in the transaction. Call this if any write fails
    /// or if you need to cancel the transaction.
    ///
    /// # Arguments
    /// * `txn_id` - Transaction identifier
    /// * `producer` - Producer state
    pub async fn abort_transaction(
        &mut self,
        txn_id: impl Into<String>,
        producer: &ProducerState,
    ) -> Result<()> {
        let request = Request::AbortTransaction {
            txn_id: txn_id.into(),
            producer_id: producer.producer_id,
            producer_epoch: producer.producer_epoch,
        };

        let response = self.send_request(request).await?;

        match response {
            Response::TransactionAborted { .. } => Ok(()),
            Response::Error { message } => Err(Error::ServerError(message)),
            _ => Err(Error::InvalidResponse),
        }
    }
}

/// State for an idempotent/transactional producer
#[derive(Debug, Clone)]
pub struct ProducerState {
    /// Producer ID assigned by the broker
    pub producer_id: u64,
    /// Current epoch (increments on reconnect)
    pub producer_epoch: u16,
    /// Per-partition sequence numbers for idempotent produce.
    /// When the client specifies a partition, the per-partition sequence
    /// is used for correct broker-side deduplication.
    /// Key: (topic, partition), Value: next sequence number.
    pub partition_sequences: std::collections::HashMap<(String, u32), i32>,
    /// Global sequence counter used when the partition is server-assigned
    /// (partition=None). The broker tracks this per producer_id.
    pub next_sequence: i32,
}

impl ProducerState {
    /// Get the next sequence number for a specific topic-partition,
    /// initializing to 1 if this is the first message to that partition.
    pub fn next_sequence_for(&mut self, topic: &str, partition: u32) -> i32 {
        let seq = self
            .partition_sequences
            .entry((topic.to_string(), partition))
            .or_insert(1);
        let current = *seq;
        *seq = seq.wrapping_add(1);
        if *seq <= 0 {
            *seq = 1;
        }
        current
    }
}

/// Result of altering topic configuration
#[derive(Debug, Clone)]
pub struct AlterTopicConfigResult {
    /// Topic name
    pub topic: String,
    /// Number of configurations changed
    pub changed_count: usize,
}

/// Result of deleting records from a partition
pub use rivven_protocol::DeleteRecordsResult;

// ============================================================================
// Authentication Session
// ============================================================================

/// Authentication session information
#[derive(Debug, Clone)]
pub struct AuthSession {
    /// Session ID for subsequent requests
    pub session_id: String,
    /// Session timeout in seconds
    pub expires_in: u64,
}

// ============================================================================
// SCRAM Helper Functions
// ============================================================================

/// Generate a random nonce for SCRAM authentication
pub(crate) fn generate_nonce() -> String {
    use rand::Rng;
    let mut rng = rand::thread_rng();
    let nonce_bytes: Vec<u8> = (0..24).map(|_| rng.gen()).collect();
    base64_encode(&nonce_bytes)
}

/// Escape username for SCRAM (RFC 5802)
pub(crate) fn escape_username(username: &str) -> String {
    username.replace('=', "=3D").replace(',', "=2C")
}

/// Parse server-first message
pub(crate) fn parse_server_first(server_first: &str) -> Result<(String, String, u32)> {
    let mut nonce = None;
    let mut salt = None;
    let mut iterations = None;

    for attr in server_first.split(',') {
        if let Some(value) = attr.strip_prefix("r=") {
            nonce = Some(value.to_string());
        } else if let Some(value) = attr.strip_prefix("s=") {
            salt = Some(value.to_string());
        } else if let Some(value) = attr.strip_prefix("i=") {
            iterations = Some(
                value
                    .parse::<u32>()
                    .map_err(|_| Error::AuthenticationFailed("Invalid iteration count".into()))?,
            );
        }
    }

    let nonce = nonce.ok_or_else(|| Error::AuthenticationFailed("Missing nonce".into()))?;
    let salt = salt.ok_or_else(|| Error::AuthenticationFailed("Missing salt".into()))?;
    let iterations =
        iterations.ok_or_else(|| Error::AuthenticationFailed("Missing iterations".into()))?;

    // RFC 7677 §4 requires a minimum of 4096 iterations for SCRAM-SHA-256.
    // A malicious server could send i=1 to weaken key derivation.
    if iterations < 4096 {
        return Err(Error::AuthenticationFailed(format!(
            "SCRAM iteration count {} is below minimum 4096 (possible downgrade attack)",
            iterations
        )));
    }

    Ok((nonce, salt, iterations))
}

/// PBKDF2-HMAC-SHA256 key derivation
pub(crate) fn pbkdf2_sha256(password: &[u8], salt: &[u8], iterations: u32) -> Vec<u8> {
    let mut result = vec![0u8; 32];

    // U1 = PRF(Password, Salt || INT(1))
    let mut u = PasswordHash::hmac_sha256(password, &[salt, &1u32.to_be_bytes()].concat());
    result.copy_from_slice(&u);

    // Ui = PRF(Password, Ui-1)
    for _ in 1..iterations {
        u = PasswordHash::hmac_sha256(password, &u);
        for (r, ui) in result.iter_mut().zip(u.iter()) {
            *r ^= ui;
        }
    }

    result
}

/// SHA-256 hash
pub(crate) fn sha256(data: &[u8]) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().to_vec()
}

/// XOR two byte arrays
pub(crate) fn xor_bytes(a: &[u8], b: &[u8]) -> Vec<u8> {
    a.iter().zip(b.iter()).map(|(x, y)| x ^ y).collect()
}

/// Base64 encode
pub(crate) fn base64_encode(data: &[u8]) -> String {
    use base64::{engine::general_purpose::STANDARD, Engine};
    STANDARD.encode(data)
}

/// Base64 decode
pub(crate) fn base64_decode(data: &str) -> std::result::Result<Vec<u8>, base64::DecodeError> {
    use base64::{engine::general_purpose::STANDARD, Engine};
    STANDARD.decode(data)
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_escape_username() {
        assert_eq!(escape_username("alice"), "alice");
        assert_eq!(escape_username("user=name"), "user=3Dname");
        assert_eq!(escape_username("user,name"), "user=2Cname");
        assert_eq!(escape_username("user=,name"), "user=3D=2Cname");
    }

    #[test]
    fn test_parse_server_first() {
        let server_first = "r=clientnonce+servernonce,s=c2FsdA==,i=4096";
        let (nonce, salt, iterations) = parse_server_first(server_first).unwrap();

        assert_eq!(nonce, "clientnonce+servernonce");
        assert_eq!(salt, "c2FsdA==");
        assert_eq!(iterations, 4096);
    }

    #[test]
    fn test_parse_server_first_missing_nonce() {
        let server_first = "s=c2FsdA==,i=4096";
        assert!(parse_server_first(server_first).is_err());
    }

    #[test]
    fn test_parse_server_first_missing_salt() {
        let server_first = "r=nonce,i=4096";
        assert!(parse_server_first(server_first).is_err());
    }

    #[test]
    fn test_parse_server_first_missing_iterations() {
        let server_first = "r=nonce,s=c2FsdA==";
        assert!(parse_server_first(server_first).is_err());
    }

    #[test]
    fn test_xor_bytes() {
        assert_eq!(xor_bytes(&[0xFF, 0x00], &[0xFF, 0xFF]), vec![0x00, 0xFF]);
        assert_eq!(xor_bytes(&[0x12, 0x34], &[0x12, 0x34]), vec![0x00, 0x00]);
    }

    #[test]
    fn test_base64_roundtrip() {
        let data = b"hello world";
        let encoded = base64_encode(data);
        let decoded = base64_decode(&encoded).unwrap();
        assert_eq!(decoded, data);
    }

    #[test]
    fn test_sha256() {
        // SHA-256 of empty string
        let hash = sha256(b"");
        assert_eq!(hash.len(), 32);
        // Known hash value
        assert_eq!(
            hex::encode(&hash),
            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
        );
    }

    #[test]
    fn test_pbkdf2_sha256() {
        // Test vector from RFC 7914 (derived from RFC 6070)
        let password = b"password";
        let salt = b"salt";
        let iterations = 1;

        let result = pbkdf2_sha256(password, salt, iterations);
        assert_eq!(result.len(), 32);
        // The result should be deterministic
        let result2 = pbkdf2_sha256(password, salt, iterations);
        assert_eq!(result, result2);
    }

    #[test]
    fn test_generate_nonce() {
        let nonce1 = generate_nonce();
        let nonce2 = generate_nonce();

        // Nonces should be non-empty
        assert!(!nonce1.is_empty());
        assert!(!nonce2.is_empty());

        // Nonces should be different (with overwhelming probability)
        assert_ne!(nonce1, nonce2);

        // Should be valid base64
        assert!(base64_decode(&nonce1).is_ok());
    }
}