//! 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.
///
/// Handshake messages (SessionSetup, SessionAck, SessionMsg3) are **not**
/// identified by a message-type byte; they are dispatched by the FSP phase
/// nibble in the common prefix (0x1, 0x2, 0x3 respectively). The 0x00-0x0F
/// range is therefore unallocated in this enum.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum SessionMessageType {
    // 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 {
            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::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).
pub(crate) 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   | SP   | Spin bit for RTT measurement    |
/// | 1-7 |      | Reserved                        |
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct FspInnerFlags {
    /// Spin bit for passive RTT measurement.
    pub spin_bit: bool,
}

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

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

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

// ============================================================================
// 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 XK handshake message for session establishment.
///
/// SessionSetup, SessionAck, and SessionMsg3 are identified by the **phase**
/// field in the FSP common prefix (0x1, 0x2, 0x3), not by a message-type byte.
/// The `msg_type` field in the encrypted inner header applies only to
/// established-phase (0x0) messages.
///
/// ## Wire Format
///
/// Encoded with FSP common prefix: `[ver_phase:1][flags:1][payload_len:2 LE][body]`,
/// where `ver_phase = 0x01` (version 0, phase MSG1) and `flags = 0` for handshake.
///
/// **Body** (after 4-byte FSP prefix):
///
/// | Offset | Field             | Size       | Description                                         |
/// |--------|-------------------|------------|-----------------------------------------------------|
/// | 0      | flags             | 1 byte     | Bit 0: REQUEST_ACK, Bit 1: BIDIRECTIONAL            |
/// | 1      | src_coords_count  | 2 bytes LE | Number of source coordinate entries                 |
/// | 3      | src_coords        | 16 × n     | Source's ancestry (NodeAddr, self → root)           |
/// | ...    | dest_coords_count | 2 bytes LE | Number of dest coordinate entries                   |
/// | ...    | dest_coords       | 16 × m     | Destination's ancestry                              |
/// | ...    | handshake_len     | 2 bytes LE | Noise payload length                                |
/// | ...    | handshake_payload | variable   | Noise XK msg1 (33 bytes — ephemeral key only)       |
#[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 IK 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
/// XK handshake response.
///
/// SessionSetup, SessionAck, and SessionMsg3 are identified by the **phase**
/// field in the FSP common prefix (0x1, 0x2, 0x3), not by a message-type byte.
///
/// ## Wire Format
///
/// Encoded with FSP common prefix: `[ver_phase:1][flags:1][payload_len:2 LE][body]`,
/// where `ver_phase = 0x02` (version 0, phase MSG2) and `flags = 0` for handshake.
///
/// **Body** (after 4-byte FSP prefix):
///
/// | Offset | Field             | Size       | Description                                                  |
/// |--------|-------------------|------------|--------------------------------------------------------------|
/// | 0      | flags             | 1 byte     | Reserved                                                     |
/// | 1      | src_coords_count  | 2 bytes LE | Number of acknowledger coordinate entries                    |
/// | 3      | src_coords        | 16 × n     | Acknowledger's ancestry (for cache warming)                  |
/// | ...    | dest_coords_count | 2 bytes LE | Number of initiator coordinate entries                       |
/// | ...    | dest_coords       | 16 × m     | Initiator's ancestry (for return-path cache warming)         |
/// | ...    | handshake_len     | 2 bytes LE | Noise payload length                                         |
/// | ...    | handshake_payload | variable   | Noise XK msg2 (57 bytes — ephemeral key + encrypted epoch)   |
#[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 IK 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 (XK Handshake Message 3)
// ============================================================================

/// XK 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 XK msg3 (73 bytes typical)    |
#[derive(Clone, Debug)]
pub struct SessionMsg3 {
    /// Reserved flags byte.
    pub flags: u8,
    /// Noise XK 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. The msg_type is in the FSP inner
/// header, so the body starts with reserved bytes.
///
/// ## Wire Format (46 bytes body, after inner header stripped)
///
/// ```text
/// [0-1]   reserved (zero)
/// [2-9]   interval_start_counter: u64 LE
/// [10-17] interval_end_counter: u64 LE
/// [18-21] interval_start_timestamp: u32 LE
/// [22-25] interval_end_timestamp: u32 LE
/// [26-29] interval_bytes_sent: u32 LE
/// [30-37] cumulative_packets_sent: u64 LE
/// [38-45] cumulative_bytes_sent: u64 LE
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionSenderReport {
    pub interval_start_counter: u64,
    pub interval_end_counter: u64,
    pub interval_start_timestamp: u32,
    pub interval_end_timestamp: u32,
    pub interval_bytes_sent: u32,
    pub cumulative_packets_sent: u64,
    pub cumulative_bytes_sent: u64,
}

/// Body size for SessionSenderReport: 2 reserved + 44 fields.
pub const SESSION_SENDER_REPORT_SIZE: usize = 46;

impl SessionSenderReport {
    /// Encode to wire format (46 bytes body).
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(SESSION_SENDER_REPORT_SIZE);
        buf.extend_from_slice(&[0u8; 2]); // reserved
        buf.extend_from_slice(&self.interval_start_counter.to_le_bytes());
        buf.extend_from_slice(&self.interval_end_counter.to_le_bytes());
        buf.extend_from_slice(&self.interval_start_timestamp.to_le_bytes());
        buf.extend_from_slice(&self.interval_end_timestamp.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.extend_from_slice(&self.cumulative_bytes_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(),
            });
        }
        // Skip 2 reserved bytes
        let p = &body[2..];
        Ok(Self {
            interval_start_counter: u64::from_le_bytes(p[0..8].try_into().unwrap()),
            interval_end_counter: u64::from_le_bytes(p[8..16].try_into().unwrap()),
            interval_start_timestamp: u32::from_le_bytes(p[16..20].try_into().unwrap()),
            interval_end_timestamp: u32::from_le_bytes(p[20..24].try_into().unwrap()),
            interval_bytes_sent: u32::from_le_bytes(p[24..28].try_into().unwrap()),
            cumulative_packets_sent: u64::from_le_bytes(p[28..36].try_into().unwrap()),
            cumulative_bytes_sent: u64::from_le_bytes(p[36..44].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.
///
/// ## Wire Format (66 bytes body, after inner header stripped)
///
/// ```text
/// [0-1]   reserved (zero)
/// [2-9]   highest_counter: u64 LE
/// [10-17] cumulative_packets_recv: u64 LE
/// [18-25] cumulative_bytes_recv: u64 LE
/// [26-29] timestamp_echo: u32 LE
/// [30-31] dwell_time: u16 LE
/// [32-33] max_burst_loss: u16 LE
/// [34-35] mean_burst_loss: u16 LE (u8.8 fixed-point)
/// [36-37] reserved: u16 LE
/// [38-41] jitter: u32 LE (microseconds)
/// [42-45] ecn_ce_count: u32 LE
/// [46-49] owd_trend: i32 LE (µs/s)
/// [50-53] burst_loss_count: u32 LE
/// [54-57] cumulative_reorder_count: u32 LE
/// [58-61] interval_packets_recv: u32 LE
/// [62-65] interval_bytes_recv: u32 LE
/// ```
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionReceiverReport {
    pub highest_counter: u64,
    pub cumulative_packets_recv: u64,
    pub cumulative_bytes_recv: u64,
    pub timestamp_echo: u32,
    pub dwell_time: u16,
    pub max_burst_loss: u16,
    pub mean_burst_loss: u16,
    pub jitter: u32,
    pub ecn_ce_count: u32,
    pub owd_trend: i32,
    pub burst_loss_count: u32,
    pub cumulative_reorder_count: u32,
    pub interval_packets_recv: u32,
    pub interval_bytes_recv: u32,
}

/// Body size for SessionReceiverReport: 2 reserved + 64 fields.
pub const SESSION_RECEIVER_REPORT_SIZE: usize = 66;

impl SessionReceiverReport {
    /// Encode to wire format (66 bytes body).
    pub fn encode(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(SESSION_RECEIVER_REPORT_SIZE);
        buf.extend_from_slice(&[0u8; 2]); // reserved
        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.timestamp_echo.to_le_bytes());
        buf.extend_from_slice(&self.dwell_time.to_le_bytes());
        buf.extend_from_slice(&self.max_burst_loss.to_le_bytes());
        buf.extend_from_slice(&self.mean_burst_loss.to_le_bytes());
        buf.extend_from_slice(&[0u8; 2]); // reserved
        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.extend_from_slice(&self.interval_packets_recv.to_le_bytes());
        buf.extend_from_slice(&self.interval_bytes_recv.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(),
            });
        }
        // Skip 2 reserved bytes
        let p = &body[2..];
        Ok(Self {
            highest_counter: u64::from_le_bytes(p[0..8].try_into().unwrap()),
            cumulative_packets_recv: u64::from_le_bytes(p[8..16].try_into().unwrap()),
            cumulative_bytes_recv: u64::from_le_bytes(p[16..24].try_into().unwrap()),
            timestamp_echo: u32::from_le_bytes(p[24..28].try_into().unwrap()),
            dwell_time: u16::from_le_bytes(p[28..30].try_into().unwrap()),
            max_burst_loss: u16::from_le_bytes(p[30..32].try_into().unwrap()),
            mean_burst_loss: u16::from_le_bytes(p[32..34].try_into().unwrap()),
            // skip 2 reserved bytes at p[34..36]
            jitter: u32::from_le_bytes(p[36..40].try_into().unwrap()),
            ecn_ce_count: u32::from_le_bytes(p[40..44].try_into().unwrap()),
            owd_trend: i32::from_le_bytes(p[44..48].try_into().unwrap()),
            burst_loss_count: u32::from_le_bytes(p[48..52].try_into().unwrap()),
            cumulative_reorder_count: u32::from_le_bytes(p[52..56].try_into().unwrap()),
            interval_packets_recv: u32::from_le_bytes(p[56..60].try_into().unwrap()),
            interval_bytes_recv: u32::from_le_bytes(p[60..64].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]]),
        })
    }
}

#[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::DataPacket,
            SessionMessageType::SenderReport,
            SessionMessageType::ReceiverReport,
            SessionMessageType::PathMtuNotification,
            SessionMessageType::CoordsWarmup,
            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);
    }

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

    #[test]
    fn test_session_setup_encode_decode() {
        let handshake = vec![0xAA; 82]; // typical Noise IK 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 IK 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_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_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!(!flags.spin_bit);
        assert_eq!(flags.to_byte(), 0x00);
    }

    #[test]
    fn test_fsp_inner_flags_roundtrip() {
        let flags = FspInnerFlags::from_byte(0x01);
        assert!(flags.spin_bit);
        assert_eq!(flags.to_byte(), 0x01);

        let flags = FspInnerFlags::from_byte(0x00);
        assert!(!flags.spin_bit);
        assert_eq!(flags.to_byte(), 0x00);
    }

    #[test]
    fn test_fsp_inner_flags_ignores_reserved() {
        let flags = FspInnerFlags::from_byte(0xFE);
        assert!(!flags.spin_bit);
        assert_eq!(flags.to_byte(), 0x00);

        let flags = FspInnerFlags::from_byte(0xFF);
        assert!(flags.spin_bit);
        assert_eq!(flags.to_byte(), 0x01);
    }

    // ===== 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_start_counter: 100,
            interval_end_counter: 200,
            interval_start_timestamp: 5000,
            interval_end_timestamp: 6000,
            interval_bytes_sent: 50_000,
            cumulative_packets_sent: 10_000,
            cumulative_bytes_sent: 5_000_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 {
            highest_counter: 195,
            cumulative_packets_recv: 9_500,
            cumulative_bytes_recv: 4_750_000,
            timestamp_echo: 5900,
            dwell_time: 5,
            max_burst_loss: 3,
            mean_burst_loss: 384,
            jitter: 1200,
            ecn_ce_count: 0,
            owd_trend: -50,
            burst_loss_count: 2,
            cumulative_reorder_count: 10,
            interval_packets_recv: 95,
            interval_bytes_recv: 47_500,
        }
    }

    #[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_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 XK 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());
    }
}