//! Session-layer message types: setup, ack, data, and error messages.

use super::ProtocolError;
use crate::NodeAddr;
use crate::tree::TreeCoordinate;
use std::fmt;

// ============================================================================
// Session Layer Message Types
// ============================================================================

/// SessionDatagram payload message type identifiers.
///
/// These messages are carried as payloads inside `SessionDatagram` (link
/// message type 0x00). Post-handshake messages (data, reports) are end-to-end
/// encrypted with session keys via the FSP pipeline. Error signals
/// (CoordsRequired, PathBroken) are plaintext messages generated by transit
/// routers that cannot establish e2e sessions with the source.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum SessionMessageType {
    // Session establishment (0x00-0x0F)
    /// Session setup with coordinates (warms router caches).
    SessionSetup = 0x00,
    /// Session acknowledgement.
    SessionAck = 0x01,

    // Data and metrics (0x10-0x1F) — encrypted, inner header msg_type
    /// Port-multiplexed service payload: `[src_port:2 LE][dst_port:2 LE][service data...]`.
    /// Port 256 = IPv6 shim (compressed header). Receiver dispatches by dst_port.
    DataPacket = 0x10,
    /// MMP sender report (metrics from sender to receiver).
    SenderReport = 0x11,
    /// MMP receiver report (metrics from receiver to sender).
    ReceiverReport = 0x12,
    /// Path MTU notification (discovered path MTU).
    PathMtuNotification = 0x13,
    /// Standalone coordinate cache warming (empty body, coords in CP flag).
    CoordsWarmup = 0x14,

    // Link-layer error signals (0x20-0x2F) — plaintext, from transit routers
    /// Router cache miss — needs coordinates (link-layer error signal).
    CoordsRequired = 0x20,
    /// Routing failure — local minimum or unreachable (link-layer error signal).
    PathBroken = 0x21,
    /// MTU exceeded — forwarded packet too large for next-hop transport (link-layer error signal).
    MtuExceeded = 0x22,
}

impl SessionMessageType {
    /// Try to convert from a byte.
    pub fn from_byte(b: u8) -> Option<Self> {
        match b {
            0x00 => Some(SessionMessageType::SessionSetup),
            0x01 => Some(SessionMessageType::SessionAck),
            0x10 => Some(SessionMessageType::DataPacket),
            0x11 => Some(SessionMessageType::SenderReport),
            0x12 => Some(SessionMessageType::ReceiverReport),
            0x13 => Some(SessionMessageType::PathMtuNotification),
            0x14 => Some(SessionMessageType::CoordsWarmup),
            0x20 => Some(SessionMessageType::CoordsRequired),
            0x21 => Some(SessionMessageType::PathBroken),
            0x22 => Some(SessionMessageType::MtuExceeded),
            _ => None,
        }
    }

    /// Convert to a byte.
    pub fn to_byte(self) -> u8 {
        self as u8
    }
}

impl fmt::Display for SessionMessageType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            SessionMessageType::SessionSetup => "SessionSetup",
            SessionMessageType::SessionAck => "SessionAck",
            SessionMessageType::DataPacket => "DataPacket",
            SessionMessageType::SenderReport => "SenderReport",
            SessionMessageType::ReceiverReport => "ReceiverReport",
            SessionMessageType::PathMtuNotification => "PathMtuNotification",
            SessionMessageType::CoordsWarmup => "CoordsWarmup",
            SessionMessageType::CoordsRequired => "CoordsRequired",
            SessionMessageType::PathBroken => "PathBroken",
            SessionMessageType::MtuExceeded => "MtuExceeded",
        };
        write!(f, "{}", name)
    }
}

// ============================================================================
// Coordinate Wire Format Helpers
// ============================================================================

/// Wire size of a TreeCoordinate in address-only format: 2 + entries × 16.
pub(crate) fn coords_wire_size(coords: &TreeCoordinate) -> usize {
    2 + coords.entries().len() * 16
}

/// Encode a TreeCoordinate as address-only wire format: count(u16 LE) + addrs(16 × n).
///
/// Session-layer messages serialize coordinates as NodeAddr arrays (16 bytes each),
/// without the sequence/timestamp metadata used by the tree gossip protocol.
pub(crate) fn encode_coords(coords: &TreeCoordinate, buf: &mut Vec<u8>) {
    let addrs: Vec<&NodeAddr> = coords.node_addrs().collect();
    let count = addrs.len() as u16;
    buf.extend_from_slice(&count.to_le_bytes());
    for addr in addrs {
        buf.extend_from_slice(addr.as_bytes());
    }
}

/// Decode a TreeCoordinate from address-only wire format.
///
/// Returns the decoded coordinate and the number of bytes consumed.
pub(crate) fn decode_coords(data: &[u8]) -> Result<(TreeCoordinate, usize), ProtocolError> {
    if data.len() < 2 {
        return Err(ProtocolError::MessageTooShort {
            expected: 2,
            got: data.len(),
        });
    }
    let count = u16::from_le_bytes([data[0], data[1]]) as usize;
    let needed = 2 + count * 16;
    if data.len() < needed {
        return Err(ProtocolError::MessageTooShort {
            expected: needed,
            got: data.len(),
        });
    }
    if count == 0 {
        return Err(ProtocolError::Malformed(
            "coordinate with zero entries".into(),
        ));
    }
    let mut addrs = Vec::with_capacity(count);
    for i in 0..count {
        let offset = 2 + i * 16;
        let mut bytes = [0u8; 16];
        bytes.copy_from_slice(&data[offset..offset + 16]);
        addrs.push(NodeAddr::from_bytes(bytes));
    }
    let coord =
        TreeCoordinate::from_addrs(addrs).map_err(|e| ProtocolError::Malformed(e.to_string()))?;
    Ok((coord, needed))
}

/// Decode an optional coordinate field (count may be 0).
///
/// Returns None if count is 0, Some(coord) otherwise, plus bytes consumed.
pub(crate) fn decode_optional_coords(
    data: &[u8],
) -> Result<(Option<TreeCoordinate>, usize), ProtocolError> {
    if data.len() < 2 {
        return Err(ProtocolError::MessageTooShort {
            expected: 2,
            got: data.len(),
        });
    }
    let count = u16::from_le_bytes([data[0], data[1]]) as usize;
    let needed = 2 + count * 16;
    if data.len() < needed {
        return Err(ProtocolError::MessageTooShort {
            expected: needed,
            got: data.len(),
        });
    }
    if count == 0 {
        return Ok((None, 2));
    }
    let mut addrs = Vec::with_capacity(count);
    for i in 0..count {
        let offset = 2 + i * 16;
        let mut bytes = [0u8; 16];
        bytes.copy_from_slice(&data[offset..offset + 16]);
        addrs.push(NodeAddr::from_bytes(bytes));
    }
    let coord =
        TreeCoordinate::from_addrs(addrs).map_err(|e| ProtocolError::Malformed(e.to_string()))?;
    Ok((Some(coord), needed))
}

/// Encode a count of zero (for empty/absent coordinate fields).
fn encode_empty_coords(buf: &mut Vec<u8>) {
    buf.extend_from_slice(&0u16.to_le_bytes());
}

// ============================================================================
// Session Flags
// ============================================================================

/// Session flags for setup options.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct SessionFlags {
    /// Request acknowledgement from destination.
    pub request_ack: bool,
    /// Set up bidirectional session.
    pub bidirectional: bool,
}

impl SessionFlags {
    /// Create default flags.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set request_ack flag.
    pub fn with_ack(mut self) -> Self {
        self.request_ack = true;
        self
    }

    /// Set bidirectional flag.
    pub fn bidirectional(mut self) -> Self {
        self.bidirectional = true;
        self
    }

    /// Convert to a byte.
    pub fn to_byte(&self) -> u8 {
        let mut flags = 0u8;
        if self.request_ack {
            flags |= 0x01;
        }
        if self.bidirectional {
            flags |= 0x02;
        }
        flags
    }

    /// Convert from a byte.
    pub fn from_byte(byte: u8) -> Self {
        Self {
            request_ack: byte & 0x01 != 0,
            bidirectional: byte & 0x02 != 0,
        }
    }
}

// ============================================================================
// FSP Packet Flags
// ============================================================================

/// FSP common prefix flags (cleartext, in outer header).
///
/// | Bit | Name | Description                                    |
/// |-----|------|------------------------------------------------|
/// | 0   | CP   | Coords present between header and ciphertext   |
/// | 1   | K    | Key epoch (for rekeying)                       |
/// | 2   | U    | Unencrypted payload (error signals)            |
/// | 3-7 |      | Reserved                                       |
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct FspFlags {
    /// Coordinates present between header and ciphertext.
    pub coords_present: bool,
    /// Key epoch bit for rekeying.
    pub key_epoch: bool,
    /// Unencrypted payload (plaintext error signals from transit routers).
    pub unencrypted: bool,
}

impl FspFlags {
    /// Create default flags (all clear).
    pub fn new() -> Self {
        Self::default()
    }

    /// Convert to a byte.
    pub fn to_byte(&self) -> u8 {
        let mut flags = 0u8;
        if self.coords_present {
            flags |= 0x01;
        }
        if self.key_epoch {
            flags |= 0x02;
        }
        if self.unencrypted {
            flags |= 0x04;
        }
        flags
    }

    /// Convert from a byte.
    pub fn from_byte(byte: u8) -> Self {
        Self {
            coords_present: byte & 0x01 != 0,
            key_epoch: byte & 0x02 != 0,
            unencrypted: byte & 0x04 != 0,
        }
    }
}

/// FSP inner header flags (encrypted, inside AEAD envelope).
///
/// | Bit | Name | Description                     |
/// |-----|------|---------------------------------|
/// | 0-7 |      | Reserved (all zero)             |
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct FspInnerFlags;

impl FspInnerFlags {
    /// Create default inner flags (all clear).
    pub fn new() -> Self {
        Self
    }

    /// Convert to a byte.
    pub fn to_byte(&self) -> u8 {
        0x00
    }

    /// Convert from a byte.
    pub fn from_byte(_byte: u8) -> Self {
        Self
    }
}

// ============================================================================
// Session Setup
// ============================================================================

/// Session setup to establish cached coordinate state.
///
/// Carried inside a SessionDatagram envelope which provides src_addr and
/// dest_addr. The SessionSetup payload contains coordinates, session flags,
/// and the Noise XX handshake message for session establishment.
///
/// ## Wire Format
///
/// | Offset | Field            | Size    | Description                         |
/// |--------|------------------|---------|-------------------------------------|
/// | 0      | msg_type         | 1 byte  | 0x00                                |
/// | 1      | flags            | 1 byte  | Bit 0: REQUEST_ACK, Bit 1: BIDIR   |
/// | 2      | src_coords_count | 2 bytes | u16 LE, number of src coord entries |
/// | 4      | src_coords       | 16 × n  | NodeAddr array (self → root)        |
/// | ...    | dest_coords_count| 2 bytes | u16 LE, number of dest coord entries|
/// | ...    | dest_coords      | 16 × m  | NodeAddr array (dest → root)        |
/// | ...    | handshake_len    | 2 bytes  | u16 LE, Noise payload length        |
/// | ...    | handshake_payload| variable| Noise XX msg1 (33 bytes — ephemeral key)    |
#[derive(Clone, Debug)]
pub struct SessionSetup {
    /// Source coordinates (for return path caching).
    pub src_coords: TreeCoordinate,
    /// Destination coordinates (for forward routing).
    pub dest_coords: TreeCoordinate,
    /// Session options.
    pub flags: SessionFlags,
    /// Noise XX handshake message 1.
    pub handshake_payload: Vec<u8>,
}

impl SessionSetup {
    /// Create a new session setup message.
    pub fn new(src_coords: TreeCoordinate, dest_coords: TreeCoordinate) -> Self {
        Self {
            src_coords,
            dest_coords,
            flags: SessionFlags::new(),
            handshake_payload: Vec::new(),
        }
    }

    /// Set session flags.
    pub fn with_flags(mut self, flags: SessionFlags) -> Self {
        self.flags = flags;
        self
    }

    /// Set the Noise handshake payload.
    pub fn with_handshake(mut self, payload: Vec<u8>) -> Self {
        self.handshake_payload = payload;
        self
    }

    /// Encode as wire format (4-byte FSP prefix + flags + coords + handshake).
    ///
    /// The 4-byte prefix: `[ver_phase:1][flags:1][payload_len:2 LE]`
    /// where ver_phase = 0x01 (version 0, phase MSG1).
    pub fn encode(&self) -> Vec<u8> {
        // Build body first to compute payload_len
        let mut body = Vec::new();
        body.push(self.flags.to_byte());
        encode_coords(&self.src_coords, &mut body);
        encode_coords(&self.dest_coords, &mut body);
        let hs_len = self.handshake_payload.len() as u16;
        body.extend_from_slice(&hs_len.to_le_bytes());
        body.extend_from_slice(&self.handshake_payload);

        // Prepend 4-byte FSP common prefix
        let payload_len = body.len() as u16;
        let mut buf = Vec::with_capacity(4 + body.len());
        buf.push(0x01); // version 0, phase 0x1 (MSG1)
        buf.push(0x00); // flags (must be zero for handshake)
        buf.extend_from_slice(&payload_len.to_le_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    /// Decode from wire format (after 4-byte FSP prefix has been consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
        if payload.is_empty() {
            return Err(ProtocolError::MessageTooShort {
                expected: 1,
                got: 0,
            });
        }
        let flags = SessionFlags::from_byte(payload[0]);
        let mut offset = 1;

        let (src_coords, consumed) = decode_coords(&payload[offset..])?;
        offset += consumed;

        let (dest_coords, consumed) = decode_coords(&payload[offset..])?;
        offset += consumed;

        if payload.len() < offset + 2 {
            return Err(ProtocolError::MessageTooShort {
                expected: offset + 2,
                got: payload.len(),
            });
        }
        let hs_len = u16::from_le_bytes([payload[offset], payload[offset + 1]]) as usize;
        offset += 2;

        if payload.len() < offset + hs_len {
            return Err(ProtocolError::MessageTooShort {
                expected: offset + hs_len,
                got: payload.len(),
            });
        }
        let handshake_payload = payload[offset..offset + hs_len].to_vec();

        Ok(Self {
            src_coords,
            dest_coords,
            flags,
            handshake_payload,
        })
    }
}

// ============================================================================
// Session Ack
// ============================================================================

/// Session acknowledgement.
///
/// Carried inside a SessionDatagram envelope which provides src_addr and
/// dest_addr. The SessionAck payload contains both the acknowledger's and
/// initiator's coordinates for route cache warming (ensuring return-path
/// transit nodes can route independently of the forward path) and the Noise
/// XX handshake response.
///
/// ## Wire Format
///
/// | Offset | Field            | Size    | Description                         |
/// |--------|------------------|---------|-------------------------------------|
/// | 0      | msg_type         | 1 byte  | 0x01                                |
/// | 1      | flags            | 1 byte  | Reserved                            |
/// | 2      | src_coords_count | 2 bytes | u16 LE                              |
/// | 4      | src_coords       | 16 × n  | Acknowledger's coords (for caching) |
/// | ...    | dest_coords_count| 2 bytes | u16 LE                              |
/// | ...    | dest_coords      | 16 × m  | Initiator's coords (for return path)|
/// | ...    | handshake_len    | 2 bytes  | u16 LE, Noise payload length        |
/// | ...    | handshake_payload| variable| Noise XX msg2 (106+ bytes — ephemeral + static + epoch + negotiation)    |
#[derive(Clone, Debug)]
pub struct SessionAck {
    /// Acknowledger's coordinates.
    pub src_coords: TreeCoordinate,
    /// Initiator's coordinates (for return-path cache warming).
    pub dest_coords: TreeCoordinate,
    /// Reserved flags byte (for forward compatibility).
    pub flags: u8,
    /// Noise XX handshake message 2.
    pub handshake_payload: Vec<u8>,
}

impl SessionAck {
    /// Create a new session acknowledgement.
    pub fn new(src_coords: TreeCoordinate, dest_coords: TreeCoordinate) -> Self {
        Self {
            src_coords,
            dest_coords,
            flags: 0,
            handshake_payload: Vec::new(),
        }
    }

    /// Set the Noise handshake payload.
    pub fn with_handshake(mut self, payload: Vec<u8>) -> Self {
        self.handshake_payload = payload;
        self
    }

    /// Encode as wire format (4-byte FSP prefix + flags + coords + handshake).
    ///
    /// The 4-byte prefix: `[ver_phase:1][flags:1][payload_len:2 LE]`
    /// where ver_phase = 0x02 (version 0, phase MSG2).
    pub fn encode(&self) -> Vec<u8> {
        // Build body first to compute payload_len
        let mut body = Vec::new();
        body.push(self.flags);
        encode_coords(&self.src_coords, &mut body);
        encode_coords(&self.dest_coords, &mut body);
        let hs_len = self.handshake_payload.len() as u16;
        body.extend_from_slice(&hs_len.to_le_bytes());
        body.extend_from_slice(&self.handshake_payload);

        // Prepend 4-byte FSP common prefix
        let payload_len = body.len() as u16;
        let mut buf = Vec::with_capacity(4 + body.len());
        buf.push(0x02); // version 0, phase 0x2 (MSG2)
        buf.push(0x00); // flags (must be zero for handshake)
        buf.extend_from_slice(&payload_len.to_le_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    /// Decode from wire format (after 4-byte FSP prefix has been consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
        if payload.is_empty() {
            return Err(ProtocolError::MessageTooShort {
                expected: 1,
                got: 0,
            });
        }
        let flags = payload[0];
        let mut offset = 1;

        let (src_coords, consumed) = decode_coords(&payload[offset..])?;
        offset += consumed;

        let (dest_coords, consumed) = decode_coords(&payload[offset..])?;
        offset += consumed;

        if payload.len() < offset + 2 {
            return Err(ProtocolError::MessageTooShort {
                expected: offset + 2,
                got: payload.len(),
            });
        }
        let hs_len = u16::from_le_bytes([payload[offset], payload[offset + 1]]) as usize;
        offset += 2;

        if payload.len() < offset + hs_len {
            return Err(ProtocolError::MessageTooShort {
                expected: offset + hs_len,
                got: payload.len(),
            });
        }
        let handshake_payload = payload[offset..offset + hs_len].to_vec();

        Ok(Self {
            src_coords,
            dest_coords,
            flags,
            handshake_payload,
        })
    }
}

// ============================================================================
// Session Msg3 (XX Handshake Message 3)
// ============================================================================

/// XX handshake message 3 (initiator -> responder).
///
/// Carries the initiator's encrypted static key and epoch. Sent by the
/// initiator after receiving msg2. The responder learns the initiator's
/// identity from this message.
///
/// ## Wire Format
///
/// | Offset | Field            | Size    | Description                         |
/// |--------|------------------|---------|-------------------------------------|
/// | 0      | flags            | 1 byte  | Reserved                            |
/// | 1      | handshake_len    | 2 bytes | u16 LE, Noise payload length        |
/// | 3      | handshake_payload| variable| Noise XX msg3 (73 bytes typical)    |
#[derive(Clone, Debug)]
pub struct SessionMsg3 {
    /// Reserved flags byte.
    pub flags: u8,
    /// Noise XX handshake message 3.
    pub handshake_payload: Vec<u8>,
}

impl SessionMsg3 {
    /// Create a new SessionMsg3 with the given handshake payload.
    pub fn new(handshake_payload: Vec<u8>) -> Self {
        Self {
            flags: 0,
            handshake_payload,
        }
    }

    /// Encode as wire format (4-byte FSP prefix + flags + handshake).
    ///
    /// The 4-byte prefix: `[ver_phase:1][flags:1][payload_len:2 LE]`
    /// where ver_phase = 0x03 (version 0, phase MSG3).
    pub fn encode(&self) -> Vec<u8> {
        // Build body first to compute payload_len
        let mut body = Vec::new();
        body.push(self.flags);
        let hs_len = self.handshake_payload.len() as u16;
        body.extend_from_slice(&hs_len.to_le_bytes());
        body.extend_from_slice(&self.handshake_payload);

        // Prepend 4-byte FSP common prefix
        let payload_len = body.len() as u16;
        let mut buf = Vec::with_capacity(4 + body.len());
        buf.push(0x03); // version 0, phase 0x3 (MSG3)
        buf.push(0x00); // flags (must be zero for handshake)
        buf.extend_from_slice(&payload_len.to_le_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    /// Decode from wire format (after 4-byte FSP prefix has been consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
        if payload.is_empty() {
            return Err(ProtocolError::MessageTooShort {
                expected: 1,
                got: 0,
            });
        }
        let flags = payload[0];
        let mut offset = 1;

        if payload.len() < offset + 2 {
            return Err(ProtocolError::MessageTooShort {
                expected: offset + 2,
                got: payload.len(),
            });
        }
        let hs_len = u16::from_le_bytes([payload[offset], payload[offset + 1]]) as usize;
        offset += 2;

        if payload.len() < offset + hs_len {
            return Err(ProtocolError::MessageTooShort {
                expected: offset + hs_len,
                got: payload.len(),
            });
        }
        let handshake_payload = payload[offset..offset + hs_len].to_vec();

        Ok(Self {
            flags,
            handshake_payload,
        })
    }
}

// ============================================================================
// Session-Layer MMP Reports
// ============================================================================

/// Session-layer sender report (msg_type 0x11).
///
/// Mirrors the FMP `SenderReport` fields but carried as an FSP session
/// message inside the AEAD envelope. Uses the same extensibility header
/// as the link-layer format: `[format_version:1][total_length:2 LE]`.
///
/// ## Wire Format (19 bytes body, after inner header stripped)
///
/// ```text
/// [0]     format_version = 0
/// [1-2]   total_length: u16 LE (= 16)
/// [3-6]   interval_packets_sent: u32 LE
/// [7-10]  interval_bytes_sent: u32 LE
/// [11-18] cumulative_packets_sent: u64 LE
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionSenderReport {
    pub interval_packets_sent: u32,
    pub interval_bytes_sent: u32,
    pub cumulative_packets_sent: u64,
}

/// Body size for SessionSenderReport: format_version(1) + total_length(2) + payload(16).
pub const SESSION_SENDER_REPORT_SIZE: usize = 19;

/// Payload size after total_length field for SessionSenderReport format v0.
const SESSION_SR_PAYLOAD: u16 = 16;

impl SessionSenderReport {
    /// Encode to wire format (19 bytes body).
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(SESSION_SENDER_REPORT_SIZE);
        buf.push(0x00); // format_version
        buf.extend_from_slice(&SESSION_SR_PAYLOAD.to_le_bytes());
        buf.extend_from_slice(&self.interval_packets_sent.to_le_bytes());
        buf.extend_from_slice(&self.interval_bytes_sent.to_le_bytes());
        buf.extend_from_slice(&self.cumulative_packets_sent.to_le_bytes());
        buf
    }

    /// Decode from body (after FSP inner header has been stripped).
    pub fn decode(body: &[u8]) -> Result<Self, ProtocolError> {
        if body.len() < SESSION_SENDER_REPORT_SIZE {
            return Err(ProtocolError::MessageTooShort {
                expected: SESSION_SENDER_REPORT_SIZE,
                got: body.len(),
            });
        }
        let _format_version = body[0];
        let total_length = u16::from_le_bytes(body[1..3].try_into().unwrap()) as usize;
        if body.len() < 3 + total_length {
            return Err(ProtocolError::MessageTooShort {
                expected: 3 + total_length,
                got: body.len(),
            });
        }
        let p = &body[3..];
        Ok(Self {
            interval_packets_sent: u32::from_le_bytes(p[0..4].try_into().unwrap()),
            interval_bytes_sent: u32::from_le_bytes(p[4..8].try_into().unwrap()),
            cumulative_packets_sent: u64::from_le_bytes(p[8..16].try_into().unwrap()),
        })
    }
}

/// Session-layer receiver report (msg_type 0x12).
///
/// Mirrors the FMP `ReceiverReport` fields but carried as an FSP session
/// message inside the AEAD envelope. Uses the same extensibility header
/// as the link-layer format: `[format_version:1][total_length:2 LE]`.
///
/// ## Wire Format (53 bytes body, after inner header stripped)
///
/// ```text
/// [0]     format_version = 0
/// [1-2]   total_length: u16 LE (= 50)
/// [3-6]   timestamp_echo: u32 LE
/// [7-8]   dwell_time: u16 LE
/// [9-16]  highest_counter: u64 LE
/// [17-24] cumulative_packets_recv: u64 LE
/// [25-32] cumulative_bytes_recv: u64 LE
/// [33-36] jitter: u32 LE (microseconds)
/// [37-40] ecn_ce_count: u32 LE
/// [41-44] owd_trend: i32 LE (µs/s)
/// [45-48] burst_loss_count: u32 LE
/// [49-52] cumulative_reorder_count: u32 LE
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionReceiverReport {
    pub timestamp_echo: u32,
    pub dwell_time: u16,
    pub highest_counter: u64,
    pub cumulative_packets_recv: u64,
    pub cumulative_bytes_recv: u64,
    pub jitter: u32,
    pub ecn_ce_count: u32,
    pub owd_trend: i32,
    pub burst_loss_count: u32,
    pub cumulative_reorder_count: u32,
}

/// Body size for SessionReceiverReport: format_version(1) + total_length(2) + payload(50).
pub const SESSION_RECEIVER_REPORT_SIZE: usize = 53;

/// Payload size after total_length field for SessionReceiverReport format v0.
const SESSION_RR_PAYLOAD: u16 = 50;

impl SessionReceiverReport {
    /// Encode to wire format (53 bytes body).
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(SESSION_RECEIVER_REPORT_SIZE);
        buf.push(0x00); // format_version
        buf.extend_from_slice(&SESSION_RR_PAYLOAD.to_le_bytes());
        buf.extend_from_slice(&self.timestamp_echo.to_le_bytes());
        buf.extend_from_slice(&self.dwell_time.to_le_bytes());
        buf.extend_from_slice(&self.highest_counter.to_le_bytes());
        buf.extend_from_slice(&self.cumulative_packets_recv.to_le_bytes());
        buf.extend_from_slice(&self.cumulative_bytes_recv.to_le_bytes());
        buf.extend_from_slice(&self.jitter.to_le_bytes());
        buf.extend_from_slice(&self.ecn_ce_count.to_le_bytes());
        buf.extend_from_slice(&self.owd_trend.to_le_bytes());
        buf.extend_from_slice(&self.burst_loss_count.to_le_bytes());
        buf.extend_from_slice(&self.cumulative_reorder_count.to_le_bytes());
        buf
    }

    /// Decode from body (after FSP inner header has been stripped).
    pub fn decode(body: &[u8]) -> Result<Self, ProtocolError> {
        if body.len() < SESSION_RECEIVER_REPORT_SIZE {
            return Err(ProtocolError::MessageTooShort {
                expected: SESSION_RECEIVER_REPORT_SIZE,
                got: body.len(),
            });
        }
        let _format_version = body[0];
        let total_length = u16::from_le_bytes(body[1..3].try_into().unwrap()) as usize;
        if body.len() < 3 + total_length {
            return Err(ProtocolError::MessageTooShort {
                expected: 3 + total_length,
                got: body.len(),
            });
        }
        let p = &body[3..];
        Ok(Self {
            timestamp_echo: u32::from_le_bytes(p[0..4].try_into().unwrap()),
            dwell_time: u16::from_le_bytes(p[4..6].try_into().unwrap()),
            highest_counter: u64::from_le_bytes(p[6..14].try_into().unwrap()),
            cumulative_packets_recv: u64::from_le_bytes(p[14..22].try_into().unwrap()),
            cumulative_bytes_recv: u64::from_le_bytes(p[22..30].try_into().unwrap()),
            jitter: u32::from_le_bytes(p[30..34].try_into().unwrap()),
            ecn_ce_count: u32::from_le_bytes(p[34..38].try_into().unwrap()),
            owd_trend: i32::from_le_bytes(p[38..42].try_into().unwrap()),
            burst_loss_count: u32::from_le_bytes(p[42..46].try_into().unwrap()),
            cumulative_reorder_count: u32::from_le_bytes(p[46..50].try_into().unwrap()),
        })
    }
}

/// Path MTU notification (msg_type 0x13).
///
/// Sent by a node that discovers a path MTU value (from transit router
/// feedback or ICMP Packet Too Big). Allows the remote endpoint to
/// adjust its sending MTU.
///
/// ## Wire Format (2 bytes body, after inner header stripped)
///
/// ```text
/// [0-1]   path_mtu: u16 LE
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PathMtuNotification {
    /// Discovered path MTU in bytes.
    pub path_mtu: u16,
}

/// Body size for PathMtuNotification.
pub const PATH_MTU_NOTIFICATION_SIZE: usize = 2;

impl PathMtuNotification {
    /// Create a new path MTU notification.
    pub fn new(path_mtu: u16) -> Self {
        Self { path_mtu }
    }

    /// Encode to wire format (2 bytes body).
    pub fn encode(&self) -> Vec<u8> {
        self.path_mtu.to_le_bytes().to_vec()
    }

    /// Decode from body (after FSP inner header has been stripped).
    pub fn decode(body: &[u8]) -> Result<Self, ProtocolError> {
        if body.len() < PATH_MTU_NOTIFICATION_SIZE {
            return Err(ProtocolError::MessageTooShort {
                expected: PATH_MTU_NOTIFICATION_SIZE,
                got: body.len(),
            });
        }
        Ok(Self {
            path_mtu: u16::from_le_bytes([body[0], body[1]]),
        })
    }
}

// ============================================================================
// Error Messages
// ============================================================================

/// Link-layer error signal indicating router cache miss.
///
/// Generated by a transit router when it cannot forward a SessionDatagram
/// due to missing cached coordinates for the destination. Carried inside
/// a new SessionDatagram addressed back to the original source
/// (src_addr=reporter, dest_addr=original_source). Plaintext — not
/// end-to-end encrypted, since the transit router has no session with
/// the source.
///
/// ## Wire Format
///
/// | Offset | Field    | Size     | Description                        |
/// |--------|----------|---------|------------------------------------|
/// | 0      | msg_type | 1 byte  | 0x20                               |
/// | 1      | flags    | 1 byte  | Reserved                           |
/// | 2      | dest_addr| 16 bytes| The node_addr we couldn't route to |
/// | 18     | reporter | 16 bytes| NodeAddr of reporting router       |
///
/// Payload: 34 bytes
#[derive(Clone, Debug)]
pub struct CoordsRequired {
    /// Destination that couldn't be routed.
    pub dest_addr: NodeAddr,
    /// Router reporting the miss.
    pub reporter: NodeAddr,
}

/// Wire size of CoordsRequired payload: msg_type(1) + flags(1) + dest_addr(16) + reporter(16).
pub const COORDS_REQUIRED_SIZE: usize = 34;

impl CoordsRequired {
    /// Create a new CoordsRequired error.
    pub fn new(dest_addr: NodeAddr, reporter: NodeAddr) -> Self {
        Self {
            dest_addr,
            reporter,
        }
    }

    /// Encode as wire format (4-byte FSP prefix + msg_type + body).
    ///
    /// Error signals use phase=0x0 with U flag set.
    pub fn encode(&self) -> Vec<u8> {
        // Body: msg_type + flags(reserved) + dest_addr + reporter
        let body_len = 1 + 1 + 16 + 16; // 34 bytes
        let mut buf = Vec::with_capacity(4 + body_len);
        // FSP prefix: version 0, phase 0x0, U flag set
        buf.push(0x00); // version 0, phase 0x0
        buf.push(0x04); // U flag
        let payload_len = body_len as u16;
        buf.extend_from_slice(&payload_len.to_le_bytes());
        // msg_type byte (after prefix, before body)
        buf.push(SessionMessageType::CoordsRequired.to_byte());
        buf.push(0x00); // reserved flags
        buf.extend_from_slice(self.dest_addr.as_bytes());
        buf.extend_from_slice(self.reporter.as_bytes());
        buf
    }

    /// Decode from wire format (after FSP prefix and msg_type byte consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
        // flags(1) + dest_addr(16) + reporter(16) = 33
        if payload.len() < 33 {
            return Err(ProtocolError::MessageTooShort {
                expected: 33,
                got: payload.len(),
            });
        }
        // payload[0] is flags (reserved, ignored)
        let mut dest_bytes = [0u8; 16];
        dest_bytes.copy_from_slice(&payload[1..17]);
        let mut reporter_bytes = [0u8; 16];
        reporter_bytes.copy_from_slice(&payload[17..33]);

        Ok(Self {
            dest_addr: NodeAddr::from_bytes(dest_bytes),
            reporter: NodeAddr::from_bytes(reporter_bytes),
        })
    }
}

/// Error indicating routing failure (local minimum or unreachable).
///
/// Carried inside a SessionDatagram addressed back to the original source.
/// The reporting router creates a new SessionDatagram with src_addr=reporter
/// and dest_addr=original_source, so the `original_src` field from the old
/// design is no longer needed — it's the SessionDatagram's dest_addr.
///
/// ## Wire Format
///
/// | Offset | Field             | Size     | Description                   |
/// |--------|-------------------|----------|-------------------------------|
/// | 0      | msg_type          | 1 byte   | 0x21                          |
/// | 1      | flags             | 1 byte   | Reserved                      |
/// | 2      | dest_addr         | 16 bytes | The unreachable node_addr     |
/// | 18     | reporter          | 16 bytes | NodeAddr of reporting router   |
/// | 34     | last_coords_count | 2 bytes  | u16 LE                        |
/// | 36     | last_known_coords | 16 × n   | Stale coords that failed      |
#[derive(Clone, Debug)]
pub struct PathBroken {
    /// Destination that couldn't be reached.
    pub dest_addr: NodeAddr,
    /// Node that detected the failure.
    pub reporter: NodeAddr,
    /// Optional: last known coordinates of destination.
    pub last_known_coords: Option<TreeCoordinate>,
}

impl PathBroken {
    /// Create a new PathBroken error.
    pub fn new(dest_addr: NodeAddr, reporter: NodeAddr) -> Self {
        Self {
            dest_addr,
            reporter,
            last_known_coords: None,
        }
    }

    /// Add last known coordinates.
    pub fn with_last_coords(mut self, coords: TreeCoordinate) -> Self {
        self.last_known_coords = Some(coords);
        self
    }

    /// Encode as wire format (4-byte FSP prefix + msg_type + body).
    ///
    /// Error signals use phase=0x0 with U flag set.
    pub fn encode(&self) -> Vec<u8> {
        // Build body first to compute length
        let mut body = Vec::new();
        body.push(SessionMessageType::PathBroken.to_byte());
        body.push(0x00); // reserved flags
        body.extend_from_slice(self.dest_addr.as_bytes());
        body.extend_from_slice(self.reporter.as_bytes());
        if let Some(ref coords) = self.last_known_coords {
            encode_coords(coords, &mut body);
        } else {
            encode_empty_coords(&mut body);
        }

        // Prepend FSP prefix: version 0, phase 0x0, U flag set
        let payload_len = body.len() as u16;
        let mut buf = Vec::with_capacity(4 + body.len());
        buf.push(0x00); // version 0, phase 0x0
        buf.push(0x04); // U flag
        buf.extend_from_slice(&payload_len.to_le_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    /// Decode from wire format (after FSP prefix and msg_type byte consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
        // flags(1) + dest_addr(16) + reporter(16) + coords_count(2) = 35 minimum
        if payload.len() < 35 {
            return Err(ProtocolError::MessageTooShort {
                expected: 35,
                got: payload.len(),
            });
        }
        // payload[0] is flags (reserved, ignored)
        let mut dest_bytes = [0u8; 16];
        dest_bytes.copy_from_slice(&payload[1..17]);
        let mut reporter_bytes = [0u8; 16];
        reporter_bytes.copy_from_slice(&payload[17..33]);

        let (last_known_coords, _consumed) = decode_optional_coords(&payload[33..])?;

        Ok(Self {
            dest_addr: NodeAddr::from_bytes(dest_bytes),
            reporter: NodeAddr::from_bytes(reporter_bytes),
            last_known_coords,
        })
    }
}

/// Error indicating a forwarded packet exceeded the next-hop transport MTU.
///
/// Generated by a transit router when `send_encrypted_link_message()`
/// fails with `TransportError::MtuExceeded`. The reporter includes the
/// bottleneck MTU so the source can immediately reduce its sending MTU.
///
/// ## Wire Format
///
/// | Offset | Field     | Size     | Description                        |
/// |--------|-----------|----------|------------------------------------|
/// | 0      | msg_type  | 1 byte   | 0x22                               |
/// | 1      | flags     | 1 byte   | Reserved                           |
/// | 2      | dest_addr | 16 bytes | The destination we were forwarding |
/// | 18     | reporter  | 16 bytes | NodeAddr of reporting router       |
/// | 34     | mtu       | 2 bytes  | Bottleneck MTU (u16 LE)            |
///
/// Payload: 36 bytes
#[derive(Clone, Debug)]
pub struct MtuExceeded {
    /// Destination that the oversized packet was heading to.
    pub dest_addr: NodeAddr,
    /// Router that detected the MTU violation.
    pub reporter: NodeAddr,
    /// Transport MTU at the bottleneck hop.
    pub mtu: u16,
}

/// Wire size of MtuExceeded payload: msg_type(1) + flags(1) + dest_addr(16) + reporter(16) + mtu(2).
pub const MTU_EXCEEDED_SIZE: usize = 36;

impl MtuExceeded {
    /// Create a new MtuExceeded error.
    pub fn new(dest_addr: NodeAddr, reporter: NodeAddr, mtu: u16) -> Self {
        Self {
            dest_addr,
            reporter,
            mtu,
        }
    }

    /// Encode as wire format (4-byte FSP prefix + msg_type + body).
    ///
    /// Error signals use phase=0x0 with U flag set.
    pub fn encode(&self) -> Vec<u8> {
        let body_len = MTU_EXCEEDED_SIZE; // 36 bytes
        let mut buf = Vec::with_capacity(4 + body_len);
        // FSP prefix: version 0, phase 0x0, U flag set
        buf.push(0x00); // version 0, phase 0x0
        buf.push(0x04); // U flag
        let payload_len = body_len as u16;
        buf.extend_from_slice(&payload_len.to_le_bytes());
        // msg_type byte
        buf.push(SessionMessageType::MtuExceeded.to_byte());
        buf.push(0x00); // reserved flags
        buf.extend_from_slice(self.dest_addr.as_bytes());
        buf.extend_from_slice(self.reporter.as_bytes());
        buf.extend_from_slice(&self.mtu.to_le_bytes());
        buf
    }

    /// Decode from wire format (after FSP prefix and msg_type byte consumed).
    pub fn decode(payload: &[u8]) -> Result<Self, ProtocolError> {
        // flags(1) + dest_addr(16) + reporter(16) + mtu(2) = 35
        if payload.len() < 35 {
            return Err(ProtocolError::MessageTooShort {
                expected: 35,
                got: payload.len(),
            });
        }
        // payload[0] is flags (reserved, ignored)
        let mut dest_bytes = [0u8; 16];
        dest_bytes.copy_from_slice(&payload[1..17]);
        let mut reporter_bytes = [0u8; 16];
        reporter_bytes.copy_from_slice(&payload[17..33]);
        let mtu = u16::from_le_bytes([payload[33], payload[34]]);

        Ok(Self {
            dest_addr: NodeAddr::from_bytes(dest_bytes),
            reporter: NodeAddr::from_bytes(reporter_bytes),
            mtu,
        })
    }
}

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

    fn make_node_addr(val: u8) -> NodeAddr {
        let mut bytes = [0u8; 16];
        bytes[0] = val;
        NodeAddr::from_bytes(bytes)
    }

    fn make_coords(ids: &[u8]) -> TreeCoordinate {
        TreeCoordinate::from_addrs(ids.iter().map(|&v| make_node_addr(v)).collect()).unwrap()
    }

    // ===== SessionMessageType Tests =====

    #[test]
    fn test_session_message_type_roundtrip() {
        let types = [
            SessionMessageType::SessionSetup,
            SessionMessageType::SessionAck,
            SessionMessageType::DataPacket,
            SessionMessageType::SenderReport,
            SessionMessageType::ReceiverReport,
            SessionMessageType::PathMtuNotification,
            SessionMessageType::CoordsRequired,
            SessionMessageType::PathBroken,
            SessionMessageType::MtuExceeded,
        ];

        for ty in types {
            let byte = ty.to_byte();
            let restored = SessionMessageType::from_byte(byte);
            assert_eq!(restored, Some(ty));
        }
    }

    #[test]
    fn test_session_message_type_invalid() {
        assert!(SessionMessageType::from_byte(0xFF).is_none());
        assert!(SessionMessageType::from_byte(0x99).is_none());
    }

    // ===== SessionFlags Tests =====

    #[test]
    fn test_session_flags() {
        let flags = SessionFlags::new().with_ack().bidirectional();

        assert!(flags.request_ack);
        assert!(flags.bidirectional);

        let byte = flags.to_byte();
        let restored = SessionFlags::from_byte(byte);

        assert_eq!(flags, restored);
    }

    #[test]
    fn test_session_flags_default() {
        let flags = SessionFlags::new();
        assert!(!flags.request_ack);
        assert!(!flags.bidirectional);
        assert_eq!(flags.to_byte(), 0);
    }

    // ===== SessionSetup Tests =====

    #[test]
    fn test_session_setup() {
        let setup = SessionSetup::new(make_coords(&[1, 0]), make_coords(&[2, 0]))
            .with_flags(SessionFlags::new().with_ack());

        assert!(setup.flags.request_ack);
        assert!(!setup.flags.bidirectional);
    }

    // ===== CoordsRequired Tests =====

    #[test]
    fn test_coords_required() {
        let err = CoordsRequired::new(make_node_addr(1), make_node_addr(2));

        assert_eq!(err.dest_addr, make_node_addr(1));
        assert_eq!(err.reporter, make_node_addr(2));
    }

    // ===== PathBroken Tests =====

    #[test]
    fn test_path_broken() {
        let err = PathBroken::new(make_node_addr(2), make_node_addr(3))
            .with_last_coords(make_coords(&[2, 0]));

        assert_eq!(err.dest_addr, make_node_addr(2));
        assert_eq!(err.reporter, make_node_addr(3));
        assert!(err.last_known_coords.is_some());
    }

    // ===== Encode/Decode Roundtrip Tests =====

    #[test]
    fn test_session_setup_encode_decode() {
        let handshake = vec![0xAA; 82]; // typical Noise XX msg1
        let setup = SessionSetup::new(make_coords(&[1, 2, 0]), make_coords(&[3, 4, 0]))
            .with_flags(SessionFlags::new().with_ack().bidirectional())
            .with_handshake(handshake.clone());

        let encoded = setup.encode();

        // Verify FSP prefix: ver_phase=0x01 (version 0, phase MSG1)
        assert_eq!(encoded[0], 0x01);
        assert_eq!(encoded[1], 0x00); // flags = 0 for handshake
        let payload_len = u16::from_le_bytes([encoded[2], encoded[3]]);
        assert_eq!(payload_len as usize, encoded.len() - 4);

        // Decode (skip 4-byte FSP prefix)
        let decoded = SessionSetup::decode(&encoded[4..]).unwrap();

        assert_eq!(decoded.flags, setup.flags);
        assert_eq!(decoded.src_coords, setup.src_coords);
        assert_eq!(decoded.dest_coords, setup.dest_coords);
        assert_eq!(decoded.handshake_payload, handshake);
    }

    #[test]
    fn test_session_setup_no_handshake() {
        let setup = SessionSetup::new(make_coords(&[5, 0]), make_coords(&[6, 0]));

        let encoded = setup.encode();
        let decoded = SessionSetup::decode(&encoded[4..]).unwrap();

        assert!(decoded.handshake_payload.is_empty());
        assert_eq!(decoded.src_coords, setup.src_coords);
        assert_eq!(decoded.dest_coords, setup.dest_coords);
    }

    #[test]
    fn test_session_ack_encode_decode() {
        let handshake = vec![0xBB; 33]; // typical Noise XX msg2
        let ack = SessionAck::new(make_coords(&[7, 8, 0]), make_coords(&[3, 4, 0]))
            .with_handshake(handshake.clone());

        let encoded = ack.encode();
        // Verify FSP prefix: ver_phase=0x02 (version 0, phase MSG2)
        assert_eq!(encoded[0], 0x02);
        assert_eq!(encoded[1], 0x00); // flags = 0 for handshake

        let decoded = SessionAck::decode(&encoded[4..]).unwrap();
        assert_eq!(decoded.src_coords, ack.src_coords);
        assert_eq!(decoded.dest_coords, ack.dest_coords);
        assert_eq!(decoded.handshake_payload, handshake);
    }

    #[test]
    fn test_coords_required_encode_decode() {
        let err = CoordsRequired::new(make_node_addr(0xAA), make_node_addr(0xBB));

        let encoded = err.encode();
        // 4 prefix + 1 msg_type + 1 flags + 16 dest + 16 reporter = 38
        assert_eq!(encoded.len(), 4 + COORDS_REQUIRED_SIZE);
        // Check FSP prefix: phase 0x0, U flag
        assert_eq!(encoded[0], 0x00);
        assert_eq!(encoded[1], 0x04); // U flag
        // msg_type after prefix
        assert_eq!(encoded[4], 0x20);

        // decode after prefix + msg_type consumed
        let decoded = CoordsRequired::decode(&encoded[5..]).unwrap();
        assert_eq!(decoded.dest_addr, err.dest_addr);
        assert_eq!(decoded.reporter, err.reporter);
    }

    #[test]
    fn test_path_broken_encode_decode_no_coords() {
        let err = PathBroken::new(make_node_addr(0xCC), make_node_addr(0xDD));

        let encoded = err.encode();
        // Check FSP prefix
        assert_eq!(encoded[0], 0x00);
        assert_eq!(encoded[1], 0x04); // U flag
        assert_eq!(encoded[4], 0x21); // msg_type

        let decoded = PathBroken::decode(&encoded[5..]).unwrap();
        assert_eq!(decoded.dest_addr, err.dest_addr);
        assert_eq!(decoded.reporter, err.reporter);
        assert!(decoded.last_known_coords.is_none());
    }

    #[test]
    fn test_path_broken_encode_decode_with_coords() {
        let coords = make_coords(&[0xCC, 0xDD, 0xEE]);
        let err = PathBroken::new(make_node_addr(0x11), make_node_addr(0x22))
            .with_last_coords(coords.clone());

        let encoded = err.encode();
        let decoded = PathBroken::decode(&encoded[5..]).unwrap();

        assert_eq!(decoded.dest_addr, err.dest_addr);
        assert_eq!(decoded.reporter, err.reporter);
        assert_eq!(decoded.last_known_coords.unwrap(), coords);
    }

    #[test]
    fn test_session_setup_decode_too_short() {
        assert!(SessionSetup::decode(&[]).is_err());
    }

    #[test]
    fn test_session_ack_decode_too_short() {
        assert!(SessionAck::decode(&[]).is_err());
    }

    #[test]
    fn test_coords_required_decode_too_short() {
        assert!(CoordsRequired::decode(&[]).is_err());
        assert!(CoordsRequired::decode(&[0x00; 10]).is_err());
    }

    #[test]
    fn test_path_broken_decode_too_short() {
        assert!(PathBroken::decode(&[]).is_err());
        assert!(PathBroken::decode(&[0x00; 20]).is_err());
    }

    #[test]
    fn test_session_setup_deep_coords() {
        // Depth-10 coordinate (11 entries: self + 10 ancestors)
        let addrs: Vec<u8> = (0..11).collect();
        let src = make_coords(&addrs);
        let dest = make_coords(&[20, 21, 22, 23, 24]);
        let setup = SessionSetup::new(src.clone(), dest.clone()).with_handshake(vec![0x55; 82]);

        let encoded = setup.encode();
        let decoded = SessionSetup::decode(&encoded[4..]).unwrap();

        assert_eq!(decoded.src_coords, src);
        assert_eq!(decoded.dest_coords, dest);
    }

    // ===== FspFlags Tests =====

    #[test]
    fn test_fsp_flags_default() {
        let flags = FspFlags::new();
        assert!(!flags.coords_present);
        assert!(!flags.key_epoch);
        assert!(!flags.unencrypted);
        assert_eq!(flags.to_byte(), 0x00);
    }

    #[test]
    fn test_fsp_flags_roundtrip() {
        // All combinations of 3 bits
        for byte in 0u8..=0x07 {
            let flags = FspFlags::from_byte(byte);
            assert_eq!(flags.to_byte(), byte);
        }
    }

    #[test]
    fn test_fsp_flags_individual_bits() {
        let cp = FspFlags::from_byte(0x01);
        assert!(cp.coords_present);
        assert!(!cp.key_epoch);
        assert!(!cp.unencrypted);

        let k = FspFlags::from_byte(0x02);
        assert!(!k.coords_present);
        assert!(k.key_epoch);
        assert!(!k.unencrypted);

        let u = FspFlags::from_byte(0x04);
        assert!(!u.coords_present);
        assert!(!u.key_epoch);
        assert!(u.unencrypted);
    }

    #[test]
    fn test_fsp_flags_ignores_reserved_bits() {
        // Reserved bits in upper 5 bits are not preserved
        let flags = FspFlags::from_byte(0xFF);
        assert!(flags.coords_present);
        assert!(flags.key_epoch);
        assert!(flags.unencrypted);
        assert_eq!(flags.to_byte(), 0x07); // only lower 3 bits
    }

    // ===== FspInnerFlags Tests =====

    #[test]
    fn test_fsp_inner_flags_default() {
        let flags = FspInnerFlags::new();
        assert_eq!(flags.to_byte(), 0x00);
    }

    #[test]
    fn test_fsp_inner_flags_from_byte() {
        // All bits are reserved; from_byte always returns default
        let flags = FspInnerFlags::from_byte(0xFF);
        assert_eq!(flags.to_byte(), 0x00);
    }

    // ===== New SessionMessageType Values =====

    #[test]
    fn test_session_message_type_new_values() {
        assert_eq!(SessionMessageType::SenderReport.to_byte(), 0x11);
        assert_eq!(SessionMessageType::ReceiverReport.to_byte(), 0x12);
        assert_eq!(SessionMessageType::PathMtuNotification.to_byte(), 0x13);
    }

    #[test]
    fn test_session_message_type_display() {
        assert_eq!(
            format!("{}", SessionMessageType::SenderReport),
            "SenderReport"
        );
        assert_eq!(
            format!("{}", SessionMessageType::ReceiverReport),
            "ReceiverReport"
        );
        assert_eq!(
            format!("{}", SessionMessageType::PathMtuNotification),
            "PathMtuNotification"
        );
    }

    // ===== SessionSenderReport Tests =====

    fn sample_session_sender_report() -> SessionSenderReport {
        SessionSenderReport {
            interval_packets_sent: 100,
            interval_bytes_sent: 50_000,
            cumulative_packets_sent: 10_000,
        }
    }

    #[test]
    fn test_session_sender_report_encode_size() {
        let sr = sample_session_sender_report();
        let encoded = sr.encode();
        assert_eq!(encoded.len(), SESSION_SENDER_REPORT_SIZE);
    }

    #[test]
    fn test_session_sender_report_roundtrip() {
        let sr = sample_session_sender_report();
        let encoded = sr.encode();
        let decoded = SessionSenderReport::decode(&encoded).unwrap();
        assert_eq!(sr, decoded);
    }

    #[test]
    fn test_session_sender_report_too_short() {
        assert!(SessionSenderReport::decode(&[0u8; 10]).is_err());
    }

    // ===== SessionReceiverReport Tests =====

    fn sample_session_receiver_report() -> SessionReceiverReport {
        SessionReceiverReport {
            timestamp_echo: 5900,
            dwell_time: 5,
            highest_counter: 195,
            cumulative_packets_recv: 9_500,
            cumulative_bytes_recv: 4_750_000,
            jitter: 1200,
            ecn_ce_count: 0,
            owd_trend: -50,
            burst_loss_count: 2,
            cumulative_reorder_count: 10,
        }
    }

    #[test]
    fn test_session_receiver_report_encode_size() {
        let rr = sample_session_receiver_report();
        let encoded = rr.encode();
        assert_eq!(encoded.len(), SESSION_RECEIVER_REPORT_SIZE);
    }

    #[test]
    fn test_session_receiver_report_roundtrip() {
        let rr = sample_session_receiver_report();
        let encoded = rr.encode();
        let decoded = SessionReceiverReport::decode(&encoded).unwrap();
        assert_eq!(rr, decoded);
    }

    #[test]
    fn test_session_receiver_report_too_short() {
        assert!(SessionReceiverReport::decode(&[0u8; 10]).is_err());
    }

    #[test]
    fn test_session_receiver_report_negative_owd_trend() {
        let rr = SessionReceiverReport {
            owd_trend: -12345,
            ..sample_session_receiver_report()
        };
        let encoded = rr.encode();
        let decoded = SessionReceiverReport::decode(&encoded).unwrap();
        assert_eq!(decoded.owd_trend, -12345);
    }

    // ===== PathMtuNotification Tests =====

    #[test]
    fn test_path_mtu_notification_encode_size() {
        let n = PathMtuNotification::new(1400);
        let encoded = n.encode();
        assert_eq!(encoded.len(), PATH_MTU_NOTIFICATION_SIZE);
    }

    #[test]
    fn test_path_mtu_notification_roundtrip() {
        let n = PathMtuNotification::new(1400);
        let encoded = n.encode();
        let decoded = PathMtuNotification::decode(&encoded).unwrap();
        assert_eq!(decoded.path_mtu, 1400);
    }

    #[test]
    fn test_path_mtu_notification_too_short() {
        assert!(PathMtuNotification::decode(&[]).is_err());
        assert!(PathMtuNotification::decode(&[0x00]).is_err());
    }

    #[test]
    fn test_path_mtu_notification_boundary_values() {
        for mtu in [0u16, 1280, 1500, u16::MAX] {
            let n = PathMtuNotification::new(mtu);
            let encoded = n.encode();
            let decoded = PathMtuNotification::decode(&encoded).unwrap();
            assert_eq!(decoded.path_mtu, mtu);
        }
    }

    // ===== MtuExceeded Tests =====

    #[test]
    fn test_mtu_exceeded_encode_size() {
        let err = MtuExceeded::new(make_node_addr(0xAA), make_node_addr(0xBB), 1400);
        let encoded = err.encode();
        // 4 prefix + 36 body = 40
        assert_eq!(encoded.len(), 4 + MTU_EXCEEDED_SIZE);
    }

    #[test]
    fn test_mtu_exceeded_encode_decode() {
        let err = MtuExceeded::new(make_node_addr(0xAA), make_node_addr(0xBB), 1400);

        let encoded = err.encode();
        // Check FSP prefix: phase 0x0, U flag
        assert_eq!(encoded[0], 0x00);
        assert_eq!(encoded[1], 0x04); // U flag
        // msg_type after prefix
        assert_eq!(encoded[4], 0x22);

        // decode after prefix + msg_type consumed
        let decoded = MtuExceeded::decode(&encoded[5..]).unwrap();
        assert_eq!(decoded.dest_addr, err.dest_addr);
        assert_eq!(decoded.reporter, err.reporter);
        assert_eq!(decoded.mtu, 1400);
    }

    #[test]
    fn test_mtu_exceeded_decode_too_short() {
        assert!(MtuExceeded::decode(&[]).is_err());
        assert!(MtuExceeded::decode(&[0x00; 20]).is_err());
        assert!(MtuExceeded::decode(&[0x00; 34]).is_err()); // exactly 1 byte short
    }

    #[test]
    fn test_mtu_exceeded_boundary_mtu_values() {
        for mtu in [0u16, 1280, 1500, u16::MAX] {
            let err = MtuExceeded::new(make_node_addr(1), make_node_addr(2), mtu);
            let encoded = err.encode();
            let decoded = MtuExceeded::decode(&encoded[5..]).unwrap();
            assert_eq!(decoded.mtu, mtu);
        }
    }

    #[test]
    fn test_mtu_exceeded_message_type_value() {
        assert_eq!(SessionMessageType::MtuExceeded.to_byte(), 0x22);
        assert_eq!(
            SessionMessageType::from_byte(0x22),
            Some(SessionMessageType::MtuExceeded)
        );
    }

    #[test]
    fn test_mtu_exceeded_display() {
        assert_eq!(
            format!("{}", SessionMessageType::MtuExceeded),
            "MtuExceeded"
        );
    }

    // ===== SessionMsg3 Tests =====

    #[test]
    fn test_session_msg3_encode_decode() {
        let handshake = vec![0xCC; 73]; // typical XX msg3
        let msg3 = SessionMsg3::new(handshake.clone());

        let encoded = msg3.encode();
        // Verify FSP prefix: ver_phase=0x03 (version 0, phase MSG3)
        assert_eq!(encoded[0], 0x03);
        assert_eq!(encoded[1], 0x00); // flags = 0 for handshake
        let payload_len = u16::from_le_bytes([encoded[2], encoded[3]]);
        assert_eq!(payload_len as usize, encoded.len() - 4);

        // Decode (skip 4-byte FSP prefix)
        let decoded = SessionMsg3::decode(&encoded[4..]).unwrap();
        assert_eq!(decoded.flags, 0);
        assert_eq!(decoded.handshake_payload, handshake);
    }

    #[test]
    fn test_session_msg3_decode_too_short() {
        assert!(SessionMsg3::decode(&[]).is_err());
        assert!(SessionMsg3::decode(&[0x00]).is_err()); // flags only, no hs_len
    }

    #[test]
    fn test_session_msg3_empty_handshake() {
        let msg3 = SessionMsg3::new(vec![]);
        let encoded = msg3.encode();
        let decoded = SessionMsg3::decode(&encoded[4..]).unwrap();
        assert!(decoded.handshake_payload.is_empty());
    }
}