//! Embeddable reliable stream machinery for Dosh-native application transports.
//!
//! The terminal client/server binaries use the same wire packets exposed here:
//! `StreamOpen`, `StreamData`, `StreamWindowAdjust`, and `StreamClose`. This
//! module keeps the stream state machine independent from UDP sockets, native
//! auth, terminal rendering, or any particular service protocol. An embedding
//! application is expected to:
//!
//! 1. authenticate and maintain a Dosh session using [`crate::protocol`],
//! 2. encrypt each [`OutgoingStreamPacket::body`] as the given packet kind,
//! 3. pass received stream packet bodies back into [`StreamMux`],
//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data and
//!    idempotent stream control packets.
//!
//! That split is intentional: Dosh roaming/reconnect stays at the authenticated
//! session layer, while this module provides the reusable fast persistent byte
//! streams that higher-level app protocols can embed. [`DoshTransport::recv`]
//! calls [`DoshTransport::maintenance`] while it is idle, so embedders using the
//! session wrapper get retransmission and keepalives in a simple `recv().await`
//! loop.

use crate::protocol::{
    self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData,
    StreamEof, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
};
use crate::udp::{is_transient_udp_error, is_transient_udp_send_error};
use anyhow::{Result, anyhow, bail};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;

pub const DEFAULT_INITIAL_WINDOW: usize = 1024 * 1024;
pub const DEFAULT_RETRANSMIT_AFTER: Duration = Duration::from_millis(200);
pub const ADAPTIVE_RETRANSMIT_PAD: Duration = Duration::from_millis(10);
pub const ADAPTIVE_RETRANSMIT_MIN: Duration = Duration::from_millis(10);
pub const DEFAULT_KEEPALIVE_AFTER: Duration = Duration::from_secs(2);
pub const DEFAULT_RETIRED_STREAM_TOMBSTONES: usize = 16 * 1024;
pub const STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS: u32 = 8;
pub const SERVICE_TARGET_PREFIX: &str = "@dosh-";
pub const MAX_STREAM_DATA_BYTES: usize = 60 * 1024;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TransportConfig {
    pub initial_window: usize,
    pub retransmit_after: Duration,
    pub keepalive_after: Duration,
    pub retired_stream_tombstones: usize,
}

impl Default for TransportConfig {
    fn default() -> Self {
        Self {
            initial_window: DEFAULT_INITIAL_WINDOW,
            retransmit_after: DEFAULT_RETRANSMIT_AFTER,
            keepalive_after: DEFAULT_KEEPALIVE_AFTER,
            retired_stream_tombstones: DEFAULT_RETIRED_STREAM_TOMBSTONES,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionRole {
    Client,
    Server,
}

impl SessionRole {
    fn send_direction(self) -> u32 {
        match self {
            Self::Client => CLIENT_TO_SERVER,
            Self::Server => SERVER_TO_CLIENT,
        }
    }

    fn recv_direction(self) -> u32 {
        match self {
            Self::Client => SERVER_TO_CLIENT,
            Self::Server => CLIENT_TO_SERVER,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionTransportConfig {
    pub role: SessionRole,
    pub conn_id: [u8; 16],
    pub session_key: [u8; 32],
    pub peer_addr: SocketAddr,
    pub initial_send_seq: u64,
    pub initial_ack: u64,
    pub stream: TransportConfig,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SessionEvent {
    Stream(TransportEvent),
    Ping,
    Pong,
    Ignored,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutgoingStreamPacket {
    pub kind: PacketKind,
    pub body: Vec<u8>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IncomingStreamData {
    pub stream_id: u64,
    pub chunks: Vec<Vec<u8>>,
    pub received_offset: u64,
    pub consumed: usize,
    pub window_adjust: OutgoingStreamPacket,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IncomingStreamOpen {
    pub stream_id: u64,
    pub target_host: String,
    pub target_port: u16,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AcceptedService {
    pub name: String,
    pub stream_id: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TransportEvent {
    Open(IncomingStreamOpen),
    DuplicateOpen {
        stream_id: u64,
        response: OutgoingStreamPacket,
    },
    OpenOk {
        stream_id: u64,
        flushed: Vec<OutgoingStreamPacket>,
    },
    OpenReject {
        stream_id: u64,
        reason: String,
    },
    Data(IncomingStreamData),
    WindowAdjust {
        stream_id: u64,
        flushed: Vec<OutgoingStreamPacket>,
    },
    Eof {
        stream_id: u64,
    },
    Close {
        stream_id: u64,
    },
    Ignored {
        stream_id: u64,
    },
}

#[derive(Debug, Clone)]
struct PendingStreamChunk {
    offset: u64,
    bytes: Vec<u8>,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Debug, Clone)]
struct PendingStreamOpen {
    target_host: String,
    target_port: u16,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Debug, Clone)]
struct PendingStreamControl {
    last_sent: Instant,
    attempts: u32,
}

#[derive(Debug, Clone)]
struct PendingWindowAdjust {
    received_offset: u64,
    bytes: usize,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Debug, Clone)]
pub struct StreamMux {
    config: TransportConfig,
    opened_streams: HashSet<u64>,
    pending_opens: HashMap<u64, PendingStreamOpen>,
    send_credit: HashMap<u64, usize>,
    pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
    sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
    next_send_offset: HashMap<u64, u64>,
    next_recv_offset: HashMap<u64, u64>,
    recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
    pending_eofs: HashMap<u64, PendingStreamControl>,
    pending_closes: HashMap<u64, PendingStreamControl>,
    pending_window_adjusts: HashMap<u64, PendingWindowAdjust>,
    retired_streams: HashSet<u64>,
    retired_stream_order: VecDeque<u64>,
    srtt: Option<Duration>,
}

#[derive(Debug, Clone, Default)]
pub struct ServiceRegistry {
    services: HashSet<String>,
}

impl ServiceRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register(&mut self, name: impl Into<String>) -> Result<()> {
        let name = name.into();
        service_target(&name)?;
        self.services.insert(name);
        Ok(())
    }

    pub fn is_registered(&self, name: &str) -> bool {
        self.services.contains(name)
    }

    pub fn match_open(&self, open: &IncomingStreamOpen) -> Result<AcceptedService> {
        let Some(name) = service_name_from_target(&open.target_host) else {
            bail!("stream target {} is not a Dosh service", open.target_host);
        };
        if !self.services.contains(name) {
            bail!("service {name:?} is not registered");
        }
        if open.target_port != 0 {
            bail!(
                "service stream {} used non-zero target port",
                open.stream_id
            );
        }
        Ok(AcceptedService {
            name: name.to_string(),
            stream_id: open.stream_id,
        })
    }
}

impl StreamMux {
    pub fn new(config: TransportConfig) -> Self {
        Self {
            config,
            opened_streams: HashSet::new(),
            pending_opens: HashMap::new(),
            send_credit: HashMap::new(),
            pending_data: HashMap::new(),
            sent_data: HashMap::new(),
            next_send_offset: HashMap::new(),
            next_recv_offset: HashMap::new(),
            recv_pending: HashMap::new(),
            pending_eofs: HashMap::new(),
            pending_closes: HashMap::new(),
            pending_window_adjusts: HashMap::new(),
            retired_streams: HashSet::new(),
            retired_stream_order: VecDeque::new(),
            srtt: None,
        }
    }

    pub fn open_stream(
        &mut self,
        stream_id: u64,
        target_host: impl Into<String>,
        target_port: u16,
    ) -> Result<OutgoingStreamPacket> {
        let target_host = target_host.into();
        if self.opened_streams.contains(&stream_id)
            || self.pending_opens.contains_key(&stream_id)
            || self.retired_streams.contains(&stream_id)
        {
            bail!("stream {stream_id} already exists");
        }
        self.send_credit
            .insert(stream_id, self.config.initial_window);
        self.next_send_offset.entry(stream_id).or_insert(0);
        self.next_recv_offset.entry(stream_id).or_insert(0);
        self.pending_opens.insert(
            stream_id,
            PendingStreamOpen {
                target_host: target_host.clone(),
                target_port,
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        encode_packet(
            PacketKind::StreamOpen,
            &StreamOpen {
                stream_id,
                target_host,
                target_port,
            },
        )
    }

    pub fn accept_open(&mut self, open: StreamOpen) -> Result<IncomingStreamOpen> {
        if self.retired_streams.contains(&open.stream_id) {
            bail!("stream {} was already closed", open.stream_id);
        }
        if self.pending_opens.contains_key(&open.stream_id) {
            bail!(
                "stream {} collides with a local pending open",
                open.stream_id
            );
        }
        self.opened_streams.insert(open.stream_id);
        self.send_credit
            .entry(open.stream_id)
            .or_insert(self.config.initial_window);
        self.next_send_offset.entry(open.stream_id).or_insert(0);
        self.next_recv_offset.entry(open.stream_id).or_insert(0);
        Ok(IncomingStreamOpen {
            stream_id: open.stream_id,
            target_host: open.target_host,
            target_port: open.target_port,
        })
    }

    pub fn open_ok(stream_id: u64) -> Result<OutgoingStreamPacket> {
        encode_packet(PacketKind::StreamOpenOk, &StreamOpenOk { stream_id })
    }

    pub fn open_reject(stream_id: u64, reason: impl Into<String>) -> Result<OutgoingStreamPacket> {
        encode_packet(
            PacketKind::StreamOpenReject,
            &StreamOpenReject {
                stream_id,
                reason: reason.into(),
            },
        )
    }

    pub fn reject_open(
        &mut self,
        stream_id: u64,
        reason: impl Into<String>,
    ) -> Result<OutgoingStreamPacket> {
        if self.has_stream_state(stream_id) {
            self.retire_stream(stream_id);
        }
        Self::open_reject(stream_id, reason)
    }

    pub fn handle_open_ok(
        &mut self,
        ok: StreamOpenOk,
    ) -> Result<Option<Vec<OutgoingStreamPacket>>> {
        let Some(pending) = self.pending_opens.remove(&ok.stream_id) else {
            return Ok(None);
        };
        if pending.attempts == 1 {
            self.observe_rtt_sample(pending.last_sent.elapsed());
        }
        self.opened_streams.insert(ok.stream_id);
        self.send_credit
            .entry(ok.stream_id)
            .or_insert(self.config.initial_window);
        self.next_send_offset.entry(ok.stream_id).or_insert(0);
        self.next_recv_offset.entry(ok.stream_id).or_insert(0);
        self.flush_pending_data(ok.stream_id).map(Some)
    }

    pub fn handle_open_reject(&mut self, reject: StreamOpenReject) -> bool {
        if !self.has_stream_state(reject.stream_id) {
            return false;
        }
        self.retire_stream(reject.stream_id);
        true
    }

    pub fn send_data(
        &mut self,
        stream_id: u64,
        bytes: impl Into<Vec<u8>>,
    ) -> Result<Vec<OutgoingStreamPacket>> {
        let bytes = bytes.into();
        if bytes.is_empty() {
            return Ok(Vec::new());
        }
        let chunks = split_stream_data_bytes(bytes, self.config.initial_window);
        let mut out = Vec::new();
        if !self.opened_streams.contains(&stream_id) && !self.pending_opens.contains_key(&stream_id)
        {
            bail!("stream {stream_id} is not open");
        }
        for chunk in chunks {
            if !self.opened_streams.contains(&stream_id) {
                self.pending_data
                    .entry(stream_id)
                    .or_default()
                    .push_back(chunk);
                continue;
            }
            if self.send_credit.get(&stream_id).copied().unwrap_or(0) < chunk.len()
                || self
                    .pending_data
                    .get(&stream_id)
                    .is_some_and(|pending| !pending.is_empty())
            {
                self.pending_data
                    .entry(stream_id)
                    .or_default()
                    .push_back(chunk);
                continue;
            }
            out.push(self.send_data_now(stream_id, chunk)?);
        }
        Ok(out)
    }

    pub fn handle_data(&mut self, data: StreamData) -> Result<Option<IncomingStreamData>> {
        let stream_id = data.stream_id;
        if !self.opened_streams.contains(&stream_id) {
            return Ok(None);
        }
        let (chunks, consumed, received_offset) = self.accept_data(data);
        let window_adjust = encode_packet(
            PacketKind::StreamWindowAdjust,
            &StreamWindowAdjust {
                stream_id,
                received_offset,
                bytes: consumed.min(u32::MAX as usize) as u32,
            },
        )?;
        if consumed > 0 {
            self.pending_window_adjusts.insert(
                stream_id,
                PendingWindowAdjust {
                    received_offset,
                    bytes: consumed,
                    last_sent: Instant::now(),
                    attempts: 1,
                },
            );
        }
        Ok(Some(IncomingStreamData {
            stream_id,
            chunks,
            consumed,
            received_offset,
            window_adjust,
        }))
    }

    pub fn handle_window_adjust(
        &mut self,
        adjust: StreamWindowAdjust,
    ) -> Result<Option<Vec<OutgoingStreamPacket>>> {
        if !self.opened_streams.contains(&adjust.stream_id) {
            return Ok(None);
        }
        let acked_bytes = self.ack_data(adjust.stream_id, adjust.received_offset);
        self.add_credit(adjust.stream_id, acked_bytes);
        self.flush_pending_data(adjust.stream_id).map(Some)
    }

    pub fn close_stream(&mut self, stream_id: u64) -> Result<OutgoingStreamPacket> {
        if !self.has_stream_state(stream_id) {
            bail!("stream {stream_id} is not open");
        }
        self.retire_stream(stream_id);
        self.pending_closes.insert(
            stream_id,
            PendingStreamControl {
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        encode_packet(PacketKind::StreamClose, &StreamClose { stream_id })
    }

    pub fn eof_stream(&mut self, stream_id: u64) -> Result<OutgoingStreamPacket> {
        if !self.has_stream_state(stream_id) {
            bail!("stream {stream_id} is not open");
        }
        self.pending_eofs.insert(
            stream_id,
            PendingStreamControl {
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        encode_packet(PacketKind::StreamEof, &StreamEof { stream_id })
    }

    pub fn handle_eof(&self, eof: StreamEof) -> bool {
        self.has_stream_state(eof.stream_id)
    }

    pub fn handle_close(&mut self, close: StreamClose) -> bool {
        let had_state = self.has_stream_state(close.stream_id);
        self.pending_closes.remove(&close.stream_id);
        if !had_state {
            return false;
        }
        self.retire_stream(close.stream_id);
        true
    }

    pub fn handle_packet(&mut self, kind: PacketKind, body: &[u8]) -> Result<TransportEvent> {
        match kind {
            PacketKind::StreamOpen => {
                let open: StreamOpen = protocol::from_body(body)?;
                if self.retired_streams.contains(&open.stream_id) {
                    return Ok(TransportEvent::Ignored {
                        stream_id: open.stream_id,
                    });
                }
                if self.opened_streams.contains(&open.stream_id) {
                    return Ok(TransportEvent::DuplicateOpen {
                        stream_id: open.stream_id,
                        response: Self::open_ok(open.stream_id)?,
                    });
                }
                self.accept_open(open).map(TransportEvent::Open)
            }
            PacketKind::StreamOpenOk => {
                let ok: StreamOpenOk = protocol::from_body(body)?;
                let stream_id = ok.stream_id;
                match self.handle_open_ok(ok)? {
                    Some(flushed) => Ok(TransportEvent::OpenOk { stream_id, flushed }),
                    None => Ok(TransportEvent::Ignored { stream_id }),
                }
            }
            PacketKind::StreamOpenReject => {
                let reject: StreamOpenReject = protocol::from_body(body)?;
                let stream_id = reject.stream_id;
                if self.handle_open_reject(reject.clone()) {
                    Ok(TransportEvent::OpenReject {
                        stream_id,
                        reason: reject.reason,
                    })
                } else {
                    Ok(TransportEvent::Ignored { stream_id })
                }
            }
            PacketKind::StreamData => {
                let data: StreamData = protocol::from_body(body)?;
                let stream_id = data.stream_id;
                match self.handle_data(data)? {
                    Some(data) => Ok(TransportEvent::Data(data)),
                    None => Ok(TransportEvent::Ignored { stream_id }),
                }
            }
            PacketKind::StreamWindowAdjust => {
                let adjust: StreamWindowAdjust = protocol::from_body(body)?;
                let stream_id = adjust.stream_id;
                match self.handle_window_adjust(adjust)? {
                    Some(flushed) => Ok(TransportEvent::WindowAdjust { stream_id, flushed }),
                    None => Ok(TransportEvent::Ignored { stream_id }),
                }
            }
            PacketKind::StreamClose => {
                let close: StreamClose = protocol::from_body(body)?;
                let stream_id = close.stream_id;
                if self.handle_close(close) {
                    Ok(TransportEvent::Close { stream_id })
                } else {
                    Ok(TransportEvent::Ignored { stream_id })
                }
            }
            PacketKind::StreamEof => {
                let eof: StreamEof = protocol::from_body(body)?;
                let stream_id = eof.stream_id;
                if self.handle_eof(eof) {
                    Ok(TransportEvent::Eof { stream_id })
                } else {
                    Ok(TransportEvent::Ignored { stream_id })
                }
            }
            _ => bail!("packet kind {kind:?} is not a stream transport packet"),
        }
    }

    pub fn tick(&mut self) -> Result<Vec<OutgoingStreamPacket>> {
        let now = Instant::now();
        let retransmit_after = self.effective_retransmit_after();
        let mut out = Vec::new();
        let mut retransmit_opens = Vec::new();
        for (stream_id, pending) in self.pending_opens.iter_mut() {
            if now.duration_since(pending.last_sent) < retransmit_after {
                continue;
            }
            pending.last_sent = now;
            pending.attempts = pending.attempts.saturating_add(1);
            retransmit_opens.push((*stream_id, pending.target_host.clone(), pending.target_port));
        }
        for (stream_id, target_host, target_port) in retransmit_opens {
            out.push(encode_packet(
                PacketKind::StreamOpen,
                &StreamOpen {
                    stream_id,
                    target_host,
                    target_port,
                },
            )?);
        }

        let mut retransmit_data = Vec::new();
        for (stream_id, chunks) in self.sent_data.iter_mut() {
            for chunk in chunks.values_mut() {
                if now.duration_since(chunk.last_sent) < retransmit_after {
                    continue;
                }
                chunk.last_sent = now;
                chunk.attempts = chunk.attempts.saturating_add(1);
                retransmit_data.push((*stream_id, chunk.offset, chunk.bytes.clone()));
            }
        }
        for (stream_id, offset, bytes) in retransmit_data {
            out.push(encode_packet(
                PacketKind::StreamData,
                &StreamData {
                    stream_id,
                    offset,
                    bytes,
                },
            )?);
        }

        let mut retransmit_eofs = Vec::new();
        for (stream_id, pending) in self.pending_eofs.iter_mut() {
            if !self.opened_streams.contains(stream_id) {
                continue;
            }
            if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
                continue;
            }
            if now.duration_since(pending.last_sent) < retransmit_after {
                continue;
            }
            pending.last_sent = now;
            pending.attempts = pending.attempts.saturating_add(1);
            retransmit_eofs.push(*stream_id);
        }
        let opened_streams = self.opened_streams.clone();
        self.pending_eofs.retain(|stream_id, pending| {
            opened_streams.contains(stream_id)
                && pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS
        });
        for stream_id in retransmit_eofs {
            out.push(encode_packet(
                PacketKind::StreamEof,
                &StreamEof { stream_id },
            )?);
        }

        let mut retransmit_closes = Vec::new();
        for (stream_id, pending) in self.pending_closes.iter_mut() {
            if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
                continue;
            }
            if now.duration_since(pending.last_sent) < retransmit_after {
                continue;
            }
            pending.last_sent = now;
            pending.attempts = pending.attempts.saturating_add(1);
            retransmit_closes.push(*stream_id);
        }
        self.pending_closes
            .retain(|_, pending| pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS);
        for stream_id in retransmit_closes {
            out.push(encode_packet(
                PacketKind::StreamClose,
                &StreamClose { stream_id },
            )?);
        }

        let mut retransmit_adjusts = Vec::new();
        for (stream_id, pending) in self.pending_window_adjusts.iter_mut() {
            if pending.attempts >= STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS {
                continue;
            }
            if now.duration_since(pending.last_sent) < retransmit_after {
                continue;
            }
            pending.last_sent = now;
            pending.attempts = pending.attempts.saturating_add(1);
            retransmit_adjusts.push((*stream_id, pending.received_offset, pending.bytes));
        }
        self.pending_window_adjusts
            .retain(|_, pending| pending.attempts < STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS);
        for (stream_id, received_offset, bytes) in retransmit_adjusts {
            out.push(encode_packet(
                PacketKind::StreamWindowAdjust,
                &StreamWindowAdjust {
                    stream_id,
                    received_offset,
                    bytes: bytes.min(u32::MAX as usize) as u32,
                },
            )?);
        }
        Ok(out)
    }

    pub fn effective_retransmit_after(&self) -> Duration {
        adaptive_retransmit_after(self.config.retransmit_after, self.srtt)
    }

    pub fn is_open(&self, stream_id: u64) -> bool {
        self.opened_streams.contains(&stream_id)
    }

    pub fn pending_bytes(&self, stream_id: u64) -> usize {
        self.pending_data
            .get(&stream_id)
            .map(|queue| queue.iter().map(Vec::len).sum())
            .unwrap_or(0)
    }

    fn send_data_now(&mut self, stream_id: u64, bytes: Vec<u8>) -> Result<OutgoingStreamPacket> {
        let credit = self
            .send_credit
            .get_mut(&stream_id)
            .ok_or_else(|| anyhow!("stream {stream_id} has no send window"))?;
        if *credit < bytes.len() {
            bail!("stream {stream_id} send window exhausted");
        }
        *credit -= bytes.len();
        let offset = *self.next_send_offset.entry(stream_id).or_default();
        self.next_send_offset
            .insert(stream_id, offset.saturating_add(bytes.len() as u64));
        self.sent_data.entry(stream_id).or_default().insert(
            offset,
            PendingStreamChunk {
                offset,
                bytes: bytes.clone(),
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        encode_packet(
            PacketKind::StreamData,
            &StreamData {
                stream_id,
                offset,
                bytes,
            },
        )
    }

    fn flush_pending_data(&mut self, stream_id: u64) -> Result<Vec<OutgoingStreamPacket>> {
        let mut out = Vec::new();
        let chunk_limit = stream_data_chunk_limit(self.config.initial_window);
        while let Some(front_len) = self
            .pending_data
            .get(&stream_id)
            .and_then(|pending| pending.front().map(Vec::len))
        {
            if front_len > chunk_limit {
                let Some(bytes) = self
                    .pending_data
                    .get_mut(&stream_id)
                    .and_then(VecDeque::pop_front)
                else {
                    break;
                };
                if let Some(pending) = self.pending_data.get_mut(&stream_id) {
                    for chunk in split_stream_data_bytes(bytes, self.config.initial_window)
                        .into_iter()
                        .rev()
                    {
                        pending.push_front(chunk);
                    }
                }
                continue;
            }
            if self.send_credit.get(&stream_id).copied().unwrap_or(0) < front_len {
                break;
            }
            let Some(bytes) = self
                .pending_data
                .get_mut(&stream_id)
                .and_then(VecDeque::pop_front)
            else {
                break;
            };
            if self
                .pending_data
                .get(&stream_id)
                .is_some_and(VecDeque::is_empty)
            {
                self.pending_data.remove(&stream_id);
            }
            out.push(self.send_data_now(stream_id, bytes)?);
        }
        Ok(out)
    }

    fn accept_data(&mut self, data: StreamData) -> (Vec<Vec<u8>>, usize, u64) {
        accept_stream_data_chunks(
            &mut self.next_recv_offset,
            &mut self.recv_pending,
            data.stream_id,
            data.offset,
            data.bytes,
            self.config.initial_window,
        )
    }

    fn ack_data(&mut self, stream_id: u64, received_offset: u64) -> usize {
        let mut acked_bytes = 0usize;
        let mut samples = Vec::new();
        let remove_stream = {
            let Some(sent) = self.sent_data.get_mut(&stream_id) else {
                return 0;
            };
            let acked_offsets = sent
                .iter()
                .filter_map(|(offset, chunk)| {
                    let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
                    (end <= received_offset).then_some(*offset)
                })
                .collect::<Vec<_>>();
            for offset in acked_offsets {
                if let Some(chunk) = sent.remove(&offset) {
                    acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
                    if chunk.attempts == 1 {
                        samples.push(chunk.last_sent.elapsed());
                    }
                }
            }
            sent.is_empty()
        };
        if remove_stream {
            self.sent_data.remove(&stream_id);
        }
        for sample in samples {
            self.observe_rtt_sample(sample);
        }
        acked_bytes
    }

    fn observe_rtt_sample(&mut self, sample: Duration) {
        observe_retransmit_rtt(&mut self.srtt, sample);
    }

    fn add_credit(&mut self, stream_id: u64, bytes: usize) {
        let credit = self.send_credit.entry(stream_id).or_default();
        *credit = credit.saturating_add(bytes).min(self.config.initial_window);
    }

    fn has_stream_state(&self, stream_id: u64) -> bool {
        self.opened_streams.contains(&stream_id)
            || self.pending_opens.contains_key(&stream_id)
            || self.send_credit.contains_key(&stream_id)
            || self.pending_data.contains_key(&stream_id)
            || self.sent_data.contains_key(&stream_id)
            || self.next_send_offset.contains_key(&stream_id)
            || self.next_recv_offset.contains_key(&stream_id)
            || self.recv_pending.contains_key(&stream_id)
            || self.pending_eofs.contains_key(&stream_id)
            || self.pending_window_adjusts.contains_key(&stream_id)
    }

    fn retire_stream(&mut self, stream_id: u64) {
        self.cleanup_stream(stream_id);
        if self.config.retired_stream_tombstones == 0 {
            return;
        }
        if self.retired_streams.insert(stream_id) {
            self.retired_stream_order.push_back(stream_id);
        }
        while self.retired_stream_order.len() > self.config.retired_stream_tombstones {
            if let Some(expired) = self.retired_stream_order.pop_front() {
                self.retired_streams.remove(&expired);
            }
        }
    }

    fn cleanup_stream(&mut self, stream_id: u64) {
        self.opened_streams.remove(&stream_id);
        self.pending_opens.remove(&stream_id);
        self.send_credit.remove(&stream_id);
        self.pending_data.remove(&stream_id);
        self.sent_data.remove(&stream_id);
        self.next_send_offset.remove(&stream_id);
        self.next_recv_offset.remove(&stream_id);
        self.recv_pending.remove(&stream_id);
        self.pending_eofs.remove(&stream_id);
        self.pending_window_adjusts.remove(&stream_id);
    }
}

pub fn adaptive_retransmit_after(configured: Duration, srtt: Option<Duration>) -> Duration {
    let Some(srtt) = srtt else {
        return configured;
    };
    let adaptive = srtt
        .saturating_add(ADAPTIVE_RETRANSMIT_PAD)
        .max(ADAPTIVE_RETRANSMIT_MIN);
    if configured < ADAPTIVE_RETRANSMIT_MIN {
        configured
    } else {
        adaptive.min(configured)
    }
}

pub fn observe_retransmit_rtt(srtt: &mut Option<Duration>, sample: Duration) {
    *srtt = Some(match *srtt {
        None => sample,
        Some(srtt) => {
            let smoothed = ((srtt.as_micros() * 7) + sample.as_micros()) / 8;
            Duration::from_micros(smoothed.min(u64::MAX as u128) as u64)
        }
    });
}

pub fn stream_data_within_receive_window(
    expected_offset: u64,
    offset: u64,
    len: usize,
    receive_window: usize,
) -> bool {
    let Ok(len) = u64::try_from(len) else {
        return false;
    };
    let window = receive_window as u64;
    if len > window {
        return false;
    }
    let Some(end) = offset.checked_add(len) else {
        return false;
    };
    end <= expected_offset.saturating_add(window)
}

fn trim_stream_data_to_offset(
    offset: u64,
    bytes: Vec<u8>,
    expected_offset: u64,
) -> Option<(u64, Vec<u8>)> {
    let len = u64::try_from(bytes.len()).ok()?;
    let end = offset.checked_add(len)?;
    if end <= expected_offset {
        return None;
    }
    if offset >= expected_offset {
        return Some((offset, bytes));
    }
    let skip = usize::try_from(expected_offset - offset).ok()?;
    Some((expected_offset, bytes.get(skip..)?.to_vec()))
}

fn insert_pending_stream_data(pending: &mut BTreeMap<u64, Vec<u8>>, offset: u64, bytes: Vec<u8>) {
    match pending.entry(offset) {
        std::collections::btree_map::Entry::Vacant(entry) => {
            entry.insert(bytes);
        }
        std::collections::btree_map::Entry::Occupied(mut entry) => {
            if bytes.len() > entry.get().len() {
                entry.insert(bytes);
            }
        }
    }
}

pub fn accept_stream_data_chunks(
    next_recv_offset: &mut HashMap<u64, u64>,
    recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
    stream_id: u64,
    offset: u64,
    bytes: Vec<u8>,
    receive_window: usize,
) -> (Vec<Vec<u8>>, usize, u64) {
    let expected = next_recv_offset.entry(stream_id).or_insert(0);
    let Some((offset, bytes)) = trim_stream_data_to_offset(offset, bytes, *expected) else {
        return (Vec::new(), 0, *expected);
    };
    if !stream_data_within_receive_window(*expected, offset, bytes.len(), receive_window) {
        return (Vec::new(), 0, *expected);
    }
    if offset > *expected {
        insert_pending_stream_data(recv_pending.entry(stream_id).or_default(), offset, bytes);
        return (Vec::new(), 0, *expected);
    }

    let mut chunks = vec![bytes];
    let mut consumed = chunks[0].len();
    *expected = expected.saturating_add(consumed as u64);
    while let Some((&pending_offset, _)) = recv_pending
        .get(&stream_id)
        .and_then(|pending| pending.range(..=*expected).next_back())
    {
        let Some(bytes) = recv_pending
            .get_mut(&stream_id)
            .and_then(|pending| pending.remove(&pending_offset))
        else {
            break;
        };
        let Some((trimmed_offset, bytes)) =
            trim_stream_data_to_offset(pending_offset, bytes, *expected)
        else {
            continue;
        };
        if trimmed_offset > *expected {
            insert_pending_stream_data(
                recv_pending.entry(stream_id).or_default(),
                trimmed_offset,
                bytes,
            );
            break;
        }
        consumed = consumed.saturating_add(bytes.len());
        *expected = expected.saturating_add(bytes.len() as u64);
        chunks.push(bytes);
    }
    if recv_pending.get(&stream_id).is_some_and(BTreeMap::is_empty) {
        recv_pending.remove(&stream_id);
    }
    (chunks, consumed, *expected)
}

pub fn stream_data_chunk_limit(receive_window: usize) -> usize {
    MAX_STREAM_DATA_BYTES.min(receive_window.max(1))
}

pub fn split_stream_data_bytes(bytes: Vec<u8>, receive_window: usize) -> Vec<Vec<u8>> {
    if bytes.is_empty() {
        return Vec::new();
    }
    let limit = stream_data_chunk_limit(receive_window);
    if bytes.len() <= limit {
        return vec![bytes];
    }
    bytes.chunks(limit).map(Vec::from).collect()
}

pub struct DoshTransport {
    socket: Arc<UdpSocket>,
    role: SessionRole,
    conn_id: [u8; 16],
    session_key: [u8; 32],
    peer_addr: SocketAddr,
    send_seq: u64,
    ack_seq: u64,
    replay: ReplayWindow,
    mux: StreamMux,
    next_stream_id: u64,
    last_contact: Instant,
    last_keepalive: Instant,
}

impl DoshTransport {
    pub fn new(socket: Arc<UdpSocket>, config: SessionTransportConfig) -> Self {
        let now = Instant::now();
        Self {
            socket,
            role: config.role,
            conn_id: config.conn_id,
            session_key: config.session_key,
            peer_addr: config.peer_addr,
            send_seq: config.initial_send_seq,
            ack_seq: config.initial_ack,
            replay: ReplayWindow::default(),
            mux: StreamMux::new(config.stream),
            next_stream_id: initial_stream_id(config.role),
            last_contact: now,
            last_keepalive: now,
        }
    }

    pub fn new_owned(socket: UdpSocket, config: SessionTransportConfig) -> Self {
        Self::new(Arc::new(socket), config)
    }

    pub fn local_addr(&self) -> Result<SocketAddr> {
        Ok(self.socket.local_addr()?)
    }

    pub fn conn_id(&self) -> [u8; 16] {
        self.conn_id
    }

    pub fn mux(&self) -> &StreamMux {
        &self.mux
    }

    pub fn mux_mut(&mut self) -> &mut StreamMux {
        &mut self.mux
    }

    pub fn peer_addr(&self) -> SocketAddr {
        self.peer_addr
    }

    pub fn set_peer_addr(&mut self, peer_addr: SocketAddr) {
        self.peer_addr = peer_addr;
    }

    pub fn last_contact(&self) -> Instant {
        self.last_contact
    }

    pub fn stale_for(&self) -> Duration {
        self.last_contact.elapsed()
    }

    pub fn allocate_stream_id(&mut self) -> u64 {
        let id = self.next_stream_id;
        self.next_stream_id = self
            .next_stream_id
            .wrapping_add(1)
            .max(initial_stream_id(self.role));
        id
    }

    pub async fn open_service(&mut self, name: &str) -> Result<u64> {
        let stream_id = self.allocate_stream_id();
        let target = service_target(name)?;
        let packet = self.mux.open_stream(stream_id, target, 0)?;
        self.send_outgoing(packet).await?;
        Ok(stream_id)
    }

    pub async fn open_target(
        &mut self,
        target_host: impl Into<String>,
        target_port: u16,
    ) -> Result<u64> {
        let stream_id = self.allocate_stream_id();
        let packet = self
            .mux
            .open_stream(stream_id, target_host.into(), target_port)?;
        self.send_outgoing(packet).await?;
        Ok(stream_id)
    }

    pub async fn accept_stream(&mut self, stream_id: u64) -> Result<()> {
        self.send_outgoing(StreamMux::open_ok(stream_id)?).await
    }

    pub async fn reject_stream(&mut self, stream_id: u64, reason: impl Into<String>) -> Result<()> {
        let packet = self.mux.reject_open(stream_id, reason)?;
        self.send_outgoing(packet).await
    }

    pub async fn send(&mut self, stream_id: u64, bytes: impl Into<Vec<u8>>) -> Result<()> {
        for packet in self.mux.send_data(stream_id, bytes)? {
            self.send_outgoing(packet).await?;
        }
        Ok(())
    }

    pub async fn eof(&mut self, stream_id: u64) -> Result<()> {
        let packet = self.mux.eof_stream(stream_id)?;
        self.send_outgoing(packet).await
    }

    pub async fn close(&mut self, stream_id: u64) -> Result<()> {
        let packet = self.mux.close_stream(stream_id)?;
        self.send_outgoing(packet).await
    }

    pub async fn maintenance(&mut self) -> Result<()> {
        for packet in self.mux.tick()? {
            self.send_outgoing(packet).await?;
        }
        if self.last_keepalive.elapsed() >= self.mux.config.keepalive_after {
            self.send_kind(PacketKind::Ping, b"").await?;
            self.last_keepalive = Instant::now();
        }
        Ok(())
    }

    pub async fn recv(&mut self) -> Result<SessionEvent> {
        let mut buf = vec![0u8; 65535];
        loop {
            let (n, peer) = match tokio::time::timeout(
                ADAPTIVE_RETRANSMIT_MIN,
                self.socket.recv_from(&mut buf),
            )
            .await
            {
                Ok(Ok(value)) => value,
                Ok(Err(err)) if transport_recv_error_is_retryable(&err) => continue,
                Ok(Err(err)) => return Err(err.into()),
                Err(_) => {
                    self.maintenance().await?;
                    continue;
                }
            };
            if let Some(event) = self.accept_received_datagram(&buf[..n], peer).await? {
                return Ok(event);
            }
        }
    }

    async fn accept_received_datagram(
        &mut self,
        datagram: &[u8],
        peer: SocketAddr,
    ) -> Result<Option<SessionEvent>> {
        let packet = match protocol::decode(datagram) {
            Ok(packet) => packet,
            Err(_) => return Ok(None),
        };
        if packet.header.conn_id != self.conn_id {
            return Ok(None);
        }
        let plain =
            match protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction()) {
                Ok(plain) => plain,
                Err(_) => return Ok(None),
            };
        if !self.replay.accept(packet.header.seq) {
            return Ok(None);
        }
        self.peer_addr = peer;
        self.ack_seq = self.ack_seq.max(packet.header.seq);
        self.last_contact = Instant::now();
        self.handle_decrypted_packet(packet.header.kind, &plain)
            .await
            .map(Some)
    }

    pub async fn handle_datagram(
        &mut self,
        datagram: &[u8],
        peer: SocketAddr,
    ) -> Result<SessionEvent> {
        match self.accept_received_datagram(datagram, peer).await? {
            Some(event) => Ok(event),
            None => Ok(SessionEvent::Ignored),
        }
    }

    async fn handle_decrypted_packet(
        &mut self,
        kind: PacketKind,
        plain: &[u8],
    ) -> Result<SessionEvent> {
        match kind {
            PacketKind::Ping => {
                self.send_kind(PacketKind::Pong, b"").await?;
                Ok(SessionEvent::Ping)
            }
            PacketKind::Pong => Ok(SessionEvent::Pong),
            PacketKind::StreamOpen
            | PacketKind::StreamOpenOk
            | PacketKind::StreamOpenReject
            | PacketKind::StreamData
            | PacketKind::StreamWindowAdjust
            | PacketKind::StreamClose
            | PacketKind::StreamEof => {
                let event = self.mux.handle_packet(kind, plain)?;
                self.send_followups(&event).await?;
                Ok(SessionEvent::Stream(event))
            }
            _ => Ok(SessionEvent::Ignored),
        }
    }

    async fn send_followups(&mut self, event: &TransportEvent) -> Result<()> {
        match event {
            TransportEvent::DuplicateOpen { response, .. } => {
                self.send_outgoing(response.clone()).await?;
            }
            TransportEvent::OpenOk { flushed, .. }
            | TransportEvent::WindowAdjust { flushed, .. } => {
                for packet in flushed.clone() {
                    self.send_outgoing(packet).await?;
                }
            }
            TransportEvent::Data(data) => {
                self.send_outgoing(data.window_adjust.clone()).await?;
            }
            TransportEvent::Open(_)
            | TransportEvent::OpenReject { .. }
            | TransportEvent::Eof { .. }
            | TransportEvent::Close { .. }
            | TransportEvent::Ignored { .. } => {}
        }
        Ok(())
    }

    async fn send_outgoing(&mut self, packet: OutgoingStreamPacket) -> Result<()> {
        self.send_kind(packet.kind, &packet.body).await
    }

    async fn send_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result<()> {
        let encoded = self.encode_kind(kind, body)?;
        if let Err(err) = self.socket.send_to(&encoded, self.peer_addr).await {
            if is_transient_udp_send_error(&err) {
                return Ok(());
            }
            return Err(err.into());
        }
        Ok(())
    }

    fn encode_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result<Vec<u8>> {
        let seq = self.send_seq;
        self.send_seq = self.send_seq.saturating_add(1).max(1);
        protocol::encode_encrypted(
            kind,
            self.conn_id,
            seq,
            self.ack_seq,
            &self.session_key,
            self.role.send_direction(),
            body,
        )
    }
}

fn transport_recv_error_is_retryable(err: &std::io::Error) -> bool {
    is_transient_udp_error(err)
}

fn initial_stream_id(role: SessionRole) -> u64 {
    match role {
        SessionRole::Client => 1,
        SessionRole::Server => 1u64 << 63,
    }
}

pub fn service_target(name: &str) -> Result<String> {
    if name.is_empty()
        || !name
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
    {
        bail!("invalid Dosh service name {name:?}");
    }
    Ok(format!("{SERVICE_TARGET_PREFIX}{name}"))
}

pub fn is_service_target(target: &str) -> bool {
    target.starts_with(SERVICE_TARGET_PREFIX)
}

pub fn service_name_from_target(target: &str) -> Option<&str> {
    target.strip_prefix(SERVICE_TARGET_PREFIX)
}

fn encode_packet<T: serde::Serialize>(kind: PacketKind, value: &T) -> Result<OutgoingStreamPacket> {
    Ok(OutgoingStreamPacket {
        kind,
        body: protocol::to_body(value)?,
    })
}

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

    fn decode<T: serde::de::DeserializeOwned>(packet: &OutgoingStreamPacket) -> T {
        protocol::from_body(&packet.body).unwrap()
    }

    #[test]
    fn service_targets_are_reserved_and_validated() {
        assert_eq!(service_target("chat_api").unwrap(), "@dosh-chat_api");
        assert!(is_service_target("@dosh-file"));
        assert_eq!(service_name_from_target("@dosh-chat"), Some("chat"));
        assert!(service_target("../bad").is_err());
    }

    #[test]
    fn transport_recv_retries_network_churn_errors() {
        for kind in [
            std::io::ErrorKind::Interrupted,
            std::io::ErrorKind::TimedOut,
            std::io::ErrorKind::WouldBlock,
        ] {
            let err = std::io::Error::from(kind);
            assert!(transport_recv_error_is_retryable(&err));
        }
        assert!(!transport_recv_error_is_retryable(&std::io::Error::from(
            std::io::ErrorKind::PermissionDenied,
        )));

        #[cfg(unix)]
        for code in [libc::ENETUNREACH, libc::EHOSTUNREACH, libc::ECONNRESET] {
            assert!(transport_recv_error_is_retryable(
                &std::io::Error::from_raw_os_error(code)
            ));
        }

        #[cfg(windows)]
        for code in [10051, 10054, 10065] {
            assert!(transport_recv_error_is_retryable(
                &std::io::Error::from_raw_os_error(code)
            ));
        }
    }

    #[test]
    fn service_registry_matches_registered_service_opens() {
        let mut registry = ServiceRegistry::new();
        registry.register("chat").unwrap();
        let accepted = registry
            .match_open(&IncomingStreamOpen {
                stream_id: 4,
                target_host: "@dosh-chat".to_string(),
                target_port: 0,
            })
            .unwrap();
        assert_eq!(
            accepted,
            AcceptedService {
                name: "chat".to_string(),
                stream_id: 4,
            }
        );
        assert!(
            registry
                .match_open(&IncomingStreamOpen {
                    stream_id: 5,
                    target_host: "@dosh-other".to_string(),
                    target_port: 0,
                })
                .is_err()
        );
    }

    #[test]
    fn opens_stream_and_flushes_queued_data_after_open_ok() {
        let mut mux = StreamMux::new(TransportConfig::default());
        let open = mux
            .open_stream(7, service_target("chat").unwrap(), 0)
            .unwrap();
        assert_eq!(open.kind, PacketKind::StreamOpen);
        assert!(mux.send_data(7, b"hello".to_vec()).unwrap().is_empty());
        assert_eq!(mux.pending_bytes(7), 5);

        let flushed = mux
            .handle_open_ok(StreamOpenOk { stream_id: 7 })
            .unwrap()
            .unwrap();
        assert_eq!(flushed.len(), 1);
        assert_eq!(flushed[0].kind, PacketKind::StreamData);
        let data: StreamData = decode(&flushed[0]);
        assert_eq!(data.offset, 0);
        assert_eq!(data.bytes, b"hello");
    }

    #[test]
    fn receives_out_of_order_data_in_order_and_acks_offsets() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.accept_open(StreamOpen {
            stream_id: 3,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        let first = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 5,
                bytes: b"world".to_vec(),
            })
            .unwrap()
            .unwrap();
        assert!(first.chunks.is_empty());
        assert_eq!(first.received_offset, 0);

        let second = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 0,
                bytes: b"hello".to_vec(),
            })
            .unwrap()
            .unwrap();
        assert_eq!(second.chunks, vec![b"hello".to_vec(), b"world".to_vec()]);
        assert_eq!(second.consumed, 10);
        assert_eq!(second.received_offset, 10);
    }

    #[test]
    fn far_future_data_is_not_buffered_beyond_receive_window() {
        let mut mux = StreamMux::new(TransportConfig {
            initial_window: 5,
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 3,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();

        let too_far = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 5,
                bytes: b"!".to_vec(),
            })
            .unwrap()
            .unwrap();
        assert!(too_far.chunks.is_empty());
        assert_eq!(too_far.consumed, 0);
        assert_eq!(too_far.received_offset, 0);
        assert!(!mux.recv_pending.contains_key(&3));

        let in_window = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 4,
                bytes: b"o".to_vec(),
            })
            .unwrap()
            .unwrap();
        assert!(in_window.chunks.is_empty());
        assert_eq!(in_window.received_offset, 0);
        assert!(
            mux.recv_pending
                .get(&3)
                .is_some_and(|pending| pending.contains_key(&4))
        );

        let oversized_current = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 0,
                bytes: b"abcdef".to_vec(),
            })
            .unwrap()
            .unwrap();
        assert!(oversized_current.chunks.is_empty());
        assert_eq!(oversized_current.received_offset, 0);
    }

    #[test]
    fn overlapping_out_of_order_data_trims_deliverable_suffixes() {
        let mut mux = StreamMux::new(TransportConfig {
            initial_window: 16,
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 3,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();

        assert!(
            mux.handle_data(StreamData {
                stream_id: 3,
                offset: 5,
                bytes: b"FGHIJ".to_vec(),
            })
            .unwrap()
            .unwrap()
            .chunks
            .is_empty()
        );
        assert!(
            mux.handle_data(StreamData {
                stream_id: 3,
                offset: 3,
                bytes: b"DEFGH".to_vec(),
            })
            .unwrap()
            .unwrap()
            .chunks
            .is_empty()
        );

        let delivered = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 0,
                bytes: b"ABC".to_vec(),
            })
            .unwrap()
            .unwrap();

        assert_eq!(
            delivered.chunks.concat(),
            b"ABCDEFGHIJ",
            "overlapping pending suffix should be trimmed and delivered"
        );
        assert_eq!(delivered.consumed, 10);
        assert_eq!(delivered.received_offset, 10);
        assert!(!mux.recv_pending.contains_key(&3));
    }

    #[test]
    fn same_offset_pending_data_keeps_longest_range() {
        let mut next = HashMap::from([(1, 0)]);
        let mut pending = HashMap::new();

        let first = accept_stream_data_chunks(&mut next, &mut pending, 1, 3, b"DE".to_vec(), 16);
        assert!(first.0.is_empty());
        let longer = accept_stream_data_chunks(&mut next, &mut pending, 1, 3, b"DEFG".to_vec(), 16);
        assert!(longer.0.is_empty());

        let delivered =
            accept_stream_data_chunks(&mut next, &mut pending, 1, 0, b"ABC".to_vec(), 16);
        assert_eq!(delivered.0.concat(), b"ABCDEFG");
        assert_eq!(delivered.1, 7);
        assert_eq!(delivered.2, 7);
        assert!(pending.is_empty());
    }

    #[test]
    fn receive_window_helper_rejects_overflowing_ranges() {
        assert!(stream_data_within_receive_window(0, 4, 1, 5));
        assert!(!stream_data_within_receive_window(0, 5, 1, 5));
        assert!(!stream_data_within_receive_window(0, 0, 6, 5));
        assert!(!stream_data_within_receive_window(
            u64::MAX - 1,
            u64::MAX,
            1,
            5
        ));
    }

    #[test]
    fn window_adjust_acks_sent_data_and_releases_credit() {
        let mut mux = StreamMux::new(TransportConfig {
            initial_window: 5,
            retransmit_after: Duration::from_secs(1),
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();
        let sent = mux.send_data(1, b"hello".to_vec()).unwrap();
        assert_eq!(sent.len(), 1);
        assert!(mux.send_data(1, b"!".to_vec()).unwrap().is_empty());
        assert_eq!(mux.pending_bytes(1), 1);
        let flushed = mux
            .handle_window_adjust(StreamWindowAdjust {
                stream_id: 1,
                received_offset: 5,
                bytes: 5,
            })
            .unwrap()
            .unwrap();
        assert_eq!(flushed.len(), 1);
        let data: StreamData = decode(&flushed[0]);
        assert_eq!(data.offset, 5);
        assert_eq!(data.bytes, b"!");
    }

    #[test]
    fn send_data_splits_large_writes_into_packet_sized_chunks() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();

        let bytes = vec![7; MAX_STREAM_DATA_BYTES * 2 + 13];
        let sent = mux.send_data(1, bytes).unwrap();

        assert_eq!(sent.len(), 3);
        let first: StreamData = decode(&sent[0]);
        let second: StreamData = decode(&sent[1]);
        let third: StreamData = decode(&sent[2]);
        assert_eq!(first.offset, 0);
        assert_eq!(first.bytes.len(), MAX_STREAM_DATA_BYTES);
        assert_eq!(second.offset, MAX_STREAM_DATA_BYTES as u64);
        assert_eq!(second.bytes.len(), MAX_STREAM_DATA_BYTES);
        assert_eq!(third.offset, (MAX_STREAM_DATA_BYTES * 2) as u64);
        assert_eq!(third.bytes.len(), 13);
    }

    #[test]
    fn queued_large_write_flushes_in_window_sized_chunks_after_open() {
        let mut mux = StreamMux::new(TransportConfig {
            initial_window: 5,
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        assert!(mux.send_data(1, b"hello!".to_vec()).unwrap().is_empty());

        let flushed = mux
            .handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();

        assert_eq!(flushed.len(), 1);
        let first: StreamData = decode(&flushed[0]);
        assert_eq!(first.offset, 0);
        assert_eq!(first.bytes, b"hello");
        assert_eq!(mux.pending_bytes(1), 1);

        let flushed = mux
            .handle_window_adjust(StreamWindowAdjust {
                stream_id: 1,
                received_offset: 5,
                bytes: 5,
            })
            .unwrap()
            .unwrap();
        assert_eq!(flushed.len(), 1);
        let second: StreamData = decode(&flushed[0]);
        assert_eq!(second.offset, 5);
        assert_eq!(second.bytes, b"!");
    }

    #[test]
    fn stream_open_ack_lowers_retransmit_timeout_from_observed_rtt() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_secs(1),
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-fast", 0).unwrap();
        mux.pending_opens.get_mut(&1).unwrap().last_sent =
            Instant::now() - Duration::from_millis(30);
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();

        assert!(mux.effective_retransmit_after() < Duration::from_millis(100));
        assert!(mux.effective_retransmit_after() >= Duration::from_millis(10));

        mux.open_stream(2, "@dosh-fast", 0).unwrap();
        mux.pending_opens.get_mut(&2).unwrap().last_sent =
            Instant::now() - Duration::from_millis(50);
        let retransmits = mux.tick().unwrap();
        assert_eq!(retransmits.len(), 1);
        assert_eq!(retransmits[0].kind, PacketKind::StreamOpen);
        let open: StreamOpen = decode(&retransmits[0]);
        assert_eq!(open.stream_id, 2);
    }

    #[test]
    fn data_ack_lowers_retransmit_timeout_from_first_attempt_rtt() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_secs(1),
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-fast", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();
        mux.send_data(1, b"hello".to_vec()).unwrap();
        mux.sent_data
            .get_mut(&1)
            .unwrap()
            .get_mut(&0)
            .unwrap()
            .last_sent = Instant::now() - Duration::from_millis(25);

        mux.handle_window_adjust(StreamWindowAdjust {
            stream_id: 1,
            received_offset: 5,
            bytes: 5,
        })
        .unwrap()
        .unwrap();

        assert!(mux.effective_retransmit_after() < Duration::from_millis(100));
    }

    #[test]
    fn retransmitted_data_ack_does_not_update_rtt_estimate() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_secs(1),
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-fast", 0).unwrap();
        mux.pending_opens.get_mut(&1).unwrap().last_sent =
            Instant::now() - Duration::from_millis(30);
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();
        let learned = mux.effective_retransmit_after();

        mux.send_data(1, b"hello".to_vec()).unwrap();
        let chunk = mux.sent_data.get_mut(&1).unwrap().get_mut(&0).unwrap();
        chunk.attempts = 2;
        chunk.last_sent = Instant::now() - Duration::from_millis(900);
        mux.handle_window_adjust(StreamWindowAdjust {
            stream_id: 1,
            received_offset: 5,
            bytes: 5,
        })
        .unwrap()
        .unwrap();

        assert_eq!(mux.effective_retransmit_after(), learned);
    }

    #[test]
    fn cumulative_window_adjust_restores_credit_even_if_delta_was_lost() {
        let mut mux = StreamMux::new(TransportConfig {
            initial_window: 5,
            retransmit_after: Duration::from_secs(1),
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();
        assert_eq!(mux.send_data(1, b"hello".to_vec()).unwrap().len(), 1);
        assert_eq!(mux.send_credit[&1], 0);

        let flushed = mux
            .handle_window_adjust(StreamWindowAdjust {
                stream_id: 1,
                received_offset: 5,
                bytes: 0,
            })
            .unwrap()
            .unwrap();
        assert!(flushed.is_empty());
        assert_eq!(mux.send_credit[&1], 5);
        assert_eq!(mux.send_data(1, b"!".to_vec()).unwrap().len(), 1);
    }

    #[test]
    fn tick_retransmits_pending_stream_eof() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_millis(200),
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();

        let eof = mux.eof_stream(1).unwrap();
        assert_eq!(eof.kind, PacketKind::StreamEof);
        mux.pending_eofs.get_mut(&1).unwrap().last_sent =
            Instant::now() - Duration::from_millis(250);

        let retransmits = mux.tick().unwrap();
        assert_eq!(retransmits.len(), 1);
        assert_eq!(retransmits[0].kind, PacketKind::StreamEof);
        let eof: StreamEof = decode(&retransmits[0]);
        assert_eq!(eof.stream_id, 1);
        assert_eq!(mux.pending_eofs[&1].attempts, 2);
    }

    #[test]
    fn tick_retransmits_pending_stream_close() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_millis(200),
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();

        let close = mux.close_stream(1).unwrap();
        assert_eq!(close.kind, PacketKind::StreamClose);
        mux.pending_closes.get_mut(&1).unwrap().last_sent =
            Instant::now() - Duration::from_millis(250);

        let retransmits = mux.tick().unwrap();
        assert_eq!(retransmits.len(), 1);
        assert_eq!(retransmits[0].kind, PacketKind::StreamClose);
        let close: StreamClose = decode(&retransmits[0]);
        assert_eq!(close.stream_id, 1);
        assert_eq!(mux.pending_closes[&1].attempts, 2);
    }

    #[test]
    fn tick_retransmits_pending_window_adjust() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_millis(200),
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();

        let incoming = mux
            .handle_data(StreamData {
                stream_id: 1,
                offset: 0,
                bytes: b"hello".to_vec(),
            })
            .unwrap()
            .unwrap();
        assert_eq!(incoming.consumed, 5);
        assert_eq!(incoming.received_offset, 5);
        assert_eq!(incoming.window_adjust.kind, PacketKind::StreamWindowAdjust);
        mux.pending_window_adjusts.get_mut(&1).unwrap().last_sent =
            Instant::now() - Duration::from_millis(250);

        let retransmits = mux.tick().unwrap();
        assert_eq!(retransmits.len(), 1);
        assert_eq!(retransmits[0].kind, PacketKind::StreamWindowAdjust);
        let adjust: StreamWindowAdjust = decode(&retransmits[0]);
        assert_eq!(adjust.stream_id, 1);
        assert_eq!(adjust.received_offset, 5);
        assert_eq!(adjust.bytes, 5);
        assert_eq!(mux.pending_window_adjusts[&1].attempts, 2);
    }

    #[test]
    fn stream_control_retransmits_expire_after_cap() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_millis(200),
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        mux.close_stream(1).unwrap();
        let pending = mux.pending_closes.get_mut(&1).unwrap();
        pending.last_sent = Instant::now() - Duration::from_millis(250);
        pending.attempts = STREAM_CONTROL_RETRANSMIT_MAX_ATTEMPTS;

        assert!(mux.tick().unwrap().is_empty());
        assert!(!mux.pending_closes.contains_key(&1));
    }

    #[test]
    fn cleanup_stream_removes_reliable_controls() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.accept_open(StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        mux.eof_stream(1).unwrap();
        mux.handle_data(StreamData {
            stream_id: 1,
            offset: 0,
            bytes: b"hello".to_vec(),
        })
        .unwrap()
        .unwrap();

        mux.retire_stream(1);

        assert!(!mux.pending_eofs.contains_key(&1));
        assert!(!mux.pending_window_adjusts.contains_key(&1));
        assert!(!mux.is_open(1));
    }

    #[test]
    fn peer_close_after_local_close_stops_close_retransmits() {
        let mut mux = StreamMux::new(TransportConfig {
            retransmit_after: Duration::from_millis(200),
            ..TransportConfig::default()
        });
        mux.accept_open(StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        mux.close_stream(1).unwrap();

        assert!(!mux.handle_close(StreamClose { stream_id: 1 }));
        assert!(!mux.pending_closes.contains_key(&1));
        assert!(mux.tick().unwrap().is_empty());
    }

    #[test]
    fn handle_packet_decodes_and_dispatches_stream_packets() {
        let mut mux = StreamMux::new(TransportConfig::default());
        let body = protocol::to_body(&StreamOpen {
            stream_id: 9,
            target_host: "@dosh-app".to_string(),
            target_port: 0,
        })
        .unwrap();
        let event = mux.handle_packet(PacketKind::StreamOpen, &body).unwrap();
        assert_eq!(
            event,
            TransportEvent::Open(IncomingStreamOpen {
                stream_id: 9,
                target_host: "@dosh-app".to_string(),
                target_port: 0,
            })
        );
    }

    #[test]
    fn duplicate_stream_open_reacks_without_second_open_event() {
        let mut mux = StreamMux::new(TransportConfig::default());
        let open = StreamOpen {
            stream_id: 9,
            target_host: "@dosh-app".to_string(),
            target_port: 0,
        };
        let body = protocol::to_body(&open).unwrap();
        assert!(matches!(
            mux.handle_packet(PacketKind::StreamOpen, &body).unwrap(),
            TransportEvent::Open(_)
        ));

        let duplicate = mux.handle_packet(PacketKind::StreamOpen, &body).unwrap();
        match duplicate {
            TransportEvent::DuplicateOpen {
                stream_id,
                response,
            } => {
                assert_eq!(stream_id, 9);
                assert_eq!(response.kind, PacketKind::StreamOpenOk);
                let ok: StreamOpenOk = decode(&response);
                assert_eq!(ok.stream_id, 9);
            }
            other => panic!("unexpected duplicate open event {other:?}"),
        }
    }

    #[test]
    fn send_data_to_closed_or_unknown_stream_fails() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        assert!(mux.send_data(1, b"queued".to_vec()).unwrap().is_empty());
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();
        mux.close_stream(1).unwrap();

        assert!(mux.send_data(1, b"lost".to_vec()).is_err());
        assert!(mux.send_data(99, b"lost".to_vec()).is_err());
    }

    #[test]
    fn stream_eof_is_delivered_without_closing_stream_state() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();

        let eof = mux.eof_stream(1).unwrap();
        assert_eq!(eof.kind, PacketKind::StreamEof);
        assert!(mux.is_open(1));

        let body = protocol::to_body(&StreamEof { stream_id: 1 }).unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamEof, &body).unwrap(),
            TransportEvent::Eof { stream_id: 1 }
        );
        assert!(mux.is_open(1));
    }

    #[test]
    fn rejecting_incoming_stream_retires_local_state() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.accept_open(StreamOpen {
            stream_id: 7,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        assert!(mux.is_open(7));

        let reject = mux.reject_open(7, "nope").unwrap();

        assert_eq!(reject.kind, PacketKind::StreamOpenReject);
        assert!(!mux.is_open(7));
        assert!(mux.retired_streams.contains(&7));
        assert!(mux.send_data(7, b"must-not-send".to_vec()).is_err());
        assert!(
            mux.accept_open(StreamOpen {
                stream_id: 7,
                target_host: "@dosh-test".to_string(),
                target_port: 0,
            })
            .is_err()
        );
    }

    #[test]
    fn stale_packets_after_close_do_not_recreate_stream_state() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 })
            .unwrap()
            .unwrap();
        assert_eq!(mux.send_data(1, b"hello".to_vec()).unwrap().len(), 1);
        mux.close_stream(1).unwrap();

        let stale_data = protocol::to_body(&StreamData {
            stream_id: 1,
            offset: 0,
            bytes: b"stale".to_vec(),
        })
        .unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamData, &stale_data)
                .unwrap(),
            TransportEvent::Ignored { stream_id: 1 }
        );

        let stale_open_ok = protocol::to_body(&StreamOpenOk { stream_id: 1 }).unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamOpenOk, &stale_open_ok)
                .unwrap(),
            TransportEvent::Ignored { stream_id: 1 }
        );

        let stale_adjust = protocol::to_body(&StreamWindowAdjust {
            stream_id: 1,
            received_offset: 5,
            bytes: 5,
        })
        .unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamWindowAdjust, &stale_adjust)
                .unwrap(),
            TransportEvent::Ignored { stream_id: 1 }
        );

        let stale_open = protocol::to_body(&StreamOpen {
            stream_id: 1,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamOpen, &stale_open)
                .unwrap(),
            TransportEvent::Ignored { stream_id: 1 }
        );

        let stale_close = protocol::to_body(&StreamClose { stream_id: 1 }).unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamClose, &stale_close)
                .unwrap(),
            TransportEvent::Ignored { stream_id: 1 }
        );

        assert!(!mux.is_open(1));
        assert!(!mux.pending_opens.contains_key(&1));
        assert!(!mux.send_credit.contains_key(&1));
        assert!(!mux.pending_data.contains_key(&1));
        assert!(!mux.sent_data.contains_key(&1));
        assert!(!mux.next_send_offset.contains_key(&1));
        assert!(!mux.next_recv_offset.contains_key(&1));
        assert!(!mux.recv_pending.contains_key(&1));
        assert!(mux.retired_streams.contains(&1));
    }

    #[test]
    fn retired_stream_tombstones_are_bounded() {
        let mut mux = StreamMux::new(TransportConfig {
            retired_stream_tombstones: 2,
            ..TransportConfig::default()
        });
        for stream_id in 1..=3 {
            mux.open_stream(stream_id, "@dosh-test", 0).unwrap();
            mux.close_stream(stream_id).unwrap();
        }

        assert_eq!(mux.retired_streams.len(), 2);
        assert!(!mux.retired_streams.contains(&1));
        assert!(mux.retired_streams.contains(&2));
        assert!(mux.retired_streams.contains(&3));

        let stale_open = protocol::to_body(&StreamOpen {
            stream_id: 3,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        assert_eq!(
            mux.handle_packet(PacketKind::StreamOpen, &stale_open)
                .unwrap(),
            TransportEvent::Ignored { stream_id: 3 }
        );
    }

    #[tokio::test]
    async fn session_transport_partitions_client_and_server_stream_ids() {
        let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_addr = client_socket.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [44u8; 32];
        let conn_id = [8u8; 16];
        let mut client = DoshTransport::new_owned(
            client_socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: server_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: client_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );

        assert_eq!(client.allocate_stream_id(), 1);
        assert_eq!(client.allocate_stream_id(), 2);
        assert_eq!(server.allocate_stream_id(), 1u64 << 63);
        assert_eq!(server.allocate_stream_id(), (1u64 << 63) + 1);
    }

    #[tokio::test]
    async fn session_transport_handles_simultaneous_opens_without_id_collision() {
        let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_addr = client_socket.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [45u8; 32];
        let conn_id = [9u8; 16];
        let mut client = DoshTransport::new_owned(
            client_socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: server_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: client_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );

        let client_stream = client.open_service("client_app").await.unwrap();
        let server_stream = server.open_service("server_app").await.unwrap();
        assert_ne!(client_stream, server_stream);

        let server_seen = server.recv().await.unwrap();
        let client_seen = client.recv().await.unwrap();
        match server_seen {
            SessionEvent::Stream(TransportEvent::Open(open)) => {
                assert_eq!(open.stream_id, client_stream);
                assert_eq!(open.target_host, "@dosh-client_app");
                server.accept_stream(open.stream_id).await.unwrap();
            }
            other => panic!("unexpected server event {other:?}"),
        }
        match client_seen {
            SessionEvent::Stream(TransportEvent::Open(open)) => {
                assert_eq!(open.stream_id, server_stream);
                assert_eq!(open.target_host, "@dosh-server_app");
                client.accept_stream(open.stream_id).await.unwrap();
            }
            other => panic!("unexpected client event {other:?}"),
        }

        assert!(matches!(
            client.recv().await.unwrap(),
            SessionEvent::Stream(TransportEvent::OpenOk { stream_id, .. }) if stream_id == client_stream
        ));
        assert!(matches!(
            server.recv().await.unwrap(),
            SessionEvent::Stream(TransportEvent::OpenOk { stream_id, .. }) if stream_id == server_stream
        ));
    }

    #[tokio::test]
    async fn session_transport_opens_service_and_moves_data_over_encrypted_udp() {
        let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_addr = client_socket.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [42u8; 32];
        let conn_id = [7u8; 16];
        let mut client = DoshTransport::new_owned(
            client_socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: server_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: client_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );

        let stream_id = client.open_service("app").await.unwrap();
        let event = server.recv().await.unwrap();
        assert_eq!(server.peer_addr(), client_addr);
        match event {
            SessionEvent::Stream(TransportEvent::Open(open)) => {
                assert_eq!(open.stream_id, stream_id);
                assert_eq!(open.target_host, "@dosh-app");
                server.accept_stream(open.stream_id).await.unwrap();
            }
            other => panic!("unexpected event {other:?}"),
        }

        assert!(matches!(
            client.recv().await.unwrap(),
            SessionEvent::Stream(TransportEvent::OpenOk { .. })
        ));
        client.send(stream_id, b"hello".to_vec()).await.unwrap();
        let event = server.recv().await.unwrap();
        match event {
            SessionEvent::Stream(TransportEvent::Data(data)) => {
                assert_eq!(data.stream_id, stream_id);
                assert_eq!(data.chunks, vec![b"hello".to_vec()]);
            }
            other => panic!("unexpected event {other:?}"),
        }
    }

    #[tokio::test]
    async fn session_transport_recv_drives_retransmit_while_waiting() {
        let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_addr = client_socket.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [48u8; 32];
        let conn_id = [12u8; 16];
        let stream_config = TransportConfig {
            retransmit_after: Duration::from_millis(20),
            keepalive_after: Duration::from_secs(60),
            ..TransportConfig::default()
        };
        let mut client = DoshTransport::new_owned(
            client_socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: server_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: stream_config.clone(),
            },
        );

        let stream_id = client.open_service("app").await.unwrap();

        let mut dropped = [0u8; 65535];
        tokio::time::timeout(
            Duration::from_millis(200),
            server_socket.recv_from(&mut dropped),
        )
        .await
        .unwrap()
        .unwrap();

        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: client_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: stream_config,
            },
        );

        let client_wait = tokio::spawn(async move { client.recv().await });
        match tokio::time::timeout(Duration::from_secs(1), server.recv())
            .await
            .unwrap()
            .unwrap()
        {
            SessionEvent::Stream(TransportEvent::Open(open)) => {
                assert_eq!(open.stream_id, stream_id);
                server.accept_stream(open.stream_id).await.unwrap();
            }
            other => panic!("unexpected server event after lost open: {other:?}"),
        }

        assert!(matches!(
            tokio::time::timeout(Duration::from_secs(1), client_wait)
                .await
                .unwrap()
                .unwrap()
                .unwrap(),
            SessionEvent::Stream(TransportEvent::OpenOk { stream_id: ok, .. }) if ok == stream_id
        ));
    }

    #[tokio::test]
    async fn session_transport_reject_stream_retires_server_state() {
        let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_addr = client_socket.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [46u8; 32];
        let conn_id = [10u8; 16];
        let mut client = DoshTransport::new_owned(
            client_socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: server_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: client_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );

        let stream_id = client.open_service("app").await.unwrap();
        match server.recv().await.unwrap() {
            SessionEvent::Stream(TransportEvent::Open(open)) => {
                assert_eq!(open.stream_id, stream_id);
                assert!(server.mux().is_open(stream_id));
                server.reject_stream(stream_id, "disabled").await.unwrap();
            }
            other => panic!("unexpected event {other:?}"),
        }

        assert!(!server.mux().is_open(stream_id));
        assert!(
            server
                .send(stream_id, b"must-not-send".to_vec())
                .await
                .is_err()
        );
        assert!(matches!(
            client.recv().await.unwrap(),
            SessionEvent::Stream(TransportEvent::OpenReject { stream_id: rejected, reason })
                if rejected == stream_id && reason == "disabled"
        ));
    }

    #[tokio::test]
    async fn session_transport_roams_to_authenticated_new_peer_addr() {
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let original_client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let roaming_client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let original_addr = original_client.local_addr().unwrap();
        let roaming_addr = roaming_client.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [9u8; 32];
        let conn_id = [3u8; 16];
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: original_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let ping = protocol::encode_encrypted(
            PacketKind::Ping,
            conn_id,
            1,
            0,
            &key,
            CLIENT_TO_SERVER,
            b"",
        )
        .unwrap();
        roaming_client.send_to(&ping, server_addr).await.unwrap();
        assert!(matches!(server.recv().await.unwrap(), SessionEvent::Ping));
        assert_eq!(server.peer_addr(), roaming_addr);
    }

    #[tokio::test]
    async fn unauthenticated_datagram_does_not_poison_replay_window() {
        let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let peer: SocketAddr = "127.0.0.1:9".parse().unwrap();
        let key = [11u8; 32];
        let wrong_key = [12u8; 32];
        let conn_id = [4u8; 16];
        let mut transport = DoshTransport::new_owned(
            socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: peer,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let bad = protocol::encode_encrypted(
            PacketKind::Pong,
            conn_id,
            9,
            0,
            &wrong_key,
            SERVER_TO_CLIENT,
            b"",
        )
        .unwrap();
        let valid = protocol::encode_encrypted(
            PacketKind::Pong,
            conn_id,
            9,
            0,
            &key,
            SERVER_TO_CLIENT,
            b"",
        )
        .unwrap();

        assert!(matches!(
            transport
                .handle_datagram(b"not a dosh packet", peer)
                .await
                .unwrap(),
            SessionEvent::Ignored
        ));
        assert!(matches!(
            transport.handle_datagram(&bad, peer).await.unwrap(),
            SessionEvent::Ignored
        ));
        assert!(matches!(
            transport.handle_datagram(&valid, peer).await.unwrap(),
            SessionEvent::Pong
        ));
    }
}