velo 0.12.0

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Zero-copy TCP framing codec for ActiveMessage transport
//!
//! Wire format (7-15 bytes overhead):
//! ```text
//! [u16 BE: schema_version][u8: frame_type][u32 BE: header_len][u32 BE: payload_len][header bytes][payload bytes]
//! ```
//!
//! The codec uses `BytesMut` for receiving and `Bytes` for output, enabling
//! zero-copy buffer slicing where header and payload share the underlying buffer.

use bytes::{Buf, Bytes, BytesMut};
use std::io;
use std::io::Write;
use tokio::io::{AsyncWrite, AsyncWriteExt};
use tokio_util::codec::Decoder;

use velo_ext::MessageType;

/// Current schema version
const SCHEMA_VERSION_V1: u16 = 1;

/// Default maximum frame size (16 MB)
pub(crate) const DEFAULT_MAX_FRAME_SIZE: u32 = 16 * 1024 * 1024;

/// Minimum frame header size (version + type + 2 lengths)
pub(crate) const MIN_HEADER_SIZE: usize = 2 + 1 + 4 + 4; // 11 bytes

/// Threshold below which a frame is worth copying into a single buffer so it
/// can go out as one `write_all`. Above it we keep the three-segment path so
/// the extra memcpy doesn't dominate large payloads.
///
/// Also the point at which [`crate::transports::coalesce`] stops staging a
/// frame into a shared batch, for the same reason.
pub(crate) const COALESCE_THRESHOLD: usize = 64 * 1024;

/// Cap on the stack buffer that stages preamble + header ahead of an
/// above-threshold payload, so the direct path issues two writes instead of
/// three. Sized for the headers that actually cross this path — control
/// metadata of a few tens of bytes — with room to spare; a header that does
/// not fit falls back to the three-segment write, where its own size
/// amortizes the extra syscall.
pub(crate) const DIRECT_PREFIX_CAP: usize = 256;

/// Stage `preamble` and `header` contiguously into `prefix`, returning the
/// staged length, or `None` when they do not fit [`DIRECT_PREFIX_CAP`].
///
/// On a `TCP_NODELAY` socket the alternative — writing the 11-byte preamble
/// and the header as separate `write_all`s — costs two syscalls and can put
/// two undersized segments on the wire before the payload.
#[inline]
pub(crate) fn stage_direct_prefix(
    preamble: &[u8; MIN_HEADER_SIZE],
    header: &[u8],
    prefix: &mut [u8; DIRECT_PREFIX_CAP],
) -> Option<usize> {
    let len = MIN_HEADER_SIZE + header.len();
    if len > DIRECT_PREFIX_CAP {
        return None;
    }
    prefix[..MIN_HEADER_SIZE].copy_from_slice(preamble);
    prefix[MIN_HEADER_SIZE..len].copy_from_slice(header);
    Some(len)
}

/// Fallback shrink threshold (8 MB) when neither the builder nor an env var
/// supplies one. Shared across the TCP / UDS / streaming TCP code paths that
/// all decode through this codec via `tokio_util::codec::Framed`.
pub(crate) const DEFAULT_SHRINK_THRESHOLD: usize = 8 * 1024 * 1024;

/// Capacity a per-connection read buffer is reset to after exceeding
/// [`DEFAULT_SHRINK_THRESHOLD`] (or a transport-specific threshold) and
/// emptying. Sized so a few typical frames still fit without re-growing.
pub(crate) const SHRINK_RESET_CAPACITY: usize = 256 * 1024;

/// Pure parser for shrink-threshold env values. Returns
/// [`DEFAULT_SHRINK_THRESHOLD`] when input is None, empty, or unparseable.
/// Shared by every transport that wires an env override.
pub(crate) fn parse_shrink_threshold(raw: Option<&str>) -> usize {
    raw.and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(DEFAULT_SHRINK_THRESHOLD)
}

/// Reclaim memory if a one-time large frame inflated the read buffer.
///
/// Replaces `buf` with a fresh [`SHRINK_RESET_CAPACITY`]-sized buffer when it
/// has emptied after a one-time growth. Called after each successfully decoded
/// frame with the read buffer of a `Framed`/`FramedRead` running
/// [`TcpFrameCodec`] — TCP, UDS, the dialed-connection readers, and the
/// streaming TCP pump all share this codec and the same buffer-growth pattern.
///
/// `last_frame_total_size` is the header + payload byte count of the frame
/// just decoded (the preamble's 11 bytes are negligible and can be omitted).
/// It's used to skip reclamation under sustained large traffic where shrinking
/// would just be undone by the next frame.
///
/// The four conditions together prevent per-frame reallocation churn:
///   1. `buf.is_empty()` — never lose bytes belonging to the next frame.
///   2. `buf.capacity() > threshold` — user has opted in to reclaiming this
///      buffer size.
///   3. `buf.capacity() > 2 * SHRINK_RESET_CAPACITY` — the realloc actually
///      saves meaningful memory. Stops a small user-configured `threshold`
///      from triggering pointless 256 KB → 256 KB reallocations.
///   4. `last_frame_total_size * 2 <= buf.capacity()` — the most recent
///      frame was small enough relative to current capacity that resetting
///      won't be immediately undone by the next frame. Avoids churn when a
///      caller sustains large frames above the threshold.
#[inline]
pub(crate) fn maybe_shrink_read_buffer(
    buf: &mut BytesMut,
    threshold: usize,
    last_frame_total_size: usize,
) {
    if buf.is_empty()
        && buf.capacity() > threshold
        && buf.capacity() > 2 * SHRINK_RESET_CAPACITY
        && last_frame_total_size.saturating_mul(2) <= buf.capacity()
    {
        *buf = BytesMut::with_capacity(SHRINK_RESET_CAPACITY);
    }
}

/// Zero-copy frame decoder for TCP transport
///
/// This decoder maintains state across multiple calls to support partial
/// frame reception. It decodes frames into (MessageType, header: Bytes, payload: Bytes)
/// where header and payload are zero-copy slices of the receive buffer.
#[derive(Debug, Clone)]
pub struct TcpFrameCodec {
    state: DecodeState,
    max_frame_size: u32,
}

#[derive(Debug, Clone, Copy)]
enum DecodeState {
    /// Waiting for frame header (version + type + lengths)
    AwaitingHeader,
    /// Waiting for frame data (header + payload), with known lengths
    AwaitingData {
        frame_type: MessageType,
        header_len: u32,
        payload_len: u32,
    },
}

impl TcpFrameCodec {
    /// Create a new frame codec with the default max frame size (16 MB)
    pub fn new() -> Self {
        Self {
            state: DecodeState::AwaitingHeader,
            max_frame_size: DEFAULT_MAX_FRAME_SIZE,
        }
    }

    /// Create a new frame codec with a custom max frame size
    pub fn with_max_frame_size(max_frame_size: u32) -> Self {
        Self {
            state: DecodeState::AwaitingHeader,
            max_frame_size,
        }
    }

    /// Build the frame preamble (metadata header)
    ///
    /// Returns a fixed-size preamble containing version, message type, and lengths.
    #[inline]
    pub fn build_preamble(
        msg_type: MessageType,
        header_len: u32,
        payload_len: u32,
    ) -> io::Result<[u8; MIN_HEADER_SIZE]> {
        // Validate lengths before building preamble
        Self::validate_lengths(header_len, payload_len)?;

        let mut preamble = [0u8; MIN_HEADER_SIZE];

        // Layout:
        // [0..2) = version
        // [2]    = msg_type
        // [3..7) = header_len
        // [7..11)= payload_len  (total 11 bytes)
        preamble[0..2].copy_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
        preamble[2] = msg_type.as_u8();
        preamble[3..7].copy_from_slice(&header_len.to_be_bytes());
        preamble[7..11].copy_from_slice(&payload_len.to_be_bytes());

        Ok(preamble)
    }

    /// Parse message type from a preamble
    ///
    /// Validates the schema version and extracts the message type from the preamble.
    #[inline]
    pub fn parse_message_type_from_preamble(preamble: &[u8]) -> io::Result<MessageType> {
        if preamble.len() < MIN_HEADER_SIZE {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Preamble too short",
            ));
        }

        // Validate schema version
        let schema_version = u16::from_be_bytes([preamble[0], preamble[1]]);
        if schema_version != SCHEMA_VERSION_V1 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!(
                    "Unsupported schema version: {} (expected {})",
                    schema_version, SCHEMA_VERSION_V1
                ),
            ));
        }

        // Extract and validate message type
        MessageType::from_u8(preamble[2]).ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Invalid message type: {}", preamble[2]),
            )
        })
    }

    /// Encode and write a frame asynchronously.
    ///
    /// Small frames (header + payload <= COALESCE_THRESHOLD) are packed into a
    /// single buffer and written with one `write_all` — 3× fewer syscalls per
    /// frame on the hot path for typical control/event traffic. Above the
    /// threshold the payload is never copied: preamble + header are staged
    /// into one stack buffer and written ahead of it (two writes), falling
    /// back to three segments only for a header too large for
    /// [`DIRECT_PREFIX_CAP`]. `write_all` is used throughout because TCP
    /// `write_vectored()` is allowed to short-write past ~128KB.
    #[inline]
    pub async fn encode_frame<W: AsyncWrite + Unpin>(
        writer: &mut W,
        msg_type: MessageType,
        header: &[u8],
        payload: &[u8],
    ) -> tokio::io::Result<()> {
        let preamble = Self::build_preamble(msg_type, header.len() as u32, payload.len() as u32)?;
        if header.len() + payload.len() <= COALESCE_THRESHOLD {
            let mut buf = BytesMut::with_capacity(MIN_HEADER_SIZE + header.len() + payload.len());
            buf.extend_from_slice(&preamble);
            buf.extend_from_slice(header);
            buf.extend_from_slice(payload);
            writer.write_all(&buf).await?;
        } else {
            let mut prefix = [0u8; DIRECT_PREFIX_CAP];
            match stage_direct_prefix(&preamble, header, &mut prefix) {
                Some(len) => {
                    writer.write_all(&prefix[..len]).await?;
                    writer.write_all(payload).await?;
                }
                None => {
                    writer.write_all(&preamble).await?;
                    writer.write_all(header).await?;
                    writer.write_all(payload).await?;
                }
            }
        }
        Ok(())
    }

    /// Append one encoded frame to `buf` without writing anything.
    ///
    /// The bytes produced are byte-identical to what [`encode_frame`] would
    /// write, so a peer decoding with [`TcpFrameCodec`] cannot tell whether a
    /// frame arrived alone or packed alongside others. This is what makes
    /// write coalescing wire-compatible in both directions — see
    /// `transports::coalesce`.
    #[inline]
    pub(crate) fn append_frame(
        buf: &mut BytesMut,
        msg_type: MessageType,
        header: &[u8],
        payload: &[u8],
    ) -> io::Result<()> {
        let preamble = Self::build_preamble(msg_type, header.len() as u32, payload.len() as u32)?;
        buf.reserve(MIN_HEADER_SIZE + header.len() + payload.len());
        buf.extend_from_slice(&preamble);
        buf.extend_from_slice(header);
        buf.extend_from_slice(payload);
        Ok(())
    }

    /// Encode and write a frame synchronously. See [`encode_frame`] for the
    /// coalescing rationale.
    #[inline]
    pub fn encode_frame_sync<W: Write>(
        writer: &mut W,
        msg_type: MessageType,
        header: &[u8],
        payload: &[u8],
    ) -> std::io::Result<()> {
        let preamble = Self::build_preamble(msg_type, header.len() as u32, payload.len() as u32)?;
        if header.len() + payload.len() <= COALESCE_THRESHOLD {
            let mut buf = BytesMut::with_capacity(MIN_HEADER_SIZE + header.len() + payload.len());
            buf.extend_from_slice(&preamble);
            buf.extend_from_slice(header);
            buf.extend_from_slice(payload);
            writer.write_all(&buf)?;
        } else {
            let mut prefix = [0u8; DIRECT_PREFIX_CAP];
            match stage_direct_prefix(&preamble, header, &mut prefix) {
                Some(len) => {
                    writer.write_all(&prefix[..len])?;
                    writer.write_all(payload)?;
                }
                None => {
                    writer.write_all(&preamble)?;
                    writer.write_all(header)?;
                    writer.write_all(payload)?;
                }
            }
        }
        Ok(())
    }

    /// Validate that lengths are reasonable (uses default max frame size).
    ///
    /// Used by static encode methods for backward compatibility.
    fn validate_lengths(header_len: u32, payload_len: u32) -> io::Result<()> {
        Self::validate_lengths_limit(header_len, payload_len, DEFAULT_MAX_FRAME_SIZE)
    }

    /// Validate that lengths are reasonable with a configurable max frame size.
    fn validate_lengths_limit(
        header_len: u32,
        payload_len: u32,
        max_frame_size: u32,
    ) -> io::Result<()> {
        let total_len = header_len
            .checked_add(payload_len)
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Frame size overflow"))?;

        if total_len > max_frame_size {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!(
                    "Frame size {} exceeds maximum {}",
                    total_len, max_frame_size
                ),
            ));
        }

        Ok(())
    }
}

impl Default for TcpFrameCodec {
    fn default() -> Self {
        Self::new()
    }
}

impl Decoder for TcpFrameCodec {
    type Item = (MessageType, Bytes, Bytes);
    type Error = io::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        loop {
            match self.state {
                DecodeState::AwaitingHeader => {
                    // Need at least MIN_HEADER_SIZE bytes
                    if src.len() < MIN_HEADER_SIZE {
                        return Ok(None);
                    }

                    // Parse header without consuming bytes yet
                    let schema_version = u16::from_be_bytes([src[0], src[1]]);
                    let frame_type_byte = src[2];
                    let header_len = u32::from_be_bytes([src[3], src[4], src[5], src[6]]);
                    let payload_len = u32::from_be_bytes([src[7], src[8], src[9], src[10]]);

                    // Validate schema version
                    if schema_version != SCHEMA_VERSION_V1 {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            format!(
                                "Unsupported schema version: {} (expected {})",
                                schema_version, SCHEMA_VERSION_V1
                            ),
                        ));
                    }

                    // Parse frame type
                    let frame_type = MessageType::from_u8(frame_type_byte).ok_or_else(|| {
                        io::Error::new(
                            io::ErrorKind::InvalidData,
                            format!("Invalid frame type: {}", frame_type_byte),
                        )
                    })?;

                    // Validate lengths before allocating/waiting
                    Self::validate_lengths_limit(header_len, payload_len, self.max_frame_size)?;

                    // Advance buffer past header
                    src.advance(MIN_HEADER_SIZE);

                    // Transition to data state
                    self.state = DecodeState::AwaitingData {
                        frame_type,
                        header_len,
                        payload_len,
                    };
                }

                DecodeState::AwaitingData {
                    frame_type,
                    header_len,
                    payload_len,
                    ..
                } => {
                    let total_data_len = (header_len + payload_len) as usize;

                    // Wait for full data. Reserve the whole remainder up
                    // front: `Framed` itself only ever grows the buffer by
                    // doubling from 8 KB, which for a large frame means a
                    // realloc + memcpy of everything received so far at every
                    // step. One reservation sized from the preamble replaces
                    // that with a single realloc (and lets each read pull as
                    // much as the kernel has buffered).
                    if src.len() < total_data_len {
                        src.reserve(total_data_len - src.len());
                        return Ok(None);
                    }

                    // Zero-copy: split buffer into header and payload slices
                    let header = src.split_to(header_len as usize).freeze();
                    let payload = src.split_to(payload_len as usize).freeze();

                    // Reset state for next frame
                    self.state = DecodeState::AwaitingHeader;

                    return Ok(Some((frame_type, header, payload)));
                }
            }
        }
    }
}

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

    /// Test helper to encode a frame into a Vec<u8> for verification (async)
    async fn encode_frame_to_bytes(
        msg_type: MessageType,
        header: &[u8],
        payload: &[u8],
    ) -> io::Result<Vec<u8>> {
        let mut buf = Vec::new();
        TcpFrameCodec::encode_frame(&mut buf, msg_type, header, payload).await?;
        Ok(buf)
    }

    /// Test helper to encode a frame into a Vec<u8> for verification (sync)
    fn encode_frame_to_bytes_sync(
        msg_type: MessageType,
        header: &[u8],
        payload: &[u8],
    ) -> io::Result<Vec<u8>> {
        let mut buf = Vec::new();
        TcpFrameCodec::encode_frame_sync(&mut buf, msg_type, header, payload)?;
        Ok(buf)
    }

    /// Helper to create raw frames with arbitrary parameters for negative testing.
    ///
    /// This function bypasses normal validation and encoding logic to create
    /// intentionally invalid frames (wrong schema version, oversized frames, etc.)
    /// for testing error handling paths. Use `encode_frame_to_bytes()` for
    /// testing valid frame construction.
    fn create_unsafe_frame(
        schema_version: u16,
        frame_type: MessageType,
        header: &[u8],
        payload: &[u8],
    ) -> BytesMut {
        let mut buf = BytesMut::new();
        buf.extend_from_slice(&schema_version.to_be_bytes());
        buf.extend_from_slice(&[frame_type.as_u8()]);
        buf.extend_from_slice(&(header.len() as u32).to_be_bytes());
        buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
        buf.extend_from_slice(header);
        buf.extend_from_slice(payload);
        buf
    }

    #[test]
    fn test_decode_message_frame() {
        let mut codec = TcpFrameCodec::new();
        let header = b"test-header";
        let payload = b"test-payload-data";

        let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
        let mut buf = BytesMut::from(&framed[..]);

        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_some());

        let (msg_type, decoded_header, decoded_payload) = result.unwrap();
        assert_eq!(msg_type, MessageType::Message);
        assert_eq!(decoded_header, Bytes::from(header.as_ref()));
        assert_eq!(decoded_payload, Bytes::from(payload.as_ref()));
    }

    #[test]
    fn test_decode_all_frame_types() {
        let frame_types = [
            MessageType::Message,
            MessageType::Response,
            MessageType::Ack,
            MessageType::Event,
        ];

        for frame_type in &frame_types {
            let mut codec = TcpFrameCodec::new();
            let header = b"header";
            let payload = b"payload";

            let framed = encode_frame_to_bytes_sync(*frame_type, header, payload).unwrap();
            let mut buf = BytesMut::from(&framed[..]);

            let result = codec.decode(&mut buf).unwrap();
            assert!(result.is_some());

            let (decoded_type, _, _) = result.unwrap();
            assert_eq!(decoded_type, *frame_type);
        }
    }

    #[test]
    fn test_decode_empty_payload() {
        let mut codec = TcpFrameCodec::new();
        let header = b"ack-header";
        let payload = b"";

        let framed = encode_frame_to_bytes_sync(MessageType::Ack, header, payload).unwrap();
        let mut buf = BytesMut::from(&framed[..]);

        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_some());

        let (msg_type, decoded_header, decoded_payload) = result.unwrap();
        assert_eq!(msg_type, MessageType::Ack);
        assert_eq!(&decoded_header[..], header);
        assert_eq!(decoded_payload.len(), 0);
    }

    #[test]
    fn test_decode_partial_frame() {
        let mut codec = TcpFrameCodec::new();
        let header = b"test-header";
        let payload = b"test-payload";

        let full_frame = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();

        // Send only first 5 bytes (partial header)
        let mut buf = BytesMut::from(&full_frame[..5]);
        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_none()); // Not enough data

        // Send rest of header
        buf.extend_from_slice(&full_frame[5..MIN_HEADER_SIZE]);
        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_none()); // Header parsed, but data not yet available

        // Send complete data
        buf.extend_from_slice(&full_frame[MIN_HEADER_SIZE..]);
        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_some());

        let (msg_type, decoded_header, decoded_payload) = result.unwrap();
        assert_eq!(msg_type, MessageType::Message);
        assert_eq!(&decoded_header[..], header);
        assert_eq!(&decoded_payload[..], payload);
    }

    #[test]
    fn test_decode_invalid_schema_version() {
        let mut codec = TcpFrameCodec::new();
        let header = b"header";
        let payload = b"payload";

        let mut buf = create_unsafe_frame(999, MessageType::Message, header, payload);

        let result = codec.decode(&mut buf);
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("Unsupported schema version")
        );
    }

    #[test]
    fn test_decode_invalid_frame_type() {
        let mut codec = TcpFrameCodec::new();
        let mut buf = BytesMut::new();

        // Create frame with invalid type byte (255)
        buf.extend_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
        buf.extend_from_slice(&[255u8]); // Invalid frame type
        buf.extend_from_slice(&10u32.to_be_bytes()); // header len
        buf.extend_from_slice(&10u32.to_be_bytes()); // payload len

        let result = codec.decode(&mut buf);
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("Invalid frame type")
        );
    }

    #[test]
    fn test_decode_frame_too_large() {
        let mut codec = TcpFrameCodec::new();
        let mut buf = BytesMut::new();

        // Create frame that exceeds DEFAULT_MAX_FRAME_SIZE
        buf.extend_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
        buf.extend_from_slice(&[MessageType::Message.as_u8()]);
        buf.extend_from_slice(&(DEFAULT_MAX_FRAME_SIZE / 2 + 1).to_be_bytes());
        buf.extend_from_slice(&(DEFAULT_MAX_FRAME_SIZE / 2 + 1).to_be_bytes());

        let result = codec.decode(&mut buf);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("exceeds maximum"));
    }

    #[test]
    fn test_decode_multiple_frames() {
        let mut codec = TcpFrameCodec::new();
        let mut buf = BytesMut::new();

        // Add two frames to buffer
        let frame1 =
            encode_frame_to_bytes_sync(MessageType::Message, b"header1", b"payload1").unwrap();
        let frame2 =
            encode_frame_to_bytes_sync(MessageType::Response, b"header2", b"payload2").unwrap();
        buf.extend_from_slice(&frame1);
        buf.extend_from_slice(&frame2);

        // Decode first frame
        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_some());
        let (msg_type, header, payload) = result.unwrap();
        assert_eq!(msg_type, MessageType::Message);
        assert_eq!(&header[..], b"header1");
        assert_eq!(&payload[..], b"payload1");

        // Decode second frame
        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_some());
        let (msg_type, header, payload) = result.unwrap();
        assert_eq!(msg_type, MessageType::Response);
        assert_eq!(&header[..], b"header2");
        assert_eq!(&payload[..], b"payload2");

        // No more frames
        assert!(buf.is_empty());
    }

    #[test]
    fn test_zero_copy_bytes_share_buffer() {
        let mut codec = TcpFrameCodec::new();
        let header = b"shared-header";
        let payload = b"shared-payload";

        let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
        let mut buf = BytesMut::from(&framed[..]);

        let result = codec.decode(&mut buf).unwrap().unwrap();
        let (_, decoded_header, decoded_payload) = result;

        // Verify the slices contain correct data
        assert_eq!(&decoded_header[..], header);
        assert_eq!(&decoded_payload[..], payload);

        // Clone should be cheap (just RC increment)
        let header_clone = decoded_header.clone();
        let payload_clone = decoded_payload.clone();

        assert_eq!(decoded_header, header_clone);
        assert_eq!(decoded_payload, payload_clone);
    }

    #[test]
    fn test_encode_frame() {
        let header = b"test-header";
        let payload = b"test-payload";

        let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();

        // Verify frame structure
        assert_eq!(framed.len(), MIN_HEADER_SIZE + header.len() + payload.len());

        // Verify header fields
        assert_eq!(
            u16::from_be_bytes([framed[0], framed[1]]),
            SCHEMA_VERSION_V1
        );
        assert_eq!(framed[2], MessageType::Message.as_u8());
        assert_eq!(
            u32::from_be_bytes([framed[3], framed[4], framed[5], framed[6]]),
            header.len() as u32
        );
        assert_eq!(
            u32::from_be_bytes([framed[7], framed[8], framed[9], framed[10]]),
            payload.len() as u32
        );

        // Verify data
        assert_eq!(
            &framed[MIN_HEADER_SIZE..MIN_HEADER_SIZE + header.len()],
            header
        );
        assert_eq!(&framed[MIN_HEADER_SIZE + header.len()..], payload);
    }

    #[test]
    fn test_encode_all_message_types() {
        let header = b"header";
        let payload = b"payload";

        for msg_type in &[
            MessageType::Message,
            MessageType::Response,
            MessageType::Ack,
            MessageType::Event,
        ] {
            let framed = encode_frame_to_bytes_sync(*msg_type, header, payload).unwrap();
            assert_eq!(framed[2], msg_type.as_u8());
        }
    }

    #[test]
    fn test_encode_empty_payload() {
        let header = b"ack-header";
        let payload = b"";

        let framed = encode_frame_to_bytes_sync(MessageType::Ack, header, payload).unwrap();

        assert_eq!(framed.len(), MIN_HEADER_SIZE + header.len());
        assert_eq!(
            u32::from_be_bytes([framed[7], framed[8], framed[9], framed[10]]),
            0
        );
    }

    #[test]
    fn test_encode_frame_too_large() {
        let header = vec![0u8; (DEFAULT_MAX_FRAME_SIZE / 2 + 1) as usize];
        let payload = vec![0u8; (DEFAULT_MAX_FRAME_SIZE / 2 + 1) as usize];

        let result = encode_frame_to_bytes_sync(MessageType::Message, &header, &payload);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("exceeds maximum"));
    }

    #[test]
    fn test_round_trip_encode_decode() {
        let mut codec = TcpFrameCodec::new();
        let header = b"round-trip-header";
        let payload = b"round-trip-payload-data";

        // Encode
        let framed = encode_frame_to_bytes_sync(MessageType::Response, header, payload).unwrap();

        // Decode
        let mut buf = BytesMut::from(&framed[..]);
        let result = codec.decode(&mut buf).unwrap();
        assert!(result.is_some());

        let (msg_type, decoded_header, decoded_payload) = result.unwrap();
        assert_eq!(msg_type, MessageType::Response);
        assert_eq!(&decoded_header[..], header);
        assert_eq!(&decoded_payload[..], payload);
    }

    #[test]
    fn test_round_trip_all_types() {
        let types = [
            MessageType::Message,
            MessageType::Response,
            MessageType::Ack,
            MessageType::Event,
        ];

        for msg_type in &types {
            let mut codec = TcpFrameCodec::new();
            let header = b"header";
            let payload = b"payload";

            let framed = encode_frame_to_bytes_sync(*msg_type, header, payload).unwrap();

            let mut buf = BytesMut::from(&framed[..]);
            let result = codec.decode(&mut buf).unwrap().unwrap();

            assert_eq!(result.0, *msg_type);
            assert_eq!(&result.1[..], header);
            assert_eq!(&result.2[..], payload);
        }
    }

    #[test]
    fn test_encode_frame_sync() {
        let header = b"sync-header";
        let payload = b"sync-payload";

        let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();

        // Verify frame structure
        assert_eq!(framed.len(), MIN_HEADER_SIZE + header.len() + payload.len());

        // Verify preamble fields
        assert_eq!(
            u16::from_be_bytes([framed[0], framed[1]]),
            SCHEMA_VERSION_V1
        );
        assert_eq!(framed[2], MessageType::Message.as_u8());
        assert_eq!(
            u32::from_be_bytes([framed[3], framed[4], framed[5], framed[6]]),
            header.len() as u32
        );
        assert_eq!(
            u32::from_be_bytes([framed[7], framed[8], framed[9], framed[10]]),
            payload.len() as u32
        );

        // Verify data
        assert_eq!(
            &framed[MIN_HEADER_SIZE..MIN_HEADER_SIZE + header.len()],
            header
        );
        assert_eq!(&framed[MIN_HEADER_SIZE + header.len()..], payload);
    }

    #[test]
    fn test_sync_async_produce_same_output() {
        let header = b"test-header";
        let payload = b"test-payload";

        // Encode with sync version
        let sync_framed =
            encode_frame_to_bytes_sync(MessageType::Response, header, payload).unwrap();

        // Encode with async version (using tokio runtime)
        let async_framed = tokio::runtime::Runtime::new()
            .unwrap()
            .block_on(encode_frame_to_bytes(
                MessageType::Response,
                header,
                payload,
            ))
            .unwrap();

        // Both should produce identical output
        assert_eq!(sync_framed, async_framed);
    }

    /// Helper: build a Framed with a manually inflated read buffer to simulate
    /// a one-time large frame having grown the buffer past `capacity_bytes`.
    fn framed_with_grown_buffer(
        capacity_bytes: usize,
    ) -> Framed<tokio::io::DuplexStream, TcpFrameCodec> {
        let (read_half, _write_half) = tokio::io::duplex(64);
        let mut framed = Framed::new(read_half, TcpFrameCodec::new());
        framed.read_buffer_mut().resize(capacity_bytes, 0);
        framed.read_buffer_mut().clear(); // len=0, capacity retained
        framed
    }

    #[tokio::test]
    async fn test_maybe_shrink_resets_oversized_empty_buffer() {
        let mut framed = framed_with_grown_buffer(2 * 1024 * 1024);
        let pre_capacity = framed.read_buffer_mut().capacity();
        assert!(pre_capacity > 2 * SHRINK_RESET_CAPACITY);
        // last_frame was small relative to capacity → shrink applies.
        maybe_shrink_read_buffer(framed.read_buffer_mut(), 1024 * 1024, 1024);
        let post_capacity = framed.read_buffer_mut().capacity();
        assert!(
            post_capacity < pre_capacity,
            "expected capacity to drop: pre={} post={}",
            pre_capacity,
            post_capacity
        );
        assert_eq!(post_capacity, SHRINK_RESET_CAPACITY);
    }

    #[tokio::test]
    async fn test_maybe_shrink_skips_when_not_empty() {
        // The helper must not discard in-flight bytes belonging to the next frame.
        let (read_half, _write_half) = tokio::io::duplex(64);
        let mut framed = Framed::new(read_half, TcpFrameCodec::new());
        framed.read_buffer_mut().resize(2 * 1024 * 1024, 0xAB);
        let pre_capacity = framed.read_buffer_mut().capacity();
        maybe_shrink_read_buffer(framed.read_buffer_mut(), 1024 * 1024, 1024);
        assert_eq!(framed.read_buffer_mut().capacity(), pre_capacity);
    }

    #[tokio::test]
    async fn test_maybe_shrink_skips_under_threshold() {
        let (read_half, _write_half) = tokio::io::duplex(64);
        let mut framed = Framed::new(read_half, TcpFrameCodec::new());
        // Default Framed buffer capacity is well under DEFAULT_SHRINK_THRESHOLD.
        let pre_capacity = framed.read_buffer_mut().capacity();
        maybe_shrink_read_buffer(framed.read_buffer_mut(), DEFAULT_SHRINK_THRESHOLD, 1024);
        assert_eq!(framed.read_buffer_mut().capacity(), pre_capacity);
    }

    /// Sustained large traffic: when the most recent frame is comparable in
    /// size to current capacity, shrinking would just be undone by the next
    /// frame. The helper must skip to avoid per-frame realloc churn.
    #[tokio::test]
    async fn test_maybe_shrink_skips_under_sustained_large_frames() {
        let cap = 2 * 1024 * 1024;
        let mut framed = framed_with_grown_buffer(cap);
        let pre = framed.read_buffer_mut().capacity();
        // last_frame_size * 2 > capacity → skip (don't churn).
        maybe_shrink_read_buffer(framed.read_buffer_mut(), 1024 * 1024, cap * 3 / 4);
        assert_eq!(framed.read_buffer_mut().capacity(), pre);
    }

    /// Small user-configured threshold must not provoke reallocations when
    /// the buffer is already close to (or below) the reset target. Without
    /// the `capacity > 2 * SHRINK_RESET_CAPACITY` gate the helper would
    /// repeatedly replace a 256 KB buffer with another 256 KB buffer.
    #[tokio::test]
    async fn test_maybe_shrink_skips_when_capacity_not_meaningfully_oversized() {
        // Buffer slightly above SHRINK_RESET_CAPACITY but below 2× — savings
        // are negligible and re-growth on the next frame is likely.
        let cap = SHRINK_RESET_CAPACITY + 4096;
        let mut framed = framed_with_grown_buffer(cap);
        let pre = framed.read_buffer_mut().capacity();
        maybe_shrink_read_buffer(framed.read_buffer_mut(), 1024, 100); // small threshold, tiny frame
        assert_eq!(framed.read_buffer_mut().capacity(), pre);
    }

    #[test]
    fn test_parse_shrink_threshold() {
        // Pure parser exercised without touching the process env (which is
        // not thread-safe and would race with parallel test execution).
        assert_eq!(parse_shrink_threshold(Some("12345")), 12345);
        assert_eq!(parse_shrink_threshold(Some("0")), 0);
        assert_eq!(
            parse_shrink_threshold(Some("not-a-number")),
            DEFAULT_SHRINK_THRESHOLD
        );
        assert_eq!(parse_shrink_threshold(Some("")), DEFAULT_SHRINK_THRESHOLD);
        assert_eq!(parse_shrink_threshold(None), DEFAULT_SHRINK_THRESHOLD);
    }

    /// The coalesced fast path (small frames) and the three-segment slow path
    /// (large frames) must produce byte-identical wire output, and decode back
    /// to the same (type, header, payload) tuple.
    #[test]
    fn test_coalesce_and_segmented_paths_agree() {
        // Sizes chosen to straddle COALESCE_THRESHOLD (64 KB):
        //   below: 100 B  → fast path
        //   at:    threshold - HEADER_SIZE → fast path edge
        //   above: threshold + 1 KB → slow path
        let sizes = [
            100usize,
            COALESCE_THRESHOLD - 16,
            COALESCE_THRESHOLD,
            COALESCE_THRESHOLD + 1024,
        ];
        for total in sizes {
            // Split total roughly half header / half payload.
            let header_len = total / 2;
            let payload_len = total - header_len;
            let header: Vec<u8> = (0..header_len).map(|i| (i % 251) as u8).collect();
            let payload: Vec<u8> = (0..payload_len).map(|i| (i % 253) as u8).collect();

            let framed =
                encode_frame_to_bytes_sync(MessageType::Message, &header, &payload).unwrap();
            assert_eq!(framed.len(), MIN_HEADER_SIZE + header_len + payload_len);

            // Roundtrip through the decoder.
            let mut codec = TcpFrameCodec::new();
            let mut buf = BytesMut::from(&framed[..]);
            let (msg_type, decoded_header, decoded_payload) =
                codec.decode(&mut buf).unwrap().unwrap();
            assert_eq!(msg_type, MessageType::Message);
            assert_eq!(&decoded_header[..], header.as_slice());
            assert_eq!(&decoded_payload[..], payload.as_slice());
        }
    }

    /// Counts `write` calls so tests can pin how many segments the encoder
    /// hands the socket.
    #[derive(Default)]
    struct CountingWriter {
        data: Vec<u8>,
        writes: usize,
    }

    impl Write for CountingWriter {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.writes += 1;
            self.data.extend_from_slice(buf);
            Ok(buf.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    /// Above the threshold with a typical small header, the encoder stages
    /// preamble + header into one buffer: two writes, byte-identical output.
    #[test]
    fn test_direct_path_stages_prefix_into_two_writes() {
        let header = vec![0x11u8; 64];
        let payload = vec![0xABu8; COALESCE_THRESHOLD + 1];

        let mut w = CountingWriter::default();
        TcpFrameCodec::encode_frame_sync(&mut w, MessageType::Message, &header, &payload).unwrap();
        assert_eq!(w.writes, 2, "staged prefix + payload");

        let reference =
            create_unsafe_frame(SCHEMA_VERSION_V1, MessageType::Message, &header, &payload);
        assert_eq!(w.data, reference.to_vec(), "wire bytes must be unchanged");
    }

    /// A header too large for the stack prefix falls back to three segments —
    /// same bytes on the wire.
    #[test]
    fn test_direct_path_oversized_header_falls_back() {
        let header = vec![0x22u8; DIRECT_PREFIX_CAP - MIN_HEADER_SIZE + 1];
        let payload = vec![0xABu8; COALESCE_THRESHOLD + 1];

        let mut w = CountingWriter::default();
        TcpFrameCodec::encode_frame_sync(&mut w, MessageType::Message, &header, &payload).unwrap();
        assert_eq!(
            w.writes, 3,
            "preamble, header, payload each their own write"
        );

        let reference =
            create_unsafe_frame(SCHEMA_VERSION_V1, MessageType::Message, &header, &payload);
        assert_eq!(w.data, reference.to_vec(), "wire bytes must be unchanged");
    }

    /// Once the preamble announces a frame, the decoder reserves the whole
    /// remainder so `Framed`'s read loop does not grow the buffer by
    /// doubling (a realloc + memcpy of everything received so far, at every
    /// step, for a large frame).
    #[test]
    fn test_decode_reserves_capacity_for_announced_frame() {
        let mut codec = TcpFrameCodec::new();
        let header_len = 64u32;
        let payload_len = 4 * 1024 * 1024u32;

        let mut buf = BytesMut::with_capacity(8 * 1024);
        buf.extend_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
        buf.extend_from_slice(&[MessageType::Message.as_u8()]);
        buf.extend_from_slice(&header_len.to_be_bytes());
        buf.extend_from_slice(&payload_len.to_be_bytes());
        buf.extend_from_slice(&[0u8; 1000]); // partial data

        assert!(codec.decode(&mut buf).unwrap().is_none());
        let total = (header_len + payload_len) as usize;
        assert!(
            buf.capacity() >= total,
            "decoder must reserve for the announced frame: capacity={} < {}",
            buf.capacity(),
            total
        );
    }
}