ringline-quic 0.3.0

use std::collections::VecDeque;
use std::net::SocketAddr;
use std::time::Instant;

use bytes::BytesMut;
use quinn_proto::{
    ClientConfig, ConnectionError, ConnectionHandle, DatagramEvent, Dir, Event, StreamEvent,
    StreamId, WriteError,
};
use slab::Slab;

use crate::config::QuicConfig;
use crate::error::Error;
use crate::event::{QuicConnId, QuicEvent};

/// A sans-IO QUIC endpoint.
///
/// Wraps [`quinn_proto::Endpoint`] and exposes an event-queue API.
/// This crate has no runtime dependency — callers are responsible for
/// sending outgoing packets (via [`poll_send`](Self::poll_send)) and
/// feeding incoming datagrams.
///
/// # Usage
///
/// 1. Feed incoming UDP datagrams via [`handle_datagram`](Self::handle_datagram).
/// 2. Drive connection timers via [`drive_timers`](Self::drive_timers).
/// 3. Poll application events via [`poll_event`](Self::poll_event).
/// 4. Drain outgoing packets via [`poll_send`](Self::poll_send).
pub struct QuicEndpoint {
    endpoint: quinn_proto::Endpoint,
    connections: Slab<QuicConnection>,
    /// Maps `ConnectionHandle.0` → slab key.  Grows as needed.
    handle_map: Vec<Option<u32>>,
    /// Application-facing event queue.
    events: VecDeque<QuicEvent>,
    /// Outgoing UDP packets waiting to be sent.
    send_queue: VecDeque<OutgoingPacket>,
    /// Scratch buffer for `poll_transmit`.
    transmit_buf: Vec<u8>,
    /// Scratch buffer for `endpoint.handle()` responses.
    response_buf: Vec<u8>,
    local_addr: SocketAddr,
    client_config: Option<ClientConfig>,
    send_queue_capacity: usize,
    max_transmit_datagrams: usize,
    /// When > 0, `drain_transmits` is a no-op. Set by [`BatchGuard`]
    /// during a batched run so quinn-proto can accumulate multiple
    /// packets before being asked to transmit — without this, each
    /// `stream_send` / `stream_finish` / `open_bi` triggers an
    /// immediate `poll_transmit` which produces one datagram and
    /// defeats GSO coalescing on the way out to the kernel. Nested
    /// so re-entrant batches still work correctly.
    batch_depth: u32,
}

struct QuicConnection {
    handle: ConnectionHandle,
    conn: quinn_proto::Connection,
    established: bool,
    outbound: bool,
    /// Whether 0-RTT was possible at the time this connection was inserted.
    /// On the client side that means we may have written application data
    /// using 0-RTT keys. We use this flag at handshake completion to
    /// decide whether to emit `QuicEvent::ZeroRttRejected` when the peer
    /// declined to accept it.
    zero_rtt_attempted: bool,
    /// Last-observed peer address. Compared against
    /// `conn.remote_address()` after each `poll_connection` round so we
    /// can emit `QuicEvent::PeerAddressChanged` on path migration.
    cached_remote: SocketAddr,
    /// Set to true once we have emitted `QuicEvent::ConnectionClosed` for
    /// this connection.  Guards against a second emission if quinn-proto
    /// somehow fires `Event::ConnectionLost` after `close_connection()` has
    /// already injected the event (defensive — the drain path should prevent
    /// this, but belt-and-suspenders).
    close_event_emitted: bool,
}

/// One outbound UDP packet produced by the endpoint.
///
/// `data` may carry a single datagram or — when `segment_size` is
/// `Some` — a Generic Segmentation Offload (GSO) buffer holding several
/// equally-sized datagrams concatenated together (the last may be
/// shorter). Runtime adapters that support `UDP_SEGMENT` should hand the
/// whole buffer to the kernel in one syscall (e.g. via
/// `UdpCtx::send_to_gso`); adapters without GSO support, or in-process
/// tests that feed datagrams back into [`QuicEndpoint::handle_datagram`],
/// can call [`datagrams()`](Self::datagrams) to iterate the per-datagram
/// slices.
pub struct OutgoingPacket {
    pub destination: SocketAddr,
    pub data: Vec<u8>,
    pub segment_size: Option<u16>,
}

impl OutgoingPacket {
    /// Iterate over the individual datagram slices contained in `data`.
    ///
    /// When `segment_size` is `None`, yields a single slice over the
    /// whole buffer. When it's `Some(s)`, yields successive `s`-byte
    /// slices (the last may be shorter).
    pub fn datagrams(&self) -> Datagrams<'_> {
        let seg = self
            .segment_size
            .map(|s| s as usize)
            .unwrap_or(self.data.len().max(1));
        Datagrams {
            data: &self.data,
            segment_size: seg,
            offset: 0,
        }
    }
}

/// Iterator over the per-datagram slices of an [`OutgoingPacket`].
pub struct Datagrams<'a> {
    data: &'a [u8],
    segment_size: usize,
    offset: usize,
}

impl<'a> Iterator for Datagrams<'a> {
    type Item = &'a [u8];

    fn next(&mut self) -> Option<&'a [u8]> {
        if self.offset >= self.data.len() {
            return None;
        }
        let end = (self.offset + self.segment_size).min(self.data.len());
        let chunk = &self.data[self.offset..end];
        self.offset = end;
        Some(chunk)
    }
}

impl QuicEndpoint {
    /// Create a new QUIC endpoint.
    ///
    /// `local_addr` is the address of the UDP socket this endpoint is bound to.
    pub fn new(config: QuicConfig, local_addr: SocketAddr) -> Self {
        let endpoint = quinn_proto::Endpoint::new(
            config.endpoint_config,
            config.server_config,
            config.allow_mtud,
            config.rng_seed,
        );

        Self {
            endpoint,
            connections: Slab::new(),
            handle_map: Vec::new(),
            events: VecDeque::new(),
            send_queue: VecDeque::new(),
            transmit_buf: Vec::with_capacity(1500),
            response_buf: Vec::with_capacity(1500),
            local_addr,
            client_config: config.client_config,
            send_queue_capacity: config.send_queue_capacity,
            max_transmit_datagrams: config.max_transmit_datagrams.max(1),
            batch_depth: 0,
        }
    }

    /// Enter a batched-send scope.
    ///
    /// While the returned [`BatchGuard`] is alive, calls to
    /// `stream_send` / `stream_finish` / `open_bi` / etc. **do not**
    /// trigger an internal `drain_transmits` — quinn-proto's
    /// `poll_transmit` is not invoked for each individual operation.
    /// On guard drop the deferred work is flushed in one shot via
    /// [`flush`](Self::flush), giving quinn-proto a chance to coalesce
    /// the batch's outgoing datagrams into a single GSO segment (which
    /// runtime adapters with `UDP_SEGMENT` support can hand to the
    /// kernel in one syscall).
    ///
    /// Without batching, every stream-altering call produces ~one UDP
    /// datagram on the wire, defeating GSO coalescing entirely. With
    /// batching, a tight loop that opens N streams, sends a request
    /// on each, and finishes can produce a single GSO buffer
    /// containing the resulting N initial packets — turning N small
    /// `sendmsg` syscalls into one.
    ///
    /// Nested batches compose: only the outermost guard's drop
    /// performs the flush.
    pub fn batch(&mut self) -> BatchGuard<'_> {
        self.batch_depth = self.batch_depth.saturating_add(1);
        BatchGuard { endpoint: self }
    }

    /// Feed an incoming UDP datagram to the QUIC state machine.
    pub fn handle_datagram(&mut self, now: Instant, data: &[u8], peer: SocketAddr) {
        let data = BytesMut::from(data);
        let event = self.endpoint.handle(
            now,
            peer,
            Some(self.local_addr.ip()),
            None, // ECN not yet supported
            data,
            &mut self.response_buf,
        );

        match event {
            Some(DatagramEvent::ConnectionEvent(ch, event)) => {
                if let Some(&Some(key)) = self.handle_map.get(ch.0) {
                    let key = key as usize;
                    self.connections[key].conn.handle_event(event);
                    self.poll_connection(key, now);
                }
            }
            Some(DatagramEvent::NewConnection(incoming)) => {
                let result = self.endpoint.accept(
                    incoming,
                    now,
                    &mut self.response_buf,
                    None, // use default server config
                );
                match result {
                    Ok((ch, conn)) => {
                        let key = self.insert_connection(ch, conn, false);
                        self.drain_transmits(key, now);
                        self.poll_connection(key, now);
                    }
                    Err(quinn_proto::AcceptError {
                        response: Some(t), ..
                    }) => {
                        // Quinn produced an explicit close response (e.g. a
                        // CONNECTION_REFUSED initial-close on handshake
                        // failure). Forward it so the peer doesn't have to
                        // wait for its own handshake timeout.
                        let data = self.response_buf[..t.size].to_vec();
                        // Stateless close response: if the send queue is full
                        // the peer just times out on its own handshake. Best
                        // effort — no in-flight state to keep consistent.
                        let _ = self.queue_packet(t.destination, data, None);
                    }
                    Err(_) => {
                        // No response packet — nothing to send back.
                    }
                }
            }
            Some(DatagramEvent::Response(transmit)) => {
                // Stateless response (e.g. version negotiation, retry).
                let data = self.response_buf[..transmit.size].to_vec();
                let _ = self.queue_packet(transmit.destination, data, None);
            }
            None => {}
        }
    }

    /// Fire expired per-connection timeouts.
    pub fn drive_timers(&mut self, now: Instant) {
        // Collect keys to avoid borrow conflict with poll_connection.
        let keys: Vec<u32> = self.connections.iter().map(|(k, _)| k as u32).collect();

        for key in keys {
            let key = key as usize;
            if !self.connections.contains(key) {
                continue;
            }
            if let Some(timeout) = self.connections[key].conn.poll_timeout()
                && timeout <= now
            {
                self.connections[key].conn.handle_timeout(now);
                self.drain_transmits(key, now);
                self.poll_connection(key, now);
            }
        }
    }

    /// Poll the next application event.
    ///
    /// Returns `None` when no more events are queued.
    pub fn poll_event(&mut self) -> Option<QuicEvent> {
        self.events.pop_front()
    }

    /// Poll the next outgoing UDP packet.
    ///
    /// Returns an [`OutgoingPacket`] (which may be a GSO-packed buffer
    /// of several datagrams when `segment_size` is `Some`) or `None`
    /// when the send queue is empty. The caller is responsible for
    /// sending the packet via their UDP socket — runtime adapters with
    /// GSO support should hand the whole buffer to the kernel in one
    /// syscall; others should iterate via
    /// [`OutgoingPacket::datagrams`].
    pub fn poll_send(&mut self) -> Option<OutgoingPacket> {
        self.send_queue.pop_front()
    }

    /// Drain pending transmits from every connection into the send queue.
    ///
    /// Call after a batch of [`stream_send`](Self::stream_send) /
    /// [`stream_finish`](Self::stream_finish) / [`open_bi`](Self::open_bi) /
    /// [`open_uni`](Self::open_uni) operations to make sure the resulting
    /// frames are exposed via [`poll_send`](Self::poll_send). Otherwise those
    /// frames stay buffered inside the connection until the next inbound
    /// datagram or expiring timer happens to flush them.
    pub fn flush(&mut self, now: Instant) {
        let keys: Vec<u32> = self.connections.iter().map(|(k, _)| k as u32).collect();
        for key in keys {
            let key = key as usize;
            if self.connections.contains(key) {
                self.drain_transmits(key, now);
            }
        }
    }

    /// Initiate an outbound QUIC connection.
    ///
    /// Returns a [`QuicConnId`] that will appear in a future
    /// [`QuicEvent::Connected`] event once the handshake completes.
    pub fn connect(
        &mut self,
        now: Instant,
        peer: SocketAddr,
        server_name: &str,
    ) -> Result<QuicConnId, Error> {
        let client_config = self.client_config.clone().ok_or(Error::ConnectionClosed)?;

        let (ch, conn) = self
            .endpoint
            .connect(now, client_config, peer, server_name)?;

        let key = self.insert_connection(ch, conn, true);
        self.drain_transmits(key, now);
        Ok(QuicConnId(key as u32))
    }

    /// Write data to a QUIC stream.
    ///
    /// Returns the number of bytes written (may be less than `data.len()` due to
    /// flow control).
    pub fn stream_send(
        &mut self,
        conn: QuicConnId,
        stream: StreamId,
        data: &[u8],
    ) -> Result<usize, Error> {
        let c = self.get_conn_mut(conn)?;
        match c.conn.send_stream(stream).write(data) {
            Ok(n) => Ok(n),
            Err(WriteError::Blocked) if c.conn.is_closed() => {
                // quinn-proto's `write_source` short-circuits to `Blocked`
                // once the connection is in Closed/Draining — but no
                // `StreamWritable` event will ever fire to unblock us.
                // Translate to a distinct error so callers don't retry.
                Err(Error::ConnectionClosing)
            }
            Err(e) => Err(Error::Write(e)),
        }
    }

    /// Send scatter-gather data on a QUIC stream without copying.
    ///
    /// Each [`Bytes`](bytes::Bytes) in `chunks` is handed to quinn-proto's
    /// [`SendStream::write_chunks`](quinn_proto::SendStream::write_chunks);
    /// partial chunks are advanced in place (their refcount drops the
    /// consumed prefix) and fully-consumed chunks are replaced with an
    /// empty `Bytes`. Returns the total bytes written across all chunks.
    ///
    /// Use this when you already have refcounted buffers (e.g. from
    /// `Bytes::from(Vec<u8>)`) and want to avoid the copy that
    /// [`stream_send`](Self::stream_send) would pay.
    pub fn stream_send_chunks(
        &mut self,
        conn: QuicConnId,
        stream: StreamId,
        chunks: &mut [bytes::Bytes],
    ) -> Result<usize, Error> {
        let c = self.get_conn_mut(conn)?;
        match c.conn.send_stream(stream).write_chunks(chunks) {
            Ok(written) => Ok(written.bytes),
            Err(WriteError::Blocked) if c.conn.is_closed() => Err(Error::ConnectionClosing),
            Err(e) => Err(Error::Write(e)),
        }
    }

    /// Read data from a QUIC stream into `buf`.
    ///
    /// Returns `(bytes_read, is_finished)`. When `is_finished` is true, the peer
    /// has finished sending on this stream.
    pub fn stream_recv(
        &mut self,
        conn: QuicConnId,
        stream: StreamId,
        buf: &mut [u8],
    ) -> Result<(usize, bool), Error> {
        let key = conn.0 as usize;
        if !self.connections.contains(key) || self.connections[key].close_event_emitted {
            return Err(Error::InvalidConnection);
        }
        let c = &mut self.connections[key];
        let mut recv = c.conn.recv_stream(stream);
        let mut chunks = recv.read(true)?;
        let mut total = 0;
        let mut finished = false;
        let mut read_err: Option<quinn_proto::ReadError> = None;

        while total < buf.len() {
            match chunks.next(buf.len() - total) {
                Ok(Some(chunk)) => {
                    let len = chunk.bytes.len();
                    buf[total..total + len].copy_from_slice(&chunk.bytes);
                    total += len;
                }
                Ok(None) => {
                    finished = true;
                    break;
                }
                Err(quinn_proto::ReadError::Blocked) => break,
                Err(e) => {
                    read_err = Some(e);
                    break;
                }
            }
        }
        // `finalize` signals that MAX_DATA / MAX_STREAM_DATA frames are now
        // queued and should be flushed soon. Without honoring the hint the
        // peer waits an extra tick before its flow-control window opens —
        // measurable on quiescent connections that don't hit `flush()` or
        // receive other datagrams quickly.
        let should_transmit = chunks.finalize().should_transmit();
        if should_transmit {
            self.drain_transmits(key, Instant::now());
        }
        if let Some(e) = read_err {
            return Err(Error::Read(e));
        }
        Ok((total, finished))
    }

    /// Read stream data, appending up to `max` bytes onto `buf`.
    ///
    /// Returns `(bytes_appended, is_finished)`. Identical semantics to
    /// [`stream_recv`](Self::stream_recv) but writes directly into the
    /// caller's `BytesMut` instead of a `&mut [u8]` scratch slice,
    /// saving one memcpy per call on hot paths that accumulate into a
    /// `BytesMut` anyway (e.g. ringline-h3's per-stream frame
    /// reassembly buffer). The freshly-appended region can be
    /// `freeze()`d and sliced as refcounted `Bytes` for zero-copy
    /// frame payload extraction.
    pub fn stream_recv_into(
        &mut self,
        conn: QuicConnId,
        stream: StreamId,
        buf: &mut BytesMut,
        max: usize,
    ) -> Result<(usize, bool), Error> {
        let key = conn.0 as usize;
        if !self.connections.contains(key) || self.connections[key].close_event_emitted {
            return Err(Error::InvalidConnection);
        }
        let c = &mut self.connections[key];
        let mut recv = c.conn.recv_stream(stream);
        let mut chunks = recv.read(true)?;
        let mut total = 0;
        let mut finished = false;
        let mut read_err: Option<quinn_proto::ReadError> = None;

        while total < max {
            match chunks.next(max - total) {
                Ok(Some(chunk)) => {
                    total += chunk.bytes.len();
                    buf.extend_from_slice(&chunk.bytes);
                }
                Ok(None) => {
                    finished = true;
                    break;
                }
                Err(quinn_proto::ReadError::Blocked) => break,
                Err(e) => {
                    read_err = Some(e);
                    break;
                }
            }
        }
        let should_transmit = chunks.finalize().should_transmit();
        if should_transmit {
            self.drain_transmits(key, Instant::now());
        }
        if let Some(e) = read_err {
            return Err(Error::Read(e));
        }
        Ok((total, finished))
    }

    /// Send FIN on a stream, indicating no more data will be sent.
    ///
    /// Returns `Err(StreamStopped(code))` if the peer issued STOP_SENDING
    /// with `code` before we finished, `Err(StreamClosed)` if the stream
    /// was already finished or reset locally. The frame is queued
    /// immediately — callers don't need to `flush()`.
    pub fn stream_finish(&mut self, conn: QuicConnId, stream: StreamId) -> Result<(), Error> {
        let key = conn.0 as usize;
        let c = self.get_conn_mut(conn)?;
        match c.conn.send_stream(stream).finish() {
            Ok(()) => {}
            Err(quinn_proto::FinishError::Stopped(code)) => return Err(Error::StreamStopped(code)),
            Err(quinn_proto::FinishError::ClosedStream) => return Err(Error::StreamClosed),
        }
        self.drain_transmits(key, Instant::now());
        Ok(())
    }

    /// Ask the peer to stop sending on `stream` (sends a STOP_SENDING frame).
    ///
    /// On the peer the corresponding send stream will surface a
    /// [`QuicEvent::StreamStopped`] with the supplied `error_code`. Locally,
    /// the recv side of `stream` is considered done once this returns.
    /// Returns `Err(StreamClosed)` if the recv side was already finished.
    pub fn stop_sending(
        &mut self,
        conn: QuicConnId,
        stream: StreamId,
        error_code: quinn_proto::VarInt,
    ) -> Result<(), Error> {
        let key = conn.0 as usize;
        let c = self.get_conn_mut(conn)?;
        match c.conn.recv_stream(stream).stop(error_code) {
            Ok(()) => {}
            Err(quinn_proto::ClosedStream { .. }) => return Err(Error::StreamClosed),
        }
        self.drain_transmits(key, Instant::now());
        Ok(())
    }

    /// Reset `stream`, indicating to the peer that no more data will be
    /// sent. The peer's recv side will surface a
    /// [`quinn_proto::ReadError::Reset`] on its next `stream_recv`.
    /// Returns `Err(StreamClosed)` if the send side was already finished
    /// or reset.
    pub fn reset_stream(
        &mut self,
        conn: QuicConnId,
        stream: StreamId,
        error_code: quinn_proto::VarInt,
    ) -> Result<(), Error> {
        let key = conn.0 as usize;
        let c = self.get_conn_mut(conn)?;
        match c.conn.send_stream(stream).reset(error_code) {
            Ok(()) => {}
            Err(quinn_proto::ClosedStream { .. }) => return Err(Error::StreamClosed),
        }
        self.drain_transmits(key, Instant::now());
        Ok(())
    }

    /// Open a bidirectional stream.
    ///
    /// Returns `None` if the peer's stream concurrency limit has been reached.
    pub fn open_bi(&mut self, conn: QuicConnId) -> Result<Option<StreamId>, Error> {
        let c = self.get_conn_mut(conn)?;
        Ok(c.conn.streams().open(Dir::Bi))
    }

    /// Earliest quinn-proto timer deadline across all live connections,
    /// or `None` if no connection currently has a pending timer.
    ///
    /// Use this to size an adaptive sleep: callers that select between
    /// "next datagram" and "wake periodically to service timers" should
    /// sleep until this deadline rather than picking an arbitrary
    /// interval. Sleeping longer than the deadline lets quinn-proto
    /// fall behind on PTO / loss-recovery timers and stalls the
    /// connection; sleeping much shorter wastes wake cycles when no
    /// packets are flowing.
    pub fn next_timer_deadline(&mut self) -> Option<Instant> {
        let mut earliest: Option<Instant> = None;
        for key in 0..self.connections.capacity() {
            if !self.connections.contains(key) {
                continue;
            }
            if let Some(t) = self.connections[key].conn.poll_timeout() {
                earliest = Some(match earliest {
                    Some(prev) => prev.min(t),
                    None => t,
                });
            }
        }
        earliest
    }

    /// Snapshot of quinn-proto connection statistics (RTT, CWND, frame
    /// counts, loss). Useful for diagnostics and for adaptive policies
    /// that want to react to congestion-control state.
    ///
    /// The returned [`quinn_proto::ConnectionStats`] is a `Copy` value
    /// — callers can hold it indefinitely without keeping a borrow on
    /// the endpoint.
    pub fn connection_stats(
        &self,
        conn: QuicConnId,
    ) -> Result<quinn_proto::ConnectionStats, Error> {
        let key = conn.0 as usize;
        if !self.connections.contains(key) {
            return Err(Error::InvalidConnection);
        }
        Ok(self.connections[key].conn.stats())
    }

    /// Open a unidirectional stream.
    ///
    /// Returns `None` if the peer's stream concurrency limit has been reached.
    pub fn open_uni(&mut self, conn: QuicConnId) -> Result<Option<StreamId>, Error> {
        let c = self.get_conn_mut(conn)?;
        Ok(c.conn.streams().open(Dir::Uni))
    }

    /// Close a QUIC connection with the given error code and reason.
    ///
    /// The resulting `CONNECTION_CLOSE` packet is queued for transmission
    /// before this call returns, so the peer is told about the closure as
    /// soon as the caller flushes its UDP socket. There is no need for the
    /// caller to also call [`flush`](Self::flush).
    pub fn close_connection(&mut self, conn: QuicConnId, code: u32, reason: &[u8]) {
        let key = conn.0 as usize;
        if !self.connections.contains(key) {
            return;
        }
        // Repeat calls are a no-op — quinn-proto's `close` is itself a no-op
        // once the connection has been closed, and emitting `ConnectionClosed`
        // twice tears the bookkeeping in higher layers (h3 in particular
        // would clear `pending_sends` / `request_streams` again).
        if self.connections[key].close_event_emitted {
            return;
        }
        let now = Instant::now();
        self.connections[key].conn.close(
            now,
            quinn_proto::VarInt::from_u32(code),
            bytes::Bytes::copy_from_slice(reason),
        );
        // Drain the CONNECTION_CLOSE packet (and any final acks) so callers
        // don't have to remember to flush after closing.
        self.drain_transmits(key, now);
        // Quinn-proto's graceful-close path never fires Event::ConnectionLost
        // for the initiating side (the connection enters Draining without
        // setting the error field that poll() checks). Emit the event here so
        // higher-level layers (e.g. H3Connection) can clean up per-connection
        // state without waiting for the 3×PTO drain timer.
        let conn_id = QuicConnId(key as u32);
        self.connections[key].close_event_emitted = true;
        self.events.push_back(QuicEvent::ConnectionClosed {
            conn: conn_id,
            reason: ConnectionError::LocallyClosed,
        });
    }

    /// Send an unreliable QUIC datagram (RFC 9221).
    ///
    /// Datagrams are not retransmitted, not flow-controlled per-stream,
    /// and may arrive out of order. They're useful for applications where
    /// timeliness beats reliability — telemetry, gaming, real-time media.
    ///
    /// If `drop_old_when_full` is true and the local outgoing buffer is
    /// full, older queued datagrams are dropped to make room. If false,
    /// returns an error and the caller should wait for the buffer to
    /// drain (peer acks clear queued datagrams) before retrying.
    ///
    /// Returns:
    /// - `Err(InvalidConnection)` if the connection id is unknown.
    /// - `Err(DatagramDisabled)` if local config has datagrams off.
    /// - `Err(DatagramUnsupportedByPeer)` if the peer didn't enable them.
    /// - `Err(DatagramTooLarge)` if `data` exceeds the peer's
    ///   `max_datagram_size`.
    /// - `Err(DatagramBlocked(data))` if the outgoing buffer is full and
    ///   `drop_old_when_full` is false — the payload is returned so the
    ///   caller can retry on `QuicEvent::DatagramsUnblocked` without an
    ///   extra clone / allocation.
    pub fn send_datagram(
        &mut self,
        conn: QuicConnId,
        data: bytes::Bytes,
        drop_old_when_full: bool,
    ) -> Result<(), Error> {
        let c = self.get_conn_mut(conn)?;
        match c.conn.datagrams().send(data, drop_old_when_full) {
            Ok(()) => {}
            Err(quinn_proto::SendDatagramError::Disabled) => return Err(Error::DatagramDisabled),
            Err(quinn_proto::SendDatagramError::UnsupportedByPeer) => {
                return Err(Error::DatagramUnsupportedByPeer);
            }
            Err(quinn_proto::SendDatagramError::TooLarge) => return Err(Error::DatagramTooLarge),
            Err(quinn_proto::SendDatagramError::Blocked(b)) => {
                return Err(Error::DatagramBlocked(b));
            }
        }
        // Make sure the resulting DATAGRAM frame actually leaves the
        // connection — see `close_connection` for the same lesson.
        let key = conn.0 as usize;
        let now = Instant::now();
        self.drain_transmits(key, now);
        Ok(())
    }

    /// Maximum size of a datagram that may be passed to `send_datagram`.
    ///
    /// Returns `None` if the peer hasn't enabled datagram support or
    /// the connection ID is unknown.
    pub fn max_datagram_size(&mut self, conn: QuicConnId) -> Option<usize> {
        self.connections
            .get_mut(conn.0 as usize)
            .and_then(|c| c.conn.datagrams().max_size())
    }

    /// Bytes still available in the outgoing datagram send buffer.
    pub fn datagram_send_buffer_space(&mut self, conn: QuicConnId) -> Option<usize> {
        self.connections
            .get_mut(conn.0 as usize)
            .map(|c| c.conn.datagrams().send_buffer_space())
    }

    /// Whether 0-RTT keys were derived for this connection (i.e. early
    /// data is/was possible).
    ///
    /// On the client side this becomes true immediately after `connect`
    /// if the configured rustls client has a session resumption ticket
    /// for the server name. On the server side it becomes true if the
    /// peer's Initial included 0-RTT material we could decrypt.
    ///
    /// Once `true`, any data written via `stream_send` /
    /// `stream_send_chunks` / `send_datagram` before the handshake
    /// completes is encrypted with 0-RTT keys and can arrive at the peer
    /// before the handshake finishes. After the handshake, check
    /// [`accepted_0rtt`](Self::accepted_0rtt) to see whether the peer
    /// kept the early data.
    pub fn has_0rtt(&self, conn: QuicConnId) -> bool {
        self.connections
            .get(conn.0 as usize)
            .is_some_and(|c| c.conn.has_0rtt())
    }

    /// Whether the peer accepted our 0-RTT data (client-side only).
    ///
    /// The value is meaningless until the connection has progressed past
    /// the handshake — usually once a `QuicEvent::Connected` has been
    /// observed for the connection. On the server side this is always
    /// false. If `has_0rtt` was true at connect time but this returns
    /// false after the handshake, the peer rejected 0-RTT and a
    /// `QuicEvent::ZeroRttRejected` event will have been emitted.
    pub fn accepted_0rtt(&self, conn: QuicConnId) -> bool {
        self.connections
            .get(conn.0 as usize)
            .is_some_and(|c| c.conn.accepted_0rtt())
    }

    /// Number of active QUIC connections.
    pub fn connection_count(&self) -> usize {
        self.connections.len()
    }

    /// Number of pending outgoing packets.
    pub fn send_queue_len(&self) -> usize {
        self.send_queue.len()
    }

    /// Peer address for a connection, if it exists.
    pub fn remote_addr(&self, conn: QuicConnId) -> Option<SocketAddr> {
        // A connection in drain (post-`ConnectionClosed`) is no longer
        // application-addressable even though the slab slot still exists
        // for quinn-proto's drain bookkeeping.
        self.connections
            .get(conn.0 as usize)
            .filter(|c| !c.close_event_emitted)
            .map(|c| c.conn.remote_address())
    }

    // ── Internal helpers ─────────────────────────────────────────────

    fn insert_connection(
        &mut self,
        ch: ConnectionHandle,
        conn: quinn_proto::Connection,
        outbound: bool,
    ) -> usize {
        // `has_0rtt()` returns whether 0-RTT keys were derived for this
        // connection. On the client side that's true when rustls had a
        // resumption ticket for the server name; on the server side it's
        // true once the client's Initial included 0-RTT material we could
        // decrypt. Capture it now so we can detect rejection later.
        let zero_rtt_attempted = conn.has_0rtt();
        let cached_remote = conn.remote_address();
        let key = self.connections.insert(QuicConnection {
            handle: ch,
            conn,
            established: false,
            outbound,
            zero_rtt_attempted,
            cached_remote,
            close_event_emitted: false,
        });

        // Grow handle_map if needed.
        let idx = ch.0;
        if idx >= self.handle_map.len() {
            self.handle_map.resize(idx + 1, None);
        }
        self.handle_map[idx] = Some(key as u32);

        key
    }

    fn get_conn_mut(&mut self, conn: QuicConnId) -> Result<&mut QuicConnection, Error> {
        let c = self
            .connections
            .get_mut(conn.0 as usize)
            .ok_or(Error::InvalidConnection)?;
        // Treat a connection that's mid-drain (post-`ConnectionClosed`) as
        // invalid for application operations: quinn-proto still owns the
        // slot for its 3×PTO drain window, but every stream op would either
        // fail or no-op, and callers should observe a clean error rather
        // than a confusing `Blocked` / `ClosedStream` return.
        if c.close_event_emitted {
            return Err(Error::InvalidConnection);
        }
        Ok(c)
    }

    /// Drain all pending transmits from a connection into the send queue.
    ///
    /// Asks quinn-proto for up to `max_transmit_datagrams` per call. If
    /// quinn returns a GSO-segmented buffer
    /// (`Transmit::segment_size = Some(s)` with `s < size`), it is
    /// queued as a single [`OutgoingPacket`] with `segment_size: Some(s)`
    /// — runtime adapters with `UDP_SEGMENT` (Linux GSO) support emit
    /// the whole buffer in one syscall, while others split per-datagram
    /// via [`OutgoingPacket::datagrams`].
    fn drain_transmits(&mut self, key: usize, now: Instant) {
        // Inside a `batch()` scope, defer until the guard's drop calls
        // `flush()`. This is what lets quinn-proto coalesce a batch of
        // stream operations into a single GSO segment instead of
        // emitting one datagram per stream_send / stream_finish /
        // open_bi.
        if self.batch_depth > 0 {
            return;
        }
        let max_datagrams = self.max_transmit_datagrams;
        loop {
            // Stop pulling from quinn the moment the send queue is full.
            // Earlier this method always called `poll_transmit` and silently
            // dropped on overflow, which left quinn's `bytes_in_flight` /
            // congestion-control state out of sync with what we'd actually
            // sent and risked losing CONNECTION_CLOSE packets on heavy socket
            // backpressure.
            if self.send_queue.len() >= self.send_queue_capacity {
                break;
            }
            self.transmit_buf.clear();
            let transmit = self.connections[key].conn.poll_transmit(
                now,
                max_datagrams,
                &mut self.transmit_buf,
            );

            match transmit {
                Some(t) => {
                    let segment_size = match t.segment_size {
                        Some(seg) if seg < t.size => u16::try_from(seg).ok(),
                        _ => None,
                    };
                    let data = self.transmit_buf[..t.size].to_vec();
                    if !self.queue_packet(t.destination, data, segment_size) {
                        // Capacity check above should keep us out of this
                        // branch, but if a future change loosens it the
                        // explicit break here keeps behavior safe.
                        break;
                    }
                }
                None => break,
            }
        }
    }

    /// Drain endpoint events and application events from a connection.
    fn poll_connection(&mut self, key: usize, now: Instant) {
        // 1. Drain endpoint events (e.g. connection ID updates).
        while let Some(event) = self.connections[key].conn.poll_endpoint_events() {
            if let Some(conn_event) = self
                .endpoint
                .handle_event(self.connections[key].handle, event)
            {
                self.connections[key].conn.handle_event(conn_event);
            }
        }

        // 2. Drain transmits generated by endpoint event handling.
        self.drain_transmits(key, now);

        // 3. Drain application events.
        let conn_id = QuicConnId(key as u32);
        while let Some(event) = self.connections[key].conn.poll() {
            match event {
                Event::Connected => {
                    self.connections[key].established = true;
                    if self.connections[key].outbound {
                        self.events.push_back(QuicEvent::Connected(conn_id));
                        // If we tried 0-RTT and the peer rejected it,
                        // surface that so the application can re-send
                        // the discarded data over 1-RTT.
                        if self.connections[key].zero_rtt_attempted
                            && !self.connections[key].conn.accepted_0rtt()
                        {
                            self.events
                                .push_back(QuicEvent::ZeroRttRejected { conn: conn_id });
                        }
                    } else {
                        self.events.push_back(QuicEvent::NewConnection(conn_id));
                    }
                }
                Event::ConnectionLost { reason } => {
                    if !self.connections[key].close_event_emitted {
                        self.events.push_back(QuicEvent::ConnectionClosed {
                            conn: conn_id,
                            reason,
                        });
                        self.connections[key].close_event_emitted = true;
                    }
                    // Do *not* remove the slot yet. quinn-proto's `Endpoint`
                    // still has CID-index and connection-slab entries that
                    // are only cleaned when the per-connection state machine
                    // emits `EndpointEventInner::Drained` via
                    // `poll_endpoint_events`. That happens once the 3×PTO
                    // drain timer fires and `is_drained()` becomes true; the
                    // reap path at step 6 below handles it. Dropping the
                    // wrapper-side connection here would orphan those quinn
                    // index entries for the lifetime of the endpoint.
                    return;
                }
                Event::Stream(stream_event) => match stream_event {
                    StreamEvent::Opened { dir } => {
                        // Accept all new streams from the peer and synthesise a
                        // StreamReadable for each. quinn-proto's on_stream_frame
                        // suppresses Readable when the opening STREAM frame also
                        // carries data (setting `opened` instead). Without this
                        // synthetic event, an application that only reads on
                        // StreamReadable would hang for short one-shot messages
                        // that fit in a single frame. A stream_recv on an empty
                        // stream is a no-op, so emitting this unconditionally
                        // is safe.
                        while let Some(stream) = self.connections[key].conn.streams().accept(dir) {
                            self.events.push_back(QuicEvent::StreamOpened {
                                conn: conn_id,
                                stream,
                                bidi: dir == Dir::Bi,
                            });
                            self.events.push_back(QuicEvent::StreamReadable {
                                conn: conn_id,
                                stream,
                            });
                        }
                    }
                    StreamEvent::Readable { id } => {
                        self.events.push_back(QuicEvent::StreamReadable {
                            conn: conn_id,
                            stream: id,
                        });
                    }
                    StreamEvent::Writable { id } => {
                        self.events.push_back(QuicEvent::StreamWritable {
                            conn: conn_id,
                            stream: id,
                        });
                    }
                    StreamEvent::Finished { id } => {
                        self.events.push_back(QuicEvent::StreamFinished {
                            conn: conn_id,
                            stream: id,
                        });
                    }
                    StreamEvent::Stopped { id, error_code } => {
                        self.events.push_back(QuicEvent::StreamStopped {
                            conn: conn_id,
                            stream: id,
                            error_code,
                        });
                    }
                    StreamEvent::Available { dir } => {
                        self.events
                            .push_back(QuicEvent::StreamsAvailable { conn: conn_id, dir });
                    }
                },
                Event::DatagramReceived => {
                    // Drain everything currently buffered. Quinn fires one
                    // event per arriving datagram, but draining
                    // unconditionally keeps us robust to ordering.
                    while let Some(data) = self.connections[key].conn.datagrams().recv() {
                        self.events.push_back(QuicEvent::DatagramReceived {
                            conn: conn_id,
                            data,
                        });
                    }
                }
                Event::DatagramsUnblocked => {
                    self.events
                        .push_back(QuicEvent::DatagramsUnblocked { conn: conn_id });
                }
                Event::HandshakeDataReady => {
                    self.events
                        .push_back(QuicEvent::HandshakeDataReady { conn: conn_id });
                }
            }
        }

        // 4. Final drain of transmits generated by event processing.
        self.drain_transmits(key, now);

        // 5. Detect peer-address change (path migration). Quinn-proto
        //    handles the PATH_CHALLENGE/PATH_RESPONSE handshake
        //    internally and updates `Connection::remote_address()` once
        //    the new path is validated. We compare against the
        //    last-observed value and surface a `PeerAddressChanged`
        //    event so applications can react.
        if self.connections.contains(key) {
            let conn_id = QuicConnId(key as u32);
            let current = self.connections[key].conn.remote_address();
            let previous = self.connections[key].cached_remote;
            if current != previous {
                self.connections[key].cached_remote = current;
                self.events.push_back(QuicEvent::PeerAddressChanged {
                    conn: conn_id,
                    previous,
                    current,
                });
            }
        }

        // 6. If the connection is drained, remove it.
        if self.connections.contains(key) && self.connections[key].conn.is_drained() {
            self.remove_connection(key);
        }
    }

    fn remove_connection(&mut self, key: usize) {
        let qc = self.connections.remove(key);
        let idx = qc.handle.0;
        if idx < self.handle_map.len() {
            self.handle_map[idx] = None;
        }
    }

    /// Append a packet to the outgoing queue. Returns `true` if it was queued,
    /// `false` if the queue was already at `send_queue_capacity` and the
    /// packet was dropped. Callers driving quinn-proto's `poll_transmit` loop
    /// should stop pulling on `false`; one-shot stateless responses (which
    /// don't track in-flight state) can ignore the return value.
    fn queue_packet(
        &mut self,
        destination: SocketAddr,
        data: Vec<u8>,
        segment_size: Option<u16>,
    ) -> bool {
        if self.send_queue.len() >= self.send_queue_capacity {
            return false;
        }
        self.send_queue.push_back(OutgoingPacket {
            destination,
            data,
            segment_size,
        });
        true
    }
}

/// RAII guard returned by [`QuicEndpoint::batch`]. Suppresses the
/// internal `drain_transmits` calls that normally fire after each
/// stream operation so quinn-proto can coalesce a whole batch into a
/// single GSO buffer. On drop, flushes the connections so the
/// accumulated work goes out as one burst.
///
/// Construct via `let g = endpoint.batch();` and dereference through
/// the guard to issue stream operations:
///
/// ```rust,ignore
/// {
///     let mut g = endpoint.batch();
///     for _ in 0..n {
///         let stream = g.open_bi(conn)?.unwrap();
///         g.stream_send(conn, stream, payload)?;
///         g.stream_finish(conn, stream)?;
///     }
///     // Drop here triggers a single flush — quinn-proto can
///     // coalesce the n datagrams into one GSO segment.
/// }
/// ```
pub struct BatchGuard<'a> {
    endpoint: &'a mut QuicEndpoint,
}

impl<'a> std::ops::Deref for BatchGuard<'a> {
    type Target = QuicEndpoint;
    fn deref(&self) -> &QuicEndpoint {
        self.endpoint
    }
}

impl<'a> std::ops::DerefMut for BatchGuard<'a> {
    fn deref_mut(&mut self) -> &mut QuicEndpoint {
        self.endpoint
    }
}

impl<'a> Drop for BatchGuard<'a> {
    fn drop(&mut self) {
        // Decrement first so the flush below actually performs the
        // drain. saturating_sub keeps the field correct even if a
        // future bug double-drops (it can't today — `BatchGuard` is
        // !Copy and constructed only by `batch()`).
        self.endpoint.batch_depth = self.endpoint.batch_depth.saturating_sub(1);
        if self.endpoint.batch_depth == 0 {
            // Use a single fresh `Instant::now` for the whole flush
            // so quinn-proto's transmission accounting sees one
            // consistent timestamp for the batch.
            self.endpoint.flush(std::time::Instant::now());
        }
    }
}