punktfunk/crates/punktfunk-core/src/packet.rs

//! Zero-copy wire framing: split an access unit into FEC blocks of MTU-sized shards,
//! and reassemble + FEC-recover them on the far side.
//!
//! ## Wire layout
//!
//! Each packet is a fixed [`PacketHeader`] followed by one FEC shard's payload. Fields
//! are host-endian for now (every target platform is little-endian); the `punktfunk/1` (P2)
//! spec will pin byte order explicitly when we talk to non-LE peers.
//!
//! ## GameStream mapping (P1)
//!
//! `frame_index`↔`frameIndex`, `stream_seq`↔`streamPacketIndex`,
//! (`block_index`,`block_count`)↔the `multiFecBlocks` nibbles, and
//! (`data_shards`,`recovery_shards`,`shard_index`)↔the `fecInfo` bitfield. We carry them
//! as explicit fields rather than bit-packing; full GameStream wire-exactness is a GameStream-host
//! concern (it also needs RTP framing + RTSP), this is the coherent internal format.

use crate::config::Config;
use crate::error::{PunktfunkError, Result};
use crate::fec::ErasureCoder;
use crate::session::Frame;
use crate::stats::StatsCounters;
use std::collections::{BTreeMap, HashMap, HashSet};
use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout};

/// Identifies a punktfunk video packet (vs. an input datagram, see [`crate::input`]).
pub const PUNKTFUNK_MAGIC: u8 = 0xC9;

// Frame flags (mirroring GameStream's FLAG_*).
pub const FLAG_PIC: u8 = 0x1;
pub const FLAG_EOF: u8 = 0x2;
pub const FLAG_SOF: u8 = 0x4;
/// Bandwidth-probe filler, not decodable video: a [`crate::quic::ProbeRequest`] speed test makes
/// the host burst access units carrying this flag so the client measures throughput/loss without
/// feeding them to the decoder. Punktfunk/1 only (GameStream never sets it).
pub const FLAG_PROBE: u8 = 0x8;

/// Application `user_flags` bit (the u32 [`PacketHeader::user_flags`] word, surfaced to the client
/// as [`crate::session::Frame::flags`]) — NOT a transport packet flag. Marks the access unit that
/// **completes an intra-refresh wave**: the picture is loss-free from here even though the frame is
/// a coded `P` (no IDR, so the decoder never sets `AV_FRAME_FLAG_KEY`). The client lifts its
/// post-loss display freeze on this bit as well as on a real keyframe — the only bitstream-invisible
/// clean point it can honor without forcing a full IDR. Lives above the low nibble because the host
/// reuses `FLAG_PIC`/`FLAG_SOF`/`FLAG_PROBE` bit values inside `user_flags`; `0x10` clears all four.
pub const USER_FLAG_RECOVERY_POINT: u32 = 0x10;

/// Application `user_flags` bit — a **definitive single-frame clean re-anchor**. Unlike
/// [`USER_FLAG_RECOVERY_POINT`] (an intra-refresh wave boundary, where the first boundary after a loss
/// is only half-healed so the client waits for the second), this marks an access unit the host coded
/// to reference a **known-good** picture on purpose — an AMD **LTR reference-frame-invalidation**
/// recovery frame (`ForceLTRReferenceBitfield`): a clean P-frame off a long-term reference the client
/// already has, not an IDR. The picture is loss-free the instant this AU decodes, so the client lifts
/// its post-loss freeze on the **first** such mark. Coded `P` (no IDR), so the decoder never sets
/// `AV_FRAME_FLAG_KEY` — this host flag is the only signal.
pub const USER_FLAG_RECOVERY_ANCHOR: u32 = 0x20;

/// Widest lost-frame range (frames, wrapping `last - first`) a reference-frame-invalidation
/// recovery may be asked to repair; anything wider goes straight to the keyframe path on BOTH
/// ends. RFI can only re-reference history the encoder still holds — NVENC keeps a 5-frame DPB,
/// AMD LTR ~1 s of marks — and a genuine loss this wide (>1 s even at 240 fps) has no valid
/// reference anywhere, so an RFI request for it is either hopeless or (worse) a phantom range
/// from a desynced counter. Shared by the host's RFI dispatch (range → keyframe fallback) and the
/// client-side gap detectors (huge gap → resync + keyframe request, no RFI).
pub const RFI_MAX_RANGE: u32 = 256;

/// Crypto framing overhead [`Session`](crate::session::Session) adds when encrypting:
/// an 8-byte sequence prefix plus the GCM tag.
pub const CRYPTO_OVERHEAD: usize = 8 + crate::crypto::TAG_LEN;

/// Largest UDP datagram the core will send or accept. `Config::validate` bounds
/// `shard_payload` so `HEADER_LEN + shard_payload + CRYPTO_OVERHEAD ≤ MAX_DATAGRAM_BYTES`.
pub const MAX_DATAGRAM_BYTES: usize = 2048;

/// How far behind the newest frame's capture pts an INCOMPLETE frame may sit before it is
/// declared lost (counted in `frames_dropped`, which triggers the client's recovery-keyframe
/// request). TIME-based, not frame-count-based, so the fuse is the same at every refresh rate: a
/// fixed index window is refresh-relative (4 frames = 66 ms at 60 fps but only 33 ms at 120 fps —
/// inside normal Wi-Fi retry/block-ack reorder timescales, where a delayed-not-lost shard can
/// trail newer frames). Observed live at 120 fps: the too-tight fuse declared merely-late frames
/// dead every few seconds, and each false loss cost a recovery-IDR burst + an inflated loss report
/// (FEC churn) — a self-sustaining latency/bitrate oscillation. 120 ms rides safely above radio
/// retry jitter while still detecting a real loss ~2× faster than the original 16-frame window did
/// at 60 fps.
const LOSS_WINDOW_NS: u64 = 120_000_000;

/// Hard cap on how many frame INDICES behind the newest an incomplete frame may sit, whatever its
/// pts claims — bounds the reassembler's memory against a corrupt/hostile pts (which
/// [`LOSS_WINDOW_NS`] alone would trust) and against pathologically high frame rates. At 120 fps,
/// 120 ms ≈ 14 indices, so 64 leaves ample slack up to ~500 fps.
const HARD_LOSS_WINDOW: u32 = 64;

/// How many frames behind the newest the reassembler remembers emitted/abandoned frame indices
/// (`completed`), so a straggler shard can neither resurrect an abandoned frame nor re-open an
/// emitted one. Must cover at least [`HARD_LOSS_WINDOW`]: stragglers can trickle in later than the
/// loss verdict.
const REORDER_WINDOW: u32 = 64;

/// Fixed per-packet header. `#[repr(C)]`, no padding, zero-copy (de)serializable.
#[repr(C)]
#[derive(Clone, Copy, Debug, FromBytes, IntoBytes, KnownLayout, Immutable)]
pub struct PacketHeader {
    pub pts_ns: u64,
    pub frame_index: u32,
    pub stream_seq: u32,
    pub frame_bytes: u32,
    pub user_flags: u32,
    pub block_index: u16,
    pub block_count: u16,
    pub data_shards: u16,
    pub recovery_shards: u16,
    pub shard_index: u16,
    pub shard_bytes: u16,
    pub magic: u8,
    pub version: u8,
    pub fec_scheme: u8,
    pub flags: u8,
}

/// Size of [`PacketHeader`] on the wire (40 bytes).
pub const HEADER_LEN: usize = std::mem::size_of::<PacketHeader>();

const _: () = assert!(HEADER_LEN == 40, "PacketHeader must be 40 bytes / unpadded");

// ---------------------------------------------------------------------------
// Host side: packetization
// ---------------------------------------------------------------------------

/// Splits encoded access units into FEC-protected shard packets. Host-side only.
///
/// Frame numbering: a caller can pass an **explicit** `frame_index` to
/// [`packetize_each`](Self::packetize_each) (the punktfunk/1 encode loop owns the video numbering
/// so the encoder's reference-frame-invalidation bookkeeping stays 1:1 with the wire across
/// encoder rebuilds/resets), or pass `None` to draw from the internal counter (the legacy path —
/// synthetic/spike/ABI sessions where no encoder cares). Speed-test probe filler draws from a
/// **separate** index space ([`alloc_probe_index`](Self::alloc_probe_index)) so a burst never
/// consumes video indexes — see [`crate::quic::VIDEO_CAP_PROBE_SEQ`].
pub struct Packetizer {
    next_frame_index: u32,
    /// Probe-space frame counter (see [`alloc_probe_index`](Self::alloc_probe_index)).
    next_probe_index: u32,
    next_seq: u32,
    shard_payload: usize,
    fec: crate::config::FecConfig,
    version: u8,
    /// Reusable zero-padded scratch for the frame's final data shard when the frame isn't an
    /// exact `shard_payload` multiple (and for the single all-zero shard of an empty frame).
    /// Every other data shard is a `shard_payload`-sized slice straight into the frame buffer —
    /// blocks are consecutive shard ranges, so only the frame's last shard can be partial.
    tail: Vec<u8>,
}

impl Packetizer {
    pub fn new(config: &Config) -> Self {
        Packetizer {
            next_frame_index: 0,
            next_probe_index: 0,
            next_seq: 0,
            shard_payload: config.shard_payload,
            fec: config.fec,
            version: config.phase as u8,
            tail: Vec::new(),
        }
    }

    /// Allocate the next **probe-space** frame index (speed-test filler). A separate counter from
    /// the video `frame_index`es so a multi-thousand-AU probe burst never advances the video
    /// numbering — the client routes [`FLAG_PROBE`]-flagged shards into its own reassembly window
    /// (see [`Reassembler`]), so the two spaces never collide. Only used against clients that
    /// advertise [`crate::quic::VIDEO_CAP_PROBE_SEQ`].
    pub fn alloc_probe_index(&mut self) -> u32 {
        let i = self.next_probe_index;
        self.next_probe_index = i.wrapping_add(1);
        i
    }

    /// Live-adjust the FEC recovery percentage (adaptive FEC). Takes effect on the next
    /// [`packetize`](Self::packetize); the wire is self-describing (each packet carries its block's
    /// data/recovery counts), so the receiver needs no notification. Clamped to ≤ 90.
    pub fn set_fec_percent(&mut self, pct: u8) {
        self.fec.fec_percent = pct.min(90);
    }

    /// The current FEC recovery percentage.
    pub fn fec_percent(&self) -> u8 {
        self.fec.fec_percent
    }

    /// Packetize one access unit into owned wire packets (header ++ shard payload each).
    /// Thin wrapper over [`packetize_each`](Self::packetize_each) — the allocation-free
    /// streaming path's reference implementation (tests and the loss harness use this).
    pub fn packetize(
        &mut self,
        frame: &[u8],
        pts_ns: u64,
        user_flags: u32,
        coder: &dyn ErasureCoder,
    ) -> Result<Vec<Vec<u8>>> {
        let mut packets = Vec::new();
        self.packetize_each(frame, pts_ns, user_flags, None, coder, |hdr, body| {
            let mut pkt = Vec::with_capacity(HEADER_LEN + body.len());
            pkt.extend_from_slice(hdr.as_bytes());
            pkt.extend_from_slice(body);
            packets.push(pkt);
            Ok(())
        })?;
        Ok(packets)
    }

    /// Packetize one access unit, yielding each packet to `emit` as a `(header, shard bytes)`
    /// pair — in exact wire order, which is also the order the session's nonce counter
    /// advances. No per-packet allocation happens here, so the caller can write header and
    /// shard straight into a pooled wire buffer and seal in place
    /// ([`Session::seal_frame`](crate::session::Session::seal_frame)). An `emit` error aborts
    /// the frame mid-way (packet numbering has already advanced — callers treat it as fatal).
    ///
    /// `frame_index`: `Some(i)` stamps the AU with the caller's index — the punktfunk/1 encode
    /// loop numbers video AUs itself so the encoder's RFI bookkeeping (LTR marks, DPB timestamps)
    /// is 1:1 with what the client sees, surviving encoder rebuilds/resets that restart internal
    /// counters. `None` draws from the internal counter (the legacy/self-numbering path). A
    /// session must not mix the two styles for the same index space.
    pub fn packetize_each(
        &mut self,
        frame: &[u8],
        pts_ns: u64,
        user_flags: u32,
        frame_index: Option<u32>,
        coder: &dyn ErasureCoder,
        mut emit: impl FnMut(&PacketHeader, &[u8]) -> Result<()>,
    ) -> Result<()> {
        let payload = self.shard_payload;
        let frame_index = frame_index.unwrap_or_else(|| {
            let i = self.next_frame_index;
            self.next_frame_index = i.wrapping_add(1);
            i
        });

        // At least one (zero-padded) data shard even for an empty frame.
        let total_data = frame.len().div_ceil(payload).max(1);
        let max_block = self.fec.max_data_per_block as usize;
        let block_count = total_data.div_ceil(max_block).max(1);
        let frame_bytes = frame.len() as u32;

        // Defend the u16 wire fields against silent truncation. `Config::validate`
        // already rejects configs that could reach these for valid frame sizes; this is
        // the belt-and-suspenders for a frame larger than the negotiated maximum.
        if payload > u16::MAX as usize {
            return Err(PunktfunkError::InvalidArg("shard_payload exceeds u16"));
        }
        if block_count > u16::MAX as usize {
            return Err(PunktfunkError::Unsupported(
                "frame too large: block count exceeds u16",
            ));
        }

        // Stage the frame's one possibly-partial shard (the last) in the reusable
        // zero-padded scratch; every full shard is referenced in place below.
        let full_shards = frame.len() / payload;
        self.tail.clear();
        self.tail.resize(payload, 0);
        let rem = frame.len() % payload;
        if rem > 0 {
            self.tail[..rem].copy_from_slice(&frame[full_shards * payload..]);
        }
        let tail = &self.tail;
        let shard_at = |s: usize| -> &[u8] {
            if s < full_shards {
                &frame[s * payload..(s + 1) * payload]
            } else {
                tail.as_slice()
            }
        };

        for b in 0..block_count {
            let first = b * max_block;
            let last = ((b + 1) * max_block).min(total_data);
            let block_data_count = last - first;

            // This block's data shards: references into `frame` (plus the staged tail).
            let data_shards: Vec<&[u8]> = (first..last).map(shard_at).collect();

            let recovery_count = self.fec.recovery_for(block_data_count);
            let recovery = coder.encode(&data_shards, recovery_count)?;
            let total_shards = block_data_count + recovery_count;
            if total_shards > u16::MAX as usize {
                return Err(PunktfunkError::Unsupported("block shard count exceeds u16"));
            }

            for shard_index in 0..total_shards {
                let body: &[u8] = if shard_index < block_data_count {
                    data_shards[shard_index]
                } else {
                    &recovery[shard_index - block_data_count]
                };

                let seq = self.next_seq;
                self.next_seq = self.next_seq.wrapping_add(1);

                let mut flags = FLAG_PIC;
                if b == 0 && shard_index == 0 {
                    flags |= FLAG_SOF;
                }
                if b + 1 == block_count && shard_index + 1 == total_shards {
                    flags |= FLAG_EOF;
                }

                let hdr = PacketHeader {
                    pts_ns,
                    frame_index,
                    stream_seq: seq,
                    frame_bytes,
                    user_flags,
                    block_index: b as u16,
                    block_count: block_count as u16,
                    data_shards: block_data_count as u16,
                    recovery_shards: recovery_count as u16,
                    shard_index: shard_index as u16,
                    shard_bytes: payload as u16,
                    magic: PUNKTFUNK_MAGIC,
                    version: self.version,
                    fec_scheme: coder.scheme() as u8,
                    flags,
                };
                emit(&hdr, body)?;
            }
        }
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Client side: reassembly + FEC recovery
// ---------------------------------------------------------------------------

struct BlockBuf {
    data_shards: usize,
    recovery_shards: usize,
    shard_bytes: usize,
    /// Length `data_shards + recovery_shards`; `Some` = received.
    shards: Vec<Option<Vec<u8>>>,
    received: usize,
    done: bool,
}

struct FrameBuf {
    frame_bytes: usize,
    block_count: usize,
    pts_ns: u64,
    user_flags: u32,
    blocks: HashMap<u16, BlockBuf>,
    /// Reconstructed payload per completed block, ordered by block index.
    block_data: BTreeMap<u16, Vec<u8>>,
}

/// Per-session bounds the reassembler enforces on every packet header *before*
/// allocating, so a hostile or corrupt header cannot drive unbounded memory use. All
/// derived from the negotiated [`Config`].
#[derive(Clone, Copy, Debug)]
pub struct ReassemblerLimits {
    /// Expected shard payload length; every shard in the stream must match exactly.
    pub shard_bytes: usize,
    /// Max data shards per block (the negotiated `max_data_per_block`).
    pub max_data_shards: usize,
    /// Max total shards per block (data + recovery), capped by the FEC scheme ceiling.
    pub max_total_shards: usize,
    /// Max FEC blocks per frame.
    pub max_blocks: usize,
    /// Max accepted access-unit size.
    pub max_frame_bytes: usize,
}

impl ReassemblerLimits {
    pub fn from_config(c: &Config) -> Self {
        let max_data = c.fec.max_data_per_block as usize;
        let max_total =
            (max_data + c.fec.recovery_for(max_data)).min(c.fec.scheme.max_total_shards());
        let total_data = c.max_frame_bytes.div_ceil(c.shard_payload.max(1)).max(1);
        ReassemblerLimits {
            shard_bytes: c.shard_payload,
            max_data_shards: max_data,
            max_total_shards: max_total,
            max_blocks: total_data.div_ceil(max_data).max(1),
            max_frame_bytes: c.max_frame_bytes,
        }
    }
}

/// One frame-index space's reassembly state: the in-flight frames, the recently-emitted memory,
/// and the loss-window anchor. The [`Reassembler`] keeps two — video and speed-test probe filler —
/// because the two ride **separate index counters** on a [`VIDEO_CAP_PROBE_SEQ`]-aware host
/// (a probe burst must neither advance the video loss window nor be dropped as "stale" against
/// it). [`VIDEO_CAP_PROBE_SEQ`]: crate::quic::VIDEO_CAP_PROBE_SEQ
#[derive(Default)]
struct ReassemblyWindow {
    frames: HashMap<u32, FrameBuf>,
    /// Recently-emitted frames, so stray/late shards can't resurrect them. Pruned to
    /// the reorder window alongside `frames`.
    completed: HashSet<u32>,
    /// The newest frame seen, as `(frame_index, capture pts)` — the loss-window anchor: an
    /// incomplete frame is declared lost once it sits [`LOSS_WINDOW_NS`] behind this pts (or
    /// [`HARD_LOSS_WINDOW`] indices, whichever trips first).
    newest_frame: Option<(u32, u64)>,
}

/// Buffers incoming shards, recovers lost ones via FEC, and emits whole access units.
/// Client-side only.
pub struct Reassembler {
    limits: ReassemblerLimits,
    /// The video stream's window — its aged-out incomplete frames count into `frames_dropped`
    /// (the client's loss-recovery trigger).
    video: ReassemblyWindow,
    /// Speed-test probe filler ([`FLAG_PROBE`] in `user_flags`). Routed by the flag, so it also
    /// captures an OLD host's probe frames (which still carry video-space indexes — they complete
    /// fine here, and keeping them out of the video window means a burst can no longer advance the
    /// video loss anchor). Aged-out probe frames are NOT `frames_dropped` — probe loss is measured
    /// bytes-wise by the probe accumulator and must not fire video recovery.
    probe: ReassemblyWindow,
}

impl Reassembler {
    pub fn new(limits: ReassemblerLimits) -> Self {
        Reassembler {
            limits,
            video: ReassemblyWindow::default(),
            probe: ReassemblyWindow::default(),
        }
    }

    /// Ingest one (already-decrypted) packet. Returns the access unit when its last
    /// block completes, otherwise `None`.
    pub fn push(
        &mut self,
        pkt: &[u8],
        coder: &dyn ErasureCoder,
        stats: &StatsCounters,
    ) -> Result<Option<Frame>> {
        // On a lossy datagram link a malformed or non-video packet is dropped, never
        // fatal: it must not abort `poll_frame`. A FEC reconstruction failure (corrupt or
        // incompatible shards that passed the header checks) likewise drops the block rather
        // than killing the whole session — the stream recovers at the next keyframe/RFI.
        if pkt.len() < HEADER_LEN {
            StatsCounters::add(&stats.packets_dropped, 1);
            return Ok(None);
        }
        let hdr = match PacketHeader::read_from_bytes(&pkt[..HEADER_LEN]) {
            Ok(h) => h,
            Err(_) => {
                StatsCounters::add(&stats.packets_dropped, 1);
                return Ok(None);
            }
        };

        let lim = self.limits;
        let shard_bytes = hdr.shard_bytes as usize;
        let data_shards = hdr.data_shards as usize;
        let recovery_shards = hdr.recovery_shards as usize;
        let total = data_shards + recovery_shards;
        let shard_index = hdr.shard_index as usize;
        let block_count = hdr.block_count as usize;
        let frame_bytes = hdr.frame_bytes as usize;

        // Bound every attacker-controllable header field against the negotiated limits
        // BEFORE allocating anything keyed on it — this is the firewall against a tiny
        // datagram triggering a huge `vec![None; total]` / `Vec::with_capacity`.
        let drop = |stats: &StatsCounters| {
            StatsCounters::add(&stats.packets_dropped, 1);
        };
        if hdr.magic != PUNKTFUNK_MAGIC
            || shard_bytes != lim.shard_bytes
            || pkt.len() < HEADER_LEN + shard_bytes
            || data_shards == 0
            || data_shards > lim.max_data_shards
            || total == 0
            || total > lim.max_total_shards
            || shard_index >= total
            || block_count == 0
            || block_count > lim.max_blocks
            || hdr.block_index as usize >= block_count
            || frame_bytes > lim.max_frame_bytes
        {
            drop(stats);
            return Ok(None);
        }
        let payload = pkt[HEADER_LEN..HEADER_LEN + shard_bytes].to_vec();

        // Route by index space: speed-test probe filler (FLAG_PROBE in user_flags) reassembles in
        // its own window so its indexes never interact with the video loss window — a probe burst
        // can neither advance the video anchor nor be dropped as stale against it (and its aged-out
        // frames never count as `frames_dropped`, which would fire video loss recovery).
        let is_probe = hdr.user_flags & (FLAG_PROBE as u32) != 0;
        let win = if is_probe {
            &mut self.probe
        } else {
            &mut self.video
        };
        win.advance_window(hdr.frame_index, hdr.pts_ns, stats, !is_probe);

        // Drop shards for frames we've already emitted (e.g. the recovery shards of a
        // frame that completed early via the all-originals-present fast path) or that
        // have fallen out of the loss window.
        if win.completed.contains(&hdr.frame_index) || win.is_stale(hdr.frame_index, hdr.pts_ns) {
            drop(stats);
            return Ok(None);
        }

        // First packet of a frame establishes its geometry; later packets must agree.
        let frame = win
            .frames
            .entry(hdr.frame_index)
            .or_insert_with(|| FrameBuf {
                frame_bytes,
                block_count,
                pts_ns: hdr.pts_ns,
                user_flags: hdr.user_flags,
                blocks: HashMap::new(),
                block_data: BTreeMap::new(),
            });
        if frame.block_count != block_count || frame.frame_bytes != frame_bytes {
            drop(stats);
            return Ok(None);
        }

        if frame.block_data.contains_key(&hdr.block_index) {
            return Ok(None); // block already reconstructed; late/duplicate shard
        }

        // First packet of a block sizes its shard vector; later packets must match its
        // (data, recovery, shard_bytes) geometry, so `shard_index` is always in bounds.
        frame
            .blocks
            .entry(hdr.block_index)
            .or_insert_with(|| BlockBuf {
                data_shards,
                recovery_shards,
                shard_bytes,
                shards: vec![None; total],
                received: 0,
                done: false,
            });
        let block = frame.blocks.get_mut(&hdr.block_index).unwrap();
        if block.data_shards != data_shards
            || block.recovery_shards != recovery_shards
            || block.shard_bytes != shard_bytes
        {
            drop(stats);
            return Ok(None);
        }

        if block.shards[shard_index].is_none() {
            block.shards[shard_index] = Some(payload);
            block.received += 1;
        }

        // Reconstruct as soon as we hold enough shards.
        if !block.done && block.received >= block.data_shards {
            let present_data = block.shards[..block.data_shards]
                .iter()
                .filter(|s| s.is_some())
                .count();
            let recovered = match coder.reconstruct(
                block.data_shards,
                block.recovery_shards,
                &mut block.shards,
            ) {
                Ok(r) => r,
                Err(_) => {
                    // Corrupt/incompatible shards that slipped past the header checks: discard this
                    // block (mark done so later shards for it are ignored) and keep the session
                    // alive — a lossy link must not be torn down by one unrecoverable block; the
                    // frame stays incomplete and the client recovers at the next keyframe/RFI.
                    block.done = true;
                    StatsCounters::add(&stats.packets_dropped, 1);
                    return Ok(None);
                }
            };
            block.done = true;
            StatsCounters::add(
                &stats.fec_recovered_shards,
                (block.data_shards - present_data) as u64,
            );

            // Concatenate the block's data shards into its contiguous payload.
            let mut block_payload = Vec::with_capacity(block.data_shards * block.shard_bytes);
            for shard in &recovered {
                block_payload.extend_from_slice(shard);
            }
            frame.block_data.insert(hdr.block_index, block_payload);
            frame.blocks.remove(&hdr.block_index);
        }

        // Whole frame ready?
        if frame.block_data.len() == frame.block_count {
            let frame = win.frames.remove(&hdr.frame_index).unwrap();
            win.completed.insert(hdr.frame_index);
            // Reserve based on the bytes we actually hold, not the (already-bounded but
            // still caller-supplied) frame_bytes, so a small frame can't over-reserve.
            let actual: usize = frame.block_data.values().map(|b| b.len()).sum();
            let mut data = Vec::with_capacity(actual);
            for (_, block_payload) in frame.block_data.into_iter() {
                data.extend_from_slice(&block_payload);
            }
            data.truncate(frame.frame_bytes); // trim trailing-shard zero padding
            return Ok(Some(Frame {
                data,
                frame_index: hdr.frame_index,
                pts_ns: frame.pts_ns,
                flags: frame.user_flags,
            }));
        }
        Ok(None)
    }

    /// Drop all in-flight state — every partially-assembled frame and the completed/abandoned
    /// index memory, in both index spaces — as if the session just started. Used by the client's
    /// backlog flush ([`Session::flush_backlog`](crate::session::Session::flush_backlog)): after
    /// the socket backlog is discarded wholesale, the partial frames here can never complete
    /// (their remaining shards were just thrown away) and the window anchors (`newest_frame`)
    /// point into the discarded past.
    pub fn reset(&mut self) {
        self.video = ReassemblyWindow::default();
        self.probe = ReassemblyWindow::default();
    }
}

impl ReassemblyWindow {
    /// Track the newest frame, declare incomplete frames that fell out of the loss window
    /// ([`LOSS_WINDOW_NS`] behind the newest pts, or [`HARD_LOSS_WINDOW`] indices) lost — for the
    /// video window (`count_drops`) counting them dropped, which is what drives the client's
    /// recovery-keyframe request — and prune the completed-index memory to [`REORDER_WINDOW`].
    fn advance_window(
        &mut self,
        frame_index: u32,
        pts_ns: u64,
        stats: &StatsCounters,
        count_drops: bool,
    ) {
        let (newest, newest_pts) = match self.newest_frame {
            // `frame_index` is newer iff it's within the forward half of the index space.
            Some((n, p)) if frame_index.wrapping_sub(n) > u32::MAX / 2 => (n, p),
            _ => (frame_index, pts_ns),
        };
        self.newest_frame = Some((newest, newest_pts));

        let before = self.frames.len();
        let completed = &mut self.completed;
        self.frames.retain(|&idx, f| {
            let keep = newest.wrapping_sub(idx) <= HARD_LOSS_WINDOW
                && newest_pts.saturating_sub(f.pts_ns) <= LOSS_WINDOW_NS;
            if !keep {
                // Remember the abandoned index so a straggler shard is dropped (below, and in
                // `push`) instead of resurrecting the frame — which would re-allocate its buffers
                // and double-count the drop when it aged out again.
                completed.insert(idx);
            }
            keep
        });
        let pruned = before - self.frames.len();
        if pruned > 0 && count_drops {
            StatsCounters::add(&stats.frames_dropped, pruned as u64);
        }
        self.completed
            .retain(|&idx| newest.wrapping_sub(idx) <= REORDER_WINDOW);
    }

    /// True if this packet's frame lies outside the loss window (behind the newest frame by more
    /// than [`LOSS_WINDOW_NS`] of capture time or [`HARD_LOSS_WINDOW`] indices) — its shards
    /// arrive too late to be useful, and accepting one would only create a frame buffer the next
    /// [`advance_window`](Self::advance_window) immediately declares lost.
    fn is_stale(&self, frame_index: u32, pts_ns: u64) -> bool {
        match self.newest_frame {
            Some((n, newest_pts)) => {
                let behind = n.wrapping_sub(frame_index);
                behind <= u32::MAX / 2
                    && (behind > HARD_LOSS_WINDOW
                        || newest_pts.saturating_sub(pts_ns) > LOSS_WINDOW_NS)
            }
            None => false,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::FecScheme;
    use crate::fec::coder_for;

    fn limits() -> ReassemblerLimits {
        ReassemblerLimits {
            shard_bytes: 16,
            max_data_shards: 8,
            max_total_shards: 12,
            max_blocks: 4,
            max_frame_bytes: 4096,
        }
    }

    fn base_header() -> PacketHeader {
        PacketHeader {
            pts_ns: 0,
            frame_index: 0,
            stream_seq: 0,
            frame_bytes: 16,
            user_flags: 0,
            block_index: 0,
            block_count: 1,
            data_shards: 1,
            recovery_shards: 0,
            shard_index: 0,
            shard_bytes: 16,
            magic: PUNKTFUNK_MAGIC,
            version: 1,
            fec_scheme: 0,
            flags: FLAG_PIC,
        }
    }

    fn packet(h: PacketHeader) -> Vec<u8> {
        let mut p = Vec::new();
        p.extend_from_slice(h.as_bytes());
        p.extend_from_slice(&vec![0xAB; h.shard_bytes as usize]);
        p
    }

    /// A header advertising 65535+65535 shards must be dropped, not allocate gigabytes.
    #[test]
    fn rejects_oversized_shard_counts() {
        let mut r = Reassembler::new(limits());
        let coder = coder_for(FecScheme::Gf8);
        let stats = StatsCounters::default();
        let mut h = base_header();
        h.data_shards = 65535;
        h.recovery_shards = 65535;
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_none());
        assert_eq!(stats.snapshot().packets_dropped, 1);
    }

    /// A second packet for a block whose geometry differs from the first must be dropped
    /// — never index past the block's allocated shard vector (the old OOB panic).
    #[test]
    fn rejects_inconsistent_block_geometry_without_panicking() {
        let mut r = Reassembler::new(limits());
        let coder = coder_for(FecScheme::Gf8);
        let stats = StatsCounters::default();

        let mut h1 = base_header();
        h1.data_shards = 4;
        h1.recovery_shards = 2; // block sized to 6 slots
        h1.frame_bytes = 64;
        assert!(r
            .push(&packet(h1), coder.as_ref(), &stats)
            .unwrap()
            .is_none());

        // Same block, different geometry, shard_index valid for ITS total (8) but past
        // the established block's 6 slots.
        let mut h2 = base_header();
        h2.data_shards = 6;
        h2.recovery_shards = 2;
        h2.shard_index = 7;
        h2.frame_bytes = 64;
        assert!(r
            .push(&packet(h2), coder.as_ref(), &stats)
            .unwrap()
            .is_none());
        assert_eq!(stats.snapshot().packets_dropped, 1);
    }

    /// The loss window is TIME-based: an incomplete frame survives newer frames arriving within
    /// [`LOSS_WINDOW_NS`] of its capture pts (a 33 ms-late shard at 120 fps is late, not lost —
    /// the old 4-INDEX window wrongly killed it), is declared lost once the newest pts moves past
    /// the window (`frames_dropped`), and a straggler shard can't resurrect it afterwards.
    #[test]
    fn incomplete_frames_age_out_by_capture_time_not_frame_count() {
        let mut r = Reassembler::new(limits());
        let coder = coder_for(FecScheme::Gf8);
        let stats = StatsCounters::default();
        const FRAME_NS: u64 = 8_333_333; // 120 fps

        // Frame 0: one of its two shards arrives — incomplete.
        let mut h = base_header();
        h.data_shards = 2;
        h.frame_bytes = 32;
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_none());

        // Frames 1..=8 complete around it (well past the old 4-index window, inside 120 ms):
        // frame 0 must still be alive — no drop counted.
        for i in 1..=8u32 {
            let mut h = base_header();
            h.frame_index = i;
            h.pts_ns = i as u64 * FRAME_NS;
            assert!(r
                .push(&packet(h), coder.as_ref(), &stats)
                .unwrap()
                .is_some());
        }
        assert_eq!(stats.snapshot().frames_dropped, 0);

        // Frame 0's second shard arrives 8 frames late (~66 ms at 120 fps) — completes fine.
        let mut h = base_header();
        h.data_shards = 2;
        h.frame_bytes = 32;
        h.shard_index = 1;
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_some());

        // Frame 20: incomplete again; then a frame lands past the 120 ms window → declared lost.
        let mut h = base_header();
        h.frame_index = 20;
        h.pts_ns = 20 * FRAME_NS;
        h.data_shards = 2;
        h.frame_bytes = 32;
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_none());
        let mut h = base_header();
        h.frame_index = 21;
        h.pts_ns = 20 * FRAME_NS + LOSS_WINDOW_NS + 1;
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_some());
        assert_eq!(stats.snapshot().frames_dropped, 1);

        // A straggler shard for the abandoned frame 20 is dropped, never resurrected.
        let mut h = base_header();
        h.frame_index = 20;
        h.pts_ns = 20 * FRAME_NS;
        h.data_shards = 2;
        h.frame_bytes = 32;
        h.shard_index = 1;
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_none());
        assert_eq!(stats.snapshot().frames_dropped, 1, "no double-count");
    }

    /// The explicit-index path stamps the caller's `frame_index` and leaves the internal video
    /// counter untouched — the punktfunk/1 encode loop owns the numbering, and mixing must not
    /// perturb the legacy self-numbering path (tests/ABI/synthetic).
    #[test]
    fn explicit_frame_index_is_stamped_and_internal_counter_untouched() {
        use crate::config::{FecConfig, FecScheme, ProtocolPhase, Role};
        let cfg = Config {
            role: Role::Host,
            phase: ProtocolPhase::P2Punktfunk,
            fec: FecConfig {
                scheme: FecScheme::Gf16,
                fec_percent: 0,
                max_data_per_block: 8,
            },
            shard_payload: 16,
            max_frame_bytes: 4096,
            encrypt: false,
            key: [0u8; 16],
            salt: [0u8; 4],
            loopback_drop_period: 0,
        };
        let coder = coder_for(FecScheme::Gf16);
        let mut pk = Packetizer::new(&cfg);
        let mut seen = Vec::new();
        pk.packetize_each(&[1u8; 16], 0, 0, Some(4242), coder.as_ref(), |hdr, _| {
            seen.push(hdr.frame_index);
            Ok(())
        })
        .unwrap();
        assert_eq!(seen, vec![4242]);
        // The legacy wrapper still numbers from the untouched internal counter.
        let pkts = pk.packetize(&[1u8; 16], 0, 0, coder.as_ref()).unwrap();
        let hdr = PacketHeader::read_from_bytes(&pkts[0][..HEADER_LEN]).unwrap();
        assert_eq!(hdr.frame_index, 0);
        // The probe space is a third, independent counter.
        assert_eq!(pk.alloc_probe_index(), 0);
        assert_eq!(pk.alloc_probe_index(), 1);
    }

    /// Probe filler (FLAG_PROBE in user_flags) reassembles in its OWN window: a probe frame whose
    /// index is far behind the video stream's completes anyway (an old client's single window
    /// would drop it as stale), and video frames complete undisturbed around it.
    #[test]
    fn probe_frames_reassemble_in_their_own_window() {
        let mut r = Reassembler::new(limits());
        let coder = coder_for(FecScheme::Gf8);
        let stats = StatsCounters::default();

        // Establish a video stream far into its index space.
        let mut v = base_header();
        v.frame_index = 100_000;
        v.pts_ns = 1_000_000_000;
        assert!(r
            .push(&packet(v), coder.as_ref(), &stats)
            .unwrap()
            .is_some());

        // A probe frame at index 0 — 100k "behind" the video window — must still complete.
        let mut p = base_header();
        p.frame_index = 0;
        p.pts_ns = 1_000_000_100;
        p.user_flags = FLAG_PROBE as u32;
        let got = r.push(&packet(p), coder.as_ref(), &stats).unwrap();
        assert!(got.is_some(), "probe frame must complete in its own window");
        assert_eq!(got.unwrap().flags & FLAG_PROBE as u32, FLAG_PROBE as u32);

        // The probe burst must not have advanced the VIDEO window: the next video frame is
        // contiguous and completes, with nothing counted dropped.
        let mut v2 = base_header();
        v2.frame_index = 100_001;
        v2.pts_ns = 1_000_000_200;
        assert!(r
            .push(&packet(v2), coder.as_ref(), &stats)
            .unwrap()
            .is_some());
        assert_eq!(stats.snapshot().frames_dropped, 0);
    }

    /// An incomplete probe frame aging out of the probe window is NOT a video `frames_dropped`
    /// (which would fire the client's loss recovery) — probe loss is measured bytes-wise by the
    /// probe accumulator.
    #[test]
    fn aged_out_probe_frames_do_not_count_as_dropped() {
        let mut r = Reassembler::new(limits());
        let coder = coder_for(FecScheme::Gf8);
        let stats = StatsCounters::default();

        // Probe frame 0: one of two shards — incomplete.
        let mut p = base_header();
        p.user_flags = FLAG_PROBE as u32;
        p.data_shards = 2;
        p.frame_bytes = 32;
        assert!(r
            .push(&packet(p), coder.as_ref(), &stats)
            .unwrap()
            .is_none());

        // A much newer probe frame ages it out of the probe window.
        let mut p2 = base_header();
        p2.user_flags = FLAG_PROBE as u32;
        p2.frame_index = 1;
        p2.pts_ns = LOSS_WINDOW_NS + 1;
        assert!(r
            .push(&packet(p2), coder.as_ref(), &stats)
            .unwrap()
            .is_some());
        assert_eq!(
            stats.snapshot().frames_dropped,
            0,
            "probe-window drops must not fire video loss recovery"
        );
    }

    #[test]
    fn rejects_wrong_shard_bytes_and_oversized_frame() {
        let coder = coder_for(FecScheme::Gf8);

        let mut r = Reassembler::new(limits());
        let stats = StatsCounters::default();
        let mut h = base_header();
        h.shard_bytes = 8; // != negotiated 16
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_none());
        assert_eq!(stats.snapshot().packets_dropped, 1);

        let mut r = Reassembler::new(limits());
        let stats = StatsCounters::default();
        let mut h = base_header();
        h.frame_bytes = 1_000_000; // > max_frame_bytes
        assert!(r
            .push(&packet(h), coder.as_ref(), &stats)
            .unwrap()
            .is_none());
        assert_eq!(stats.snapshot().packets_dropped, 1);
    }
}