freenet 0.2.46

//! Streaming infrastructure for incremental fragment consumption.
//!
//! This module provides the `StreamingInboundStream` type which implements
//! `futures::Stream` for incremental access to incoming data. Unlike the
//! existing `InboundStream` which waits for complete reassembly, this
//! allows consumers to process data as it arrives.
//!
//! # Components
//!
//! - [`StreamingInboundStream`]: A `futures::Stream` over incoming fragment data
//! - [`StreamHandle`]: A cloneable handle for accessing an inbound stream
//! - [`StreamRegistry`]: Global registry for transport-to-operations handoff
//!
//! # Usage
//!
//! ```ignore
//! // Get a handle from the registry
//! let handle = registry.get_stream(stream_id).await?;
//!
//! // Create a streaming view
//! let stream = handle.stream();
//!
//! // Process fragments incrementally
//! while let Some(result) = stream.next().await {
//!     match result {
//!         Ok(bytes) => process_chunk(bytes),
//!         Err(e) => handle_error(e),
//!     }
//! }
//! ```

use bytes::Bytes;
use dashmap::DashMap;
use event_listener::EventListener;
use futures::Stream;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll, Waker};

use super::StreamId;
use super::streaming_buffer::LockFreeStreamBuffer;

/// No new fragments arrived within this duration -> assume the stream is dead.
///
/// This covers cases where the sending peer's connection died but the local
/// PeerConnection hasn't been torn down yet (e.g., one-directional path failure,
/// slow liveness detection).
///
/// Production data (2026-04): p90 successful large transfer completes in 673ms,
/// p95 in 1.8s. Even the slowest legitimate transfers deliver fragments sub-second.
/// A 5-second gap between fragments is conclusive evidence the stream is dead.
/// The previous 30s value caused users to wait 30s+ before retry on stalled streams
/// (0.9% of all transfers), turning a sub-second retry into a 30-40s experience.
pub const STREAM_INACTIVITY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);

/// Error type for streaming operations.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StreamError {
    /// Stream was cancelled or connection closed.
    Cancelled,
    /// Stream not found in registry.
    NotFound,
    /// Invalid fragment received.
    InvalidFragment { message: String },
    /// No fragments arrived within the inactivity timeout.
    InactivityTimeout,
}

impl std::fmt::Display for StreamError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StreamError::Cancelled => write!(f, "stream was cancelled"),
            StreamError::NotFound => write!(f, "stream not found in registry"),
            StreamError::InvalidFragment { message } => {
                write!(f, "invalid fragment: {}", message)
            }
            StreamError::InactivityTimeout => {
                write!(f, "no fragments received within inactivity timeout")
            }
        }
    }
}

impl std::error::Error for StreamError {}

/// Synchronization state for stream handles.
///
/// The lock-free buffer is stored separately (as Arc<LockFreeStreamBuffer>)
/// so that fragment insertion doesn't require holding this lock.
/// This struct only holds state that requires synchronization.
struct SyncState {
    /// True if the stream has been cancelled.
    cancelled: bool,
    /// Wakers for poll_next calls waiting on data.
    wakers: Vec<Waker>,
}

impl SyncState {
    fn new() -> Self {
        Self {
            cancelled: false,
            wakers: Vec::new(),
        }
    }

    fn wake_all(&mut self) {
        for waker in self.wakers.drain(..) {
            waker.wake();
        }
    }
}

/// A handle to an inbound stream that can be cloned for multiple consumers.
///
/// Each clone maintains its own read position, allowing independent consumers
/// to read the stream at different rates.
///
/// The design separates the lock-free buffer from synchronization state:
/// - `buffer`: Lock-free fragment storage using `OnceLock<Bytes>` slots
/// - `sync`: Cancelled flag and wakers (requires lock)
///
/// Fragment insertion goes directly to the buffer without acquiring any lock,
/// and the buffer's built-in `Notify` handles async notification.
#[derive(Clone)]
pub struct StreamHandle {
    /// Lock-free buffer for fragment storage.
    /// Accessed without holding any lock - all operations are atomic.
    buffer: Arc<LockFreeStreamBuffer>,
    /// Synchronization state (cancelled, wakers) - requires lock.
    sync: Arc<parking_lot::RwLock<SyncState>>,
    /// Stream ID for debugging.
    stream_id: StreamId,
    /// Total expected bytes.
    total_bytes: u64,
}

impl StreamHandle {
    /// Creates a new stream handle.
    ///
    /// This is public for use by `OrphanStreamRegistry` which needs to create
    /// handles and for transport layer stream registration.
    pub fn new(stream_id: StreamId, total_bytes: u64) -> Self {
        Self {
            buffer: Arc::new(LockFreeStreamBuffer::new(total_bytes)),
            sync: Arc::new(parking_lot::RwLock::new(SyncState::new())),
            stream_id,
            total_bytes,
        }
    }

    /// Returns the stream ID.
    #[allow(dead_code)]
    pub(crate) fn stream_id(&self) -> StreamId {
        self.stream_id
    }

    /// Returns the total expected bytes.
    pub fn total_bytes(&self) -> u64 {
        self.total_bytes
    }

    /// Returns true if all fragments have been received.
    pub fn is_complete(&self) -> bool {
        self.buffer.is_complete()
    }

    /// Returns the number of fragments received so far.
    pub fn received_fragments(&self) -> usize {
        self.buffer.inserted_count()
    }

    /// Returns the total expected number of fragments.
    pub fn total_fragments(&self) -> usize {
        self.buffer.total_fragments()
    }

    /// Creates a new streaming view starting from fragment 1.
    ///
    /// Each call creates an independent stream with its own read position.
    /// Fragments are cloned but remain in the buffer, allowing multiple
    /// consumers to read the same data via `fork()`.
    ///
    /// # Memory Behavior
    ///
    /// All fragments remain in the buffer until the stream is dropped or
    /// explicitly cleared. This means:
    /// - **Pro**: Multiple consumers can read the same data independently
    /// - **Pro**: Safe to use with `fork()` for parallel processing
    /// - **Con**: Memory usage equals full stream size until completion
    ///
    /// # When to Use
    ///
    /// Use `stream()` when:
    /// - Multiple consumers need to read the same data
    /// - You need to fork the stream for parallel processing
    /// - Memory is not a concern for the stream size
    /// - You might need to re-read fragments
    ///
    /// For single-consumer scenarios with large streams, prefer
    /// [`stream_with_reclaim()`](Self::stream_with_reclaim) for better memory efficiency.
    pub fn stream(&self) -> StreamingInboundStream {
        StreamingInboundStream {
            handle: self.clone(),
            next_fragment: 1,
            bytes_read: 0,
            auto_reclaim: false,
            listener: None,
        }
    }

    /// Creates a streaming view with automatic memory reclamation.
    ///
    /// Unlike [`stream()`](Self::stream), this version takes ownership of
    /// fragments as they are read, freeing memory progressively.
    ///
    /// # Memory Behavior
    ///
    /// Each fragment is removed from the buffer immediately after being read:
    /// - **Pro**: Memory usage stays constant regardless of stream size
    /// - **Pro**: Ideal for processing large streams (10MB+) without OOM
    /// - **Con**: Fragments cannot be read again once consumed
    /// - **Con**: Incompatible with `fork()` or multiple consumers
    ///
    /// # When to Use
    ///
    /// Use `stream_with_reclaim()` when:
    /// - Only one consumer will read the stream
    /// - The stream is large and memory is a concern
    /// - Data is processed once and discarded (e.g., forwarding, hashing)
    /// - You don't need `fork()` or parallel consumers
    ///
    /// # Warning
    ///
    /// **Do not use with `fork()` or multiple consumers.** Once a fragment is
    /// read by this stream, it is permanently removed from the buffer. Other
    /// consumers (including forked handles) will wait forever for fragments
    /// that have already been consumed.
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Single consumer processing a large file
    /// let mut stream = handle.stream_with_reclaim();
    /// let mut hasher = Sha256::new();
    /// while let Some(chunk) = stream.next().await {
    ///     hasher.update(&chunk?);
    ///     // Memory for this chunk is freed immediately
    /// }
    /// ```
    pub fn stream_with_reclaim(&self) -> StreamingInboundStream {
        StreamingInboundStream {
            handle: self.clone(),
            next_fragment: 1,
            bytes_read: 0,
            auto_reclaim: true,
            listener: None,
        }
    }

    /// Forks this handle, creating an independent consumer.
    ///
    /// The forked handle shares the same underlying buffer but has its own
    /// independent `SyncState` (cancelled flag and wakers). This means:
    /// - Cancelling the original handle does NOT cancel the fork
    /// - Each consumer has its own waker list
    /// - Fragment notifications reach the fork via the buffer's `Event` notifier
    ///
    /// This independence is critical for piped stream forwarding: when the
    /// upstream `PeerConnection` drops and cancels its handles, the pipe task
    /// (which reads from a forked handle) must continue forwarding data that's
    /// already in the shared buffer.
    pub fn fork(&self) -> Self {
        // Preserve the current cancelled state so forking an already-dead
        // stream doesn't resurrect it. Future cancellations on the original
        // won't propagate to the fork (independent SyncState).
        let already_cancelled = self.sync.read().cancelled;
        let mut sync = SyncState::new();
        sync.cancelled = already_cancelled;
        Self {
            buffer: self.buffer.clone(),
            sync: Arc::new(parking_lot::RwLock::new(sync)),
            stream_id: self.stream_id,
            total_bytes: self.total_bytes,
        }
    }

    /// Inserts a fragment into the stream buffer.
    ///
    /// This is called by the transport layer when a fragment arrives.
    /// The insert operation is lock-free - only the cancelled check and
    /// waker notification require synchronization.
    ///
    /// # Returns
    ///
    /// * `Ok(true)` - Fragment was inserted successfully
    /// * `Ok(false)` - Fragment was a duplicate (no-op)
    /// * `Err(StreamError)` - Invalid fragment number or stream cancelled
    pub(crate) fn push_fragment(
        &self,
        fragment_number: u32,
        data: Bytes,
    ) -> Result<bool, StreamError> {
        // Quick check if cancelled (read lock only)
        if self.sync.read().cancelled {
            return Err(StreamError::Cancelled);
        }

        // Lock-free insert into buffer
        // The buffer's internal Notify handles async notification
        match self.buffer.insert(fragment_number, data) {
            Ok(inserted) => {
                if inserted {
                    // Wake synchronous waiters (poll_next wakers)
                    self.sync.write().wake_all();
                }
                Ok(inserted)
            }
            Err(e) => Err(StreamError::InvalidFragment {
                message: e.to_string(),
            }),
        }
    }

    /// Cancels the stream, waking all waiters.
    pub(crate) fn cancel(&self) {
        let mut sync = self.sync.write();
        sync.cancelled = true;
        sync.wake_all();
        drop(sync); // Release lock before notifying

        // Wake async waiters blocked on buffer.notifier().listen()
        self.buffer.notifier().notify(usize::MAX);
    }

    /// Assembles the complete data if all fragments are present.
    ///
    /// This is a convenience method for waiting until the stream is complete
    /// and then getting all data at once.
    pub fn try_assemble(&self) -> Option<Vec<u8>> {
        self.buffer.assemble()
    }

    /// Waits for the stream to complete and returns the assembled data.
    ///
    /// This is useful when you don't need incremental processing.
    /// Uses the buffer's built-in Notify for efficient async waiting.
    ///
    /// Includes an inactivity timeout ([`STREAM_INACTIVITY_TIMEOUT`]): if no new
    /// fragment arrives within the timeout window, returns
    /// [`StreamError::InactivityTimeout`]. This prevents operations from hanging
    /// forever when the sending peer's connection dies mid-transfer.
    pub async fn assemble(&self) -> Result<Vec<u8>, StreamError> {
        loop {
            // Check cancelled state
            if self.sync.read().cancelled {
                return Err(StreamError::Cancelled);
            }

            // Try to assemble (lock-free access to buffer)
            if self.buffer.is_complete() {
                if let Some(data) = self.buffer.assemble() {
                    return Ok(data);
                }
                // Don't fail immediately - the overflow fragment may not have arrived yet.
                // When fragment #1 has reduced payload (metadata embedded), the sender uses
                // one extra fragment that goes in the overflow slot. is_complete() fires when
                // min_complete_fragments are contiguous, but actual data may be insufficient
                // until the overflow fragment arrives. Continue waiting.
                tracing::debug!(
                    stream_id = %self.stream_id.0,
                    total_bytes = self.total_bytes,
                    received_fragments = self.buffer.inserted_count(),
                    total_fragments = self.buffer.total_fragments(),
                    "Stream marked complete but assembly insufficient, waiting for overflow fragment"
                );
            }

            // Wait for notification with inactivity timeout.
            // Each arriving fragment triggers a notification, so the timeout resets
            // on every fragment. Only fires when the stream is truly stalled.
            match tokio::time::timeout(STREAM_INACTIVITY_TIMEOUT, self.buffer.notifier().listen())
                .await
            {
                Ok(()) => { /* fragment arrived or stream cancelled — loop to check */ }
                Err(_) => {
                    // Re-check before declaring timeout — a fragment may have arrived
                    // in the race window between is_complete()/listen() above (TOCTOU).
                    if self.sync.read().cancelled {
                        return Err(StreamError::Cancelled);
                    }
                    if let Some(data) = self.buffer.assemble() {
                        return Ok(data);
                    }
                    tracing::warn!(
                        stream_id = %self.stream_id.0,
                        total_bytes = self.total_bytes,
                        received_fragments = self.buffer.inserted_count(),
                        total_fragments = self.buffer.total_fragments(),
                        timeout_secs = STREAM_INACTIVITY_TIMEOUT.as_secs(),
                        "Stream assembly timed out — no fragments received within inactivity window"
                    );
                    return Err(StreamError::InactivityTimeout);
                }
            }
        }
    }
}

impl std::fmt::Debug for StreamHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StreamHandle")
            .field("stream_id", &self.stream_id)
            .field("total_bytes", &self.total_bytes)
            .field("received", &self.buffer.inserted_count())
            .field("total_fragments", &self.buffer.total_fragments())
            .field("complete", &self.buffer.is_complete())
            .field("cancelled", &self.sync.read().cancelled)
            .finish()
    }
}

/// A `futures::Stream` that yields contiguous fragments as they become available.
///
/// This stream yields `Bytes` chunks in order, waiting for out-of-order
/// fragments to arrive before yielding subsequent chunks.
///
/// # Memory Behavior
///
/// By default (`auto_reclaim = false`), fragments are cloned but remain in the buffer,
/// allowing multiple consumers to read the same data via `fork()`.
///
/// With `auto_reclaim = true`, fragments are taken from the buffer after reading,
/// freeing memory progressively. This is ideal for single-consumer scenarios with
/// large streams.
pub struct StreamingInboundStream {
    /// Handle to the shared stream state.
    handle: StreamHandle,
    /// Next fragment index to yield (1-indexed).
    next_fragment: u32,
    /// Total bytes read so far.
    bytes_read: u64,
    /// If true, fragments are taken (removed) from the buffer after reading.
    /// This enables progressive memory reclamation for single-consumer scenarios.
    auto_reclaim: bool,
    /// Listener for buffer data_available notifications.
    /// Created when poll_next returns Pending, consumed when notified.
    /// Uses the buffer's `event_listener::Event` which is fired by every
    /// `buffer.insert()` — independent of which handle called `push_fragment`.
    listener: Option<Pin<Box<EventListener>>>,
}

impl StreamingInboundStream {
    /// Returns the next fragment number that will be yielded.
    #[allow(dead_code)]
    pub fn next_fragment_number(&self) -> u32 {
        self.next_fragment
    }

    /// Returns the total bytes read so far.
    #[allow(dead_code)]
    pub fn bytes_read(&self) -> u64 {
        self.bytes_read
    }

    /// Returns the stream ID.
    #[allow(dead_code)]
    pub(crate) fn stream_id(&self) -> StreamId {
        self.handle.stream_id
    }

    /// Returns whether auto-reclaim is enabled.
    #[allow(dead_code)]
    pub fn is_auto_reclaim(&self) -> bool {
        self.auto_reclaim
    }

    /// Tries to get the next fragment, either taking or cloning based on auto_reclaim.
    fn try_get_fragment(&self, idx: u32) -> Option<Bytes> {
        if self.auto_reclaim {
            // Take ownership and clear the slot
            self.handle.buffer.take(idx)
        } else {
            // Clone, leaving the fragment in place for other consumers
            self.handle.buffer.get(idx).cloned()
        }
    }
}

impl Stream for StreamingInboundStream {
    type Item = Result<Bytes, StreamError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let next_idx = self.next_fragment;

        // Check cancelled state (per-handle, independent for forks)
        if self.handle.sync.read().cancelled {
            return Poll::Ready(Some(Err(StreamError::Cancelled)));
        }

        // Lock-free access to buffer
        let total_fragments = self.handle.buffer.total_fragments() as u32;

        // Check if we've read all fragments
        if next_idx > total_fragments {
            return Poll::Ready(None); // Stream complete
        }

        // If the stream has enough data but this fragment slot is empty, we're done.
        // The overflow slot (allocated for potential metadata overhead in fragment #1)
        // may not be used when fragment #1 has full payload capacity.
        if self.handle.buffer.is_complete() && self.handle.buffer.get(next_idx).is_none() {
            return Poll::Ready(None);
        }

        // Try to get the next fragment (lock-free fast path)
        if let Some(data) = self.try_get_fragment(next_idx) {
            self.listener = None;
            self.next_fragment = next_idx + 1;
            self.bytes_read += data.len() as u64;
            return Poll::Ready(Some(Ok(data)));
        }

        // Fragment not yet available.
        //
        // Following the event_listener pattern: create the listener BEFORE
        // re-checking the condition. This prevents a TOCTOU race where a
        // notification fires between the check and listener creation.
        //
        // We use the buffer's `data_available` Event (fired by every
        // buffer.insert()) as the primary notification mechanism. This works
        // for both original and forked handles since buffer.insert() fires
        // regardless of which handle called push_fragment().
        //
        // cancel() also fires buffer.notifier().notify(), so the listener
        // covers both data arrival and cancellation.
        if self.listener.is_none() {
            self.listener = Some(Box::pin(self.handle.buffer.notifier().listen()));
        }

        // Re-check after listener creation (the event_listener pattern):
        // if a fragment arrived between the fast-path check and listen(),
        // we catch it here instead of missing the notification.
        if self.handle.sync.read().cancelled {
            self.listener = None;
            return Poll::Ready(Some(Err(StreamError::Cancelled)));
        }
        if let Some(data) = self.try_get_fragment(next_idx) {
            self.listener = None;
            self.next_fragment = next_idx + 1;
            self.bytes_read += data.len() as u64;
            return Poll::Ready(Some(Ok(data)));
        }

        // Poll the listener. The listener is guaranteed to be Some here
        // (set above), but we use if-let to satisfy the no-unwrap rule.
        if let Some(listener) = self.listener.as_mut() {
            match listener.as_mut().poll(cx) {
                Poll::Ready(()) => {
                    // Notified — clear listener and re-check buffer inline
                    // (avoids wake_by_ref spin-loop).
                    self.listener = None;
                    if self.handle.sync.read().cancelled {
                        return Poll::Ready(Some(Err(StreamError::Cancelled)));
                    }
                    if let Some(data) = self.try_get_fragment(next_idx) {
                        self.next_fragment = next_idx + 1;
                        self.bytes_read += data.len() as u64;
                        return Poll::Ready(Some(Ok(data)));
                    }
                    // Spurious notification. Create a fresh listener,
                    // re-check, and register waker.
                    self.listener = Some(Box::pin(self.handle.buffer.notifier().listen()));
                    if let Some(data) = self.try_get_fragment(next_idx) {
                        self.listener = None;
                        self.next_fragment = next_idx + 1;
                        self.bytes_read += data.len() as u64;
                        return Poll::Ready(Some(Ok(data)));
                    }
                    if let Some(new_listener) = self.listener.as_mut() {
                        match new_listener.as_mut().poll(cx) {
                            Poll::Ready(()) => self.listener = None,
                            Poll::Pending => {}
                        }
                    }
                    Poll::Pending
                }
                Poll::Pending => Poll::Pending,
            }
        } else {
            Poll::Pending
        }
    }
}

/// Global registry for managing active inbound streams.
///
/// The registry provides the handoff point between the transport layer
/// (which receives fragments) and the operations layer (which consumes them).
///
/// Uses `DashMap` for lock-free concurrent access, avoiding the global
/// bottleneck of `tokio::sync::RwLock<HashMap>`.
pub struct StreamRegistry {
    /// Active streams indexed by stream ID.
    /// Uses DashMap for lock-free concurrent access.
    streams: DashMap<StreamId, StreamHandle>,
}

impl StreamRegistry {
    /// Creates a new stream registry.
    pub fn new() -> Self {
        Self {
            streams: DashMap::new(),
        }
    }

    /// Registers a new stream with the given ID and expected size.
    ///
    /// Returns a handle for pushing fragments to the stream.
    ///
    /// If a stream with this ID already exists, returns the existing handle.
    pub(crate) fn register(&self, stream_id: StreamId, total_bytes: u64) -> StreamHandle {
        // Use entry API to avoid race conditions
        self.streams
            .entry(stream_id)
            .or_insert_with(|| StreamHandle::new(stream_id, total_bytes))
            .clone()
    }

    /// Gets a handle to an existing stream.
    ///
    /// Returns `None` if the stream doesn't exist.
    #[allow(dead_code)]
    pub(crate) fn get(&self, stream_id: StreamId) -> Option<StreamHandle> {
        self.streams.get(&stream_id).map(|r| r.clone())
    }

    /// Removes a stream from the registry.
    ///
    /// This should be called when a stream is complete or cancelled
    /// to prevent memory leaks.
    #[allow(dead_code)]
    pub(crate) fn remove(&self, stream_id: StreamId) -> Option<StreamHandle> {
        self.streams.remove(&stream_id).map(|(_, h)| h)
    }

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

    /// Cancels all streams and clears the registry.
    ///
    /// This prevents memory leaks by removing all entries after cancellation.
    pub fn cancel_all(&self) {
        // Cancel all streams first
        for entry in self.streams.iter() {
            entry.value().cancel();
        }
        // Clear the registry to prevent memory leaks
        self.streams.clear();
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::GlobalExecutor;
    use futures::StreamExt;

    fn make_stream_id() -> StreamId {
        StreamId::next()
    }

    #[test]
    fn test_stream_handle_creation() {
        let id = make_stream_id();
        let handle = StreamHandle::new(id, 1000);

        assert_eq!(handle.stream_id(), id);
        assert_eq!(handle.total_bytes(), 1000);
        assert!(!handle.is_complete());
    }

    #[test]
    fn test_stream_handle_push_fragment() {
        let handle = StreamHandle::new(make_stream_id(), 100);

        let result = handle.push_fragment(1, Bytes::from_static(b"hello"));
        assert!(result.is_ok());
        assert!(result.unwrap()); // First insert succeeds

        // Duplicate is no-op
        let result = handle.push_fragment(1, Bytes::from_static(b"world"));
        assert!(result.is_ok());
        assert!(!result.unwrap()); // Duplicate returns false
    }

    #[test]
    fn test_stream_handle_invalid_fragment() {
        let handle = StreamHandle::new(make_stream_id(), 100);

        // Index 0 is invalid
        let result = handle.push_fragment(0, Bytes::from_static(b"hello"));
        assert!(matches!(result, Err(StreamError::InvalidFragment { .. })));

        // Index 3 is out of bounds (total_fragments=2: 1 base + 1 overflow)
        let result = handle.push_fragment(3, Bytes::from_static(b"hello"));
        assert!(matches!(result, Err(StreamError::InvalidFragment { .. })));
    }

    #[test]
    fn test_stream_handle_cancel() {
        let handle = StreamHandle::new(make_stream_id(), 100);

        handle.cancel();

        // Pushing after cancel should fail
        let result = handle.push_fragment(1, Bytes::from_static(b"hello"));
        assert!(matches!(result, Err(StreamError::Cancelled)));
    }

    #[test]
    fn test_stream_handle_fork() {
        let handle = StreamHandle::new(make_stream_id(), 100);
        handle
            .push_fragment(1, Bytes::from_static(b"hello"))
            .unwrap();

        let forked = handle.fork();

        // Both handles see the same data
        assert!(handle.is_complete());
        assert!(forked.is_complete());
        assert_eq!(handle.try_assemble(), forked.try_assemble());
    }

    #[test]
    fn test_fork_independent_cancellation() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        // Use exact fragment size so one fragment completes the stream
        let total = FRAGMENT_PAYLOAD_SIZE as u64;
        let data = vec![42u8; FRAGMENT_PAYLOAD_SIZE];
        let handle = StreamHandle::new(make_stream_id(), total);
        handle.push_fragment(1, Bytes::from(data.clone())).unwrap();

        let forked = handle.fork();

        // Cancel the original — simulates PeerConnection::Drop
        handle.cancel();

        // Original is cancelled
        assert!(handle.sync.read().cancelled);
        assert!(matches!(
            handle.push_fragment(1, Bytes::from(vec![0u8; 10])),
            Err(StreamError::Cancelled)
        ));

        // Fork is NOT cancelled — pipe task can continue reading
        assert!(!forked.sync.read().cancelled);

        // Fork can still read data from the shared buffer
        assert!(forked.is_complete());
        assert_eq!(forked.try_assemble(), Some(data));
    }

    #[tokio::test]
    async fn test_fork_stream_reads_after_original_cancel() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push both fragments via the original handle
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();
        handle
            .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // Fork and create a stream from it
        let forked = handle.fork();
        let mut stream = forked.stream();

        // Cancel the original (simulates PeerConnection drop)
        handle.cancel();

        // The forked stream should still be able to read all fragments
        let chunk1 = stream.next().await;
        assert!(chunk1.is_some());
        assert_eq!(chunk1.unwrap().unwrap().len(), FRAGMENT_PAYLOAD_SIZE);

        let chunk2 = stream.next().await;
        assert!(chunk2.is_some());
        assert_eq!(chunk2.unwrap().unwrap().len(), FRAGMENT_PAYLOAD_SIZE);

        // Stream complete
        let chunk3 = stream.next().await;
        assert!(chunk3.is_none());
    }

    #[tokio::test]
    async fn test_fork_incremental_wakeup() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push only fragment #1
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // Fork and start reading
        let forked = handle.fork();
        let mut stream = forked.stream();

        // Read fragment #1 (available immediately)
        let chunk1 = stream.next().await;
        assert!(chunk1.is_some());
        assert_eq!(chunk1.unwrap().unwrap(), vec![1u8; FRAGMENT_PAYLOAD_SIZE]);

        // Fragment #2 arrives later via the ORIGINAL handle's push_fragment.
        // The forked stream must wake up via the buffer's EventListener
        // (not via SyncState wakers, which are on the fork's independent sync).
        let handle_clone = handle.clone();
        tokio::spawn(async move {
            tokio::task::yield_now().await;
            handle_clone
                .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
        });

        // The forked stream should receive fragment #2 via EventListener
        let chunk2 = tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()).await;
        assert!(
            chunk2.is_ok(),
            "fork should wake up when fragment arrives via original handle"
        );
        let chunk2 = chunk2.unwrap();
        assert!(chunk2.is_some());
        assert_eq!(chunk2.unwrap().unwrap(), vec![2u8; FRAGMENT_PAYLOAD_SIZE]);

        // Push fragment #3 and verify
        handle
            .push_fragment(3, Bytes::from(vec![3u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();
        let chunk3 = stream.next().await;
        assert!(chunk3.is_some());
        assert_eq!(chunk3.unwrap().unwrap(), vec![3u8; FRAGMENT_PAYLOAD_SIZE]);

        // Stream complete
        let end = stream.next().await;
        assert!(end.is_none());
    }

    #[tokio::test]
    async fn test_streaming_inbound_stream_basic() {
        let handle = StreamHandle::new(make_stream_id(), 15);

        // Push all fragments
        handle
            .push_fragment(1, Bytes::from_static(b"hello"))
            .unwrap();

        let mut stream = handle.stream();
        let chunk = stream.next().await;
        assert!(chunk.is_some());
        assert_eq!(chunk.unwrap().unwrap(), Bytes::from_static(b"hello"));

        // Stream should be exhausted
        let chunk = stream.next().await;
        assert!(chunk.is_none());
    }

    #[tokio::test]
    async fn test_streaming_inbound_stream_incremental() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);
        let mut stream = handle.stream();

        // Spawn a task to push fragments with delays
        let handle_clone = handle.clone();
        let producer = GlobalExecutor::spawn(async move {
            for i in 1..=3 {
                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                handle_clone
                    .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                    .unwrap();
            }
        });

        // Consume incrementally
        let mut count = 0;
        while let Some(result) = stream.next().await {
            count += 1;
            let bytes = result.unwrap();
            assert_eq!(bytes.len(), FRAGMENT_PAYLOAD_SIZE);
        }

        producer.await.unwrap();
        assert_eq!(count, 3);
    }

    #[tokio::test]
    async fn test_streaming_cancelled() {
        let handle = StreamHandle::new(make_stream_id(), 100);
        let mut stream = handle.stream();

        // Cancel the stream
        handle.cancel();

        // Next poll should return Cancelled error
        let result = stream.next().await;
        assert!(matches!(result, Some(Err(StreamError::Cancelled))));
    }

    #[tokio::test]
    async fn test_stream_registry_register_and_get() {
        let registry = StreamRegistry::new();
        let id = make_stream_id();

        let handle = registry.register(id, 1000);
        assert_eq!(handle.stream_id(), id);

        let retrieved = registry.get(id);
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().stream_id(), id);
    }

    #[tokio::test]
    async fn test_stream_registry_register_existing() {
        let registry = StreamRegistry::new();
        let id = make_stream_id();

        let handle1 = registry.register(id, 1000);
        let handle2 = registry.register(id, 2000); // Different size

        // Should return the same handle (first registration wins)
        assert_eq!(handle1.total_bytes(), handle2.total_bytes());
    }

    #[tokio::test]
    async fn test_stream_registry_remove() {
        let registry = StreamRegistry::new();
        let id = make_stream_id();

        registry.register(id, 1000);
        assert_eq!(registry.stream_count(), 1);

        let removed = registry.remove(id);
        assert!(removed.is_some());
        assert_eq!(registry.stream_count(), 0);

        // Get should return None after removal
        assert!(registry.get(id).is_none());
    }

    #[tokio::test]
    async fn test_stream_handle_assemble_async() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Spawn producer
        let handle_clone = handle.clone();
        let producer = GlobalExecutor::spawn(async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
            handle_clone
                .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
            handle_clone
                .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
        });

        // Wait for complete assembly
        let result = handle.assemble().await;
        producer.await.unwrap();

        assert!(result.is_ok());
        let data = result.unwrap();
        assert_eq!(data.len(), total as usize);
    }

    #[tokio::test]
    async fn test_assemble_inactivity_timeout() {
        // Create a stream that will never receive all its fragments.
        // assemble() should return InactivityTimeout after STREAM_INACTIVITY_TIMEOUT.
        let handle = StreamHandle::new(make_stream_id(), 1000);

        // Push one fragment but not enough to complete (needs ~1000 bytes)
        handle
            .push_fragment(1, Bytes::from(vec![0u8; 100]))
            .unwrap();

        // Use tokio::time::pause to avoid actually waiting 30 seconds
        tokio::time::pause();

        let result = handle.assemble().await;
        assert!(
            matches!(result, Err(StreamError::InactivityTimeout)),
            "Expected InactivityTimeout, got: {:?}",
            result
        );
    }

    #[tokio::test]
    async fn test_assemble_timeout_resets_per_fragment() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        // Verify the timeout is per-fragment (resets on each arrival), not absolute.
        // Send fragments with 3-second gaps (under the 5s timeout) -- should succeed.
        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        tokio::time::pause();

        let handle_clone = handle.clone();
        let producer = GlobalExecutor::spawn(async move {
            for i in 1..=3u32 {
                // 3 seconds between fragments -- under the 5s inactivity timeout
                tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
                handle_clone
                    .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                    .unwrap();
            }
        });

        // Total wall time: 9 seconds (3 x 3s), but no single gap exceeds 5s
        let result = handle.assemble().await;
        producer.await.unwrap();

        assert!(
            result.is_ok(),
            "Slow-but-active stream should succeed: {:?}",
            result
        );
        assert_eq!(result.unwrap().len(), total as usize);
    }

    /// Regression test: a stream that receives some fragments then goes silent for longer
    /// than STREAM_INACTIVITY_TIMEOUT (5s) must be killed with InactivityTimeout.
    /// Previously this timeout was 30s, causing users to wait 30-40s before retry on
    /// stalled streams. Production data showed p90 successful transfer completes in 673ms,
    /// making 30s of silence conclusive evidence of a dead stream.
    #[tokio::test]
    async fn test_stalled_stream_killed_within_inactivity_timeout() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 5) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        tokio::time::pause();

        // Deliver 2 of 5 fragments, then go silent (simulating a stalled sender).
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();
        handle
            .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        let start = tokio::time::Instant::now();
        let result = handle.assemble().await;

        // Must fail with InactivityTimeout, not hang for 30s.
        assert!(
            matches!(result, Err(StreamError::InactivityTimeout)),
            "Stalled stream should be killed by inactivity timeout, got: {:?}",
            result
        );

        // Verify the timeout fired at ~5s, not ~30s.
        let elapsed = start.elapsed();
        assert!(
            elapsed < std::time::Duration::from_secs(10),
            "Inactivity timeout should fire within ~5s, took {:?}",
            elapsed
        );
    }

    #[tokio::test]
    async fn test_multiple_independent_streams() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push first fragment
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // Create two independent streams
        let mut stream1 = handle.stream();
        let mut stream2 = handle.stream();

        // Both can read the first fragment independently
        let chunk1 = stream1.next().await.unwrap().unwrap();
        let chunk2 = stream2.next().await.unwrap().unwrap();
        assert_eq!(chunk1, chunk2);

        // Both are now waiting for fragment 2
        assert_eq!(stream1.next_fragment_number(), 2);
        assert_eq!(stream2.next_fragment_number(), 2);

        // Push second fragment
        handle
            .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // Both can read independently
        let chunk1 = stream1.next().await.unwrap().unwrap();
        let chunk2 = stream2.next().await.unwrap().unwrap();
        assert_eq!(chunk1, chunk2);

        // Both streams are exhausted
        assert!(stream1.next().await.is_none());
        assert!(stream2.next().await.is_none());
    }

    #[test]
    fn test_zero_byte_stream() {
        let handle = StreamHandle::new(make_stream_id(), 0);

        assert_eq!(handle.total_bytes(), 0);
        assert_eq!(handle.total_fragments(), 0);
        assert!(handle.is_complete()); // Empty stream is complete
        assert_eq!(handle.try_assemble(), Some(vec![]));
    }

    #[tokio::test]
    async fn test_zero_byte_stream_streaming() {
        let handle = StreamHandle::new(make_stream_id(), 0);
        let mut stream = handle.stream();

        // Should immediately return None (no fragments)
        assert!(stream.next().await.is_none());
    }

    #[tokio::test]
    async fn test_assemble_cancelled_stream() {
        let handle = StreamHandle::new(make_stream_id(), 100);
        handle.cancel();

        let result = handle.assemble().await;
        assert!(matches!(result, Err(StreamError::Cancelled)));
    }

    #[tokio::test]
    async fn test_cancel_during_assemble() {
        let handle = StreamHandle::new(make_stream_id(), 100);
        let handle_clone = handle.clone();

        // Start assemble in background
        let assemble_task = GlobalExecutor::spawn(async move { handle_clone.assemble().await });

        // Give it time to start waiting
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        // Cancel the stream
        handle.cancel();

        // Assemble should return Cancelled
        let result = tokio::time::timeout(tokio::time::Duration::from_millis(100), assemble_task)
            .await
            .expect("timeout")
            .expect("join");
        assert!(matches!(result, Err(StreamError::Cancelled)));
    }

    #[tokio::test]
    async fn test_poll_after_stream_exhausted() {
        let handle = StreamHandle::new(make_stream_id(), 5);
        handle
            .push_fragment(1, Bytes::from_static(b"hello"))
            .unwrap();

        let mut stream = handle.stream();

        // Read the only fragment
        let chunk = stream.next().await;
        assert!(chunk.is_some());

        // Multiple polls after exhaustion should all return None
        assert!(stream.next().await.is_none());
        assert!(stream.next().await.is_none());
        assert!(stream.next().await.is_none());
    }

    #[tokio::test]
    async fn test_out_of_order_fragments_stream_waits() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push fragment 3 first (out of order)
        handle
            .push_fragment(3, Bytes::from(vec![3u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        let mut stream = handle.stream();

        // Create a task that tries to read - should block waiting for fragment 1
        let handle_clone = handle.clone();
        let read_task = GlobalExecutor::spawn(async move {
            let mut s = handle_clone.stream();
            s.next().await
        });

        // Give it time to block
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;

        // Task should still be pending
        assert!(!read_task.is_finished());

        // Now push fragment 1
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // Read task should complete
        let result = tokio::time::timeout(tokio::time::Duration::from_millis(100), read_task)
            .await
            .expect("timeout")
            .expect("join");
        assert!(result.is_some());
        assert!(result.unwrap().is_ok());

        // Now we should be able to read from our stream too
        let chunk = stream.next().await;
        assert!(chunk.is_some());
    }

    #[tokio::test]
    async fn test_cancel_while_streaming() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push first fragment
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        let mut stream = handle.stream();

        // Read first fragment successfully
        let chunk = stream.next().await;
        assert!(chunk.is_some());
        assert!(chunk.unwrap().is_ok());

        // Cancel while waiting for second fragment
        let handle_clone = handle.clone();
        let cancel_task = GlobalExecutor::spawn(async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
            handle_clone.cancel();
        });

        // Try to read - should eventually return Cancelled
        let result = stream.next().await;
        assert!(matches!(result, Some(Err(StreamError::Cancelled))));

        cancel_task.await.unwrap();
    }

    #[tokio::test]
    async fn test_bytes_read_tracking() {
        let handle = StreamHandle::new(make_stream_id(), 15);
        handle
            .push_fragment(1, Bytes::from_static(b"hello world!!!"))
            .unwrap();

        let mut stream = handle.stream();
        assert_eq!(stream.bytes_read(), 0);

        let chunk = stream.next().await.unwrap().unwrap();
        assert_eq!(stream.bytes_read(), chunk.len() as u64);
    }

    #[tokio::test]
    async fn test_registry_cancel_all() {
        let registry = StreamRegistry::new();

        let id1 = make_stream_id();
        let id2 = make_stream_id();

        let handle1 = registry.register(id1, 100);
        let handle2 = registry.register(id2, 200);

        // Cancel all (now clears the registry too)
        registry.cancel_all();

        // Both should be cancelled
        let result1 = handle1.push_fragment(1, Bytes::from_static(b"test"));
        let result2 = handle2.push_fragment(1, Bytes::from_static(b"test"));

        assert!(matches!(result1, Err(StreamError::Cancelled)));
        assert!(matches!(result2, Err(StreamError::Cancelled)));

        // Registry should be empty (memory leak fix)
        assert_eq!(registry.stream_count(), 0);
    }

    #[tokio::test]
    async fn test_registry_cleanup_on_normal_completion() {
        // This test verifies that streams can be properly removed from the registry
        // when they complete normally (simulating the cleanup done in peer_connection.rs)
        let registry = StreamRegistry::new();

        let id1 = make_stream_id();
        let id2 = make_stream_id();
        let id3 = make_stream_id();

        // Register multiple streams
        let handle1 = registry.register(id1, 100);
        let _handle2 = registry.register(id2, 200);
        let _handle3 = registry.register(id3, 300);

        assert_eq!(registry.stream_count(), 3);

        // Complete stream 1's data
        handle1
            .push_fragment(1, Bytes::from_static(b"complete data here!"))
            .unwrap();
        assert!(handle1.is_complete());

        // Simulate normal completion cleanup (as done in peer_connection.rs)
        let removed = registry.remove(id1);
        assert!(removed.is_some());
        assert_eq!(registry.stream_count(), 2);

        // Remove another stream
        let removed = registry.remove(id2);
        assert!(removed.is_some());
        assert_eq!(registry.stream_count(), 1);

        // Remove last stream
        let removed = registry.remove(id3);
        assert!(removed.is_some());
        assert_eq!(registry.stream_count(), 0);

        // Registry should be completely empty - no memory leak
        assert!(registry.get(id1).is_none());
        assert!(registry.get(id2).is_none());
        assert!(registry.get(id3).is_none());
    }

    #[test]
    fn test_registry_get_nonexistent() {
        let registry = StreamRegistry::new();
        let id = make_stream_id();

        let result = registry.get(id);
        assert!(result.is_none());
    }

    #[test]
    fn test_registry_remove_nonexistent() {
        let registry = StreamRegistry::new();
        let id = make_stream_id();

        let result = registry.remove(id);
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_concurrent_push_and_stream() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;
        use std::sync::atomic::{AtomicUsize, Ordering};

        let total = (FRAGMENT_PAYLOAD_SIZE * 10) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);
        let fragments_received = Arc::new(AtomicUsize::new(0));

        // Spawn producer
        let handle_producer = handle.clone();
        let producer = GlobalExecutor::spawn(async move {
            for i in 1..=10 {
                tokio::time::sleep(tokio::time::Duration::from_millis(5)).await;
                handle_producer
                    .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                    .unwrap();
            }
        });

        // Spawn multiple consumers
        let mut consumers = Vec::new();
        for _ in 0..3 {
            let h = handle.clone();
            let counter = Arc::clone(&fragments_received);
            consumers.push(GlobalExecutor::spawn(async move {
                let mut stream = h.stream();
                let mut local_count = 0;
                while let Some(result) = stream.next().await {
                    result.unwrap();
                    local_count += 1;
                }
                counter.fetch_add(local_count, Ordering::SeqCst);
                local_count
            }));
        }

        producer.await.unwrap();

        // Each consumer should read all 10 fragments
        for consumer in consumers {
            let count = consumer.await.unwrap();
            assert_eq!(count, 10);
        }
    }

    #[test]
    fn test_stream_handle_debug() {
        let handle = StreamHandle::new(make_stream_id(), 100);
        handle
            .push_fragment(1, Bytes::from_static(b"test"))
            .unwrap();

        let debug_str = format!("{:?}", handle);
        assert!(debug_str.contains("StreamHandle"));
        assert!(debug_str.contains("total_bytes"));
        assert!(debug_str.contains("complete"));
    }

    #[tokio::test]
    async fn test_try_assemble_before_complete() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push only first fragment
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // try_assemble should return None (incomplete)
        assert!(handle.try_assemble().is_none());

        // Push second fragment
        handle
            .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        // Now try_assemble should work
        let assembled = handle.try_assemble();
        assert!(assembled.is_some());
        assert_eq!(assembled.unwrap().len(), total as usize);
    }

    #[test]
    fn test_stream_error_display() {
        let cancelled = StreamError::Cancelled;
        assert_eq!(format!("{}", cancelled), "stream was cancelled");

        let not_found = StreamError::NotFound;
        assert_eq!(format!("{}", not_found), "stream not found in registry");

        let invalid = StreamError::InvalidFragment {
            message: "test error".into(),
        };
        assert_eq!(format!("{}", invalid), "invalid fragment: test error");

        let timeout = StreamError::InactivityTimeout;
        assert_eq!(
            format!("{}", timeout),
            "no fragments received within inactivity timeout"
        );
    }

    // ==================== Auto-Reclaim Tests ====================

    #[tokio::test]
    async fn test_stream_with_reclaim_basic() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Push all fragments
        for i in 1..=3 {
            handle
                .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
        }

        assert_eq!(handle.buffer.inserted_count(), 3);

        // Use reclaiming stream
        let mut stream = handle.stream_with_reclaim();
        assert!(stream.is_auto_reclaim());

        // Read first fragment - should be removed from buffer
        let chunk = stream.next().await.unwrap().unwrap();
        assert_eq!(chunk[0], 1);
        assert_eq!(handle.buffer.inserted_count(), 2); // One removed

        // Read remaining fragments
        let _ = stream.next().await.unwrap().unwrap();
        let _ = stream.next().await.unwrap().unwrap();

        assert_eq!(handle.buffer.inserted_count(), 0); // All removed

        // Stream is exhausted
        assert!(stream.next().await.is_none());
    }

    #[tokio::test]
    async fn test_stream_without_reclaim_preserves_data() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        for i in 1..=2 {
            handle
                .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
        }

        // Use non-reclaiming stream (default)
        let mut stream = handle.stream();
        assert!(!stream.is_auto_reclaim());

        // Read all fragments
        let _ = stream.next().await.unwrap().unwrap();
        let _ = stream.next().await.unwrap().unwrap();

        // Fragments still present in buffer
        assert_eq!(handle.buffer.inserted_count(), 2);
    }

    #[tokio::test]
    async fn test_reclaim_incremental_with_delayed_fragments() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);
        let mut stream = handle.stream_with_reclaim();

        // Push and consume fragment 1
        handle
            .push_fragment(1, Bytes::from(vec![1u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        let chunk = stream.next().await.unwrap().unwrap();
        assert_eq!(chunk[0], 1);
        assert_eq!(handle.buffer.inserted_count(), 0);

        // Push fragments 2 and 3
        handle
            .push_fragment(2, Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();
        handle
            .push_fragment(3, Bytes::from(vec![3u8; FRAGMENT_PAYLOAD_SIZE]))
            .unwrap();

        assert_eq!(handle.buffer.inserted_count(), 2);

        // Consume remaining
        let _ = stream.next().await.unwrap().unwrap();
        assert_eq!(handle.buffer.inserted_count(), 1);

        let _ = stream.next().await.unwrap().unwrap();
        assert_eq!(handle.buffer.inserted_count(), 0);
    }

    #[tokio::test]
    async fn test_reclaim_vs_fork_conflict() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        let total = (FRAGMENT_PAYLOAD_SIZE * 2) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        for i in 1..=2 {
            handle
                .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
        }

        // First consumer uses reclaim
        let mut reclaim_stream = handle.stream_with_reclaim();
        let _ = reclaim_stream.next().await.unwrap().unwrap(); // Takes fragment 1

        // Second consumer (forked) tries to read - fragment 1 is gone!
        let forked = handle.fork();
        let _forked_stream = forked.stream();

        // Fragment 1 was taken, so forked stream gets None (waits for it)
        // This demonstrates why auto_reclaim should only be used with single consumers
        // The forked stream would wait forever for fragment 1 which is gone

        // Fragment 2 is still there though
        assert!(handle.buffer.get(2).is_some());
    }

    #[tokio::test]
    async fn test_stream_memory_efficiency() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        // Simulate a large stream
        let num_fragments = 20u32;
        let total = (FRAGMENT_PAYLOAD_SIZE * num_fragments as usize) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);
        let handle_clone = handle.clone();

        // Producer pushes all fragments
        for i in 1..=num_fragments {
            handle
                .push_fragment(i, Bytes::from(vec![i as u8; FRAGMENT_PAYLOAD_SIZE]))
                .unwrap();
        }

        assert_eq!(handle_clone.buffer.inserted_count(), num_fragments as usize);

        // Consumer with reclaim
        let mut stream = handle_clone.stream_with_reclaim();

        // Read half the fragments
        for _ in 0..10 {
            let _ = stream.next().await.unwrap().unwrap();
        }

        // Memory for first half is freed
        assert_eq!(handle.buffer.inserted_count(), 10);

        // Read remaining
        for _ in 0..10 {
            let _ = stream.next().await.unwrap().unwrap();
        }

        // All memory freed
        assert_eq!(handle.buffer.inserted_count(), 0);
    }

    /// Regression test for the overflow fragment race condition.
    ///
    /// When fragment #1 has reduced payload (metadata embedded), the sender uses
    /// one extra "overflow" fragment. The buffer's is_complete() fires when
    /// min_complete_fragments are contiguous, but the overflow may not have
    /// arrived yet. Previously, assemble() returned an error immediately instead
    /// of waiting for the overflow fragment, causing "assembly failed" errors
    /// in production (blank pages in River web app).
    #[tokio::test]
    async fn test_assemble_waits_for_overflow_fragment() {
        use super::super::streaming_buffer::FRAGMENT_PAYLOAD_SIZE;

        // Simulate a stream where fragment #1 has reduced payload (metadata embedded).
        // Total data = 3 full fragments worth, but fragment #1 carries less data,
        // so 4 fragments are needed (3 base + 1 overflow actually used).
        let total = (FRAGMENT_PAYLOAD_SIZE * 3) as u64;
        let handle = StreamHandle::new(make_stream_id(), total);

        // Simulate what happens when metadata is embedded in fragment #1:
        // Fragment #1 carries less than FRAGMENT_PAYLOAD_SIZE bytes
        let reduced_payload = FRAGMENT_PAYLOAD_SIZE / 2;
        let frag1_data = Bytes::from(vec![1u8; reduced_payload]);

        // Remaining data distributed across fragments 2, 3, and 4 (overflow)
        let remaining = total as usize - reduced_payload;
        let frag2_data = Bytes::from(vec![2u8; FRAGMENT_PAYLOAD_SIZE]);
        let frag3_data = Bytes::from(vec![3u8; FRAGMENT_PAYLOAD_SIZE]);
        let overflow_size = remaining - 2 * FRAGMENT_PAYLOAD_SIZE;
        let frag4_data = Bytes::from(vec![4u8; overflow_size]);

        // Push base fragments (1, 2, 3) - this triggers is_complete()
        // since min_complete_fragments = 3
        handle.push_fragment(1, frag1_data).unwrap();
        handle.push_fragment(2, frag2_data).unwrap();
        handle.push_fragment(3, frag3_data).unwrap();

        // At this point, is_complete() returns true but assemble() would fail
        // because we only have reduced_payload + 2*FPS < total bytes
        assert!(handle.is_complete());
        assert!(handle.try_assemble().is_none()); // Not enough data yet

        // Start assemble in background - it should NOT fail, it should wait
        let handle_clone = handle.clone();
        let assemble_task = GlobalExecutor::spawn(async move { handle_clone.assemble().await });

        // Give it time to discover is_complete() == true but assemble() == None
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;

        // Task should still be running (waiting for overflow), not failed
        assert!(!assemble_task.is_finished());

        // Now push the overflow fragment
        handle.push_fragment(4, frag4_data).unwrap();

        // Assemble should now complete successfully
        let result = tokio::time::timeout(tokio::time::Duration::from_millis(200), assemble_task)
            .await
            .expect("timeout waiting for assemble")
            .expect("join error");

        assert!(result.is_ok(), "assemble should succeed: {:?}", result);
        let data = result.unwrap();
        assert_eq!(data.len(), total as usize);
    }
}