//! RFC 6455 §5 — the frame codec. Pure: bytes in, frames out, no io and
//! no actor machinery, so every edge lives in fast unit tests. The
//! connection actor (chunk 3) owns buffering and the socket.
//!
//! Server-side rules enforced here:
//! - client→server frames MUST be masked (§5.1); decode takes
//!   `require_masked` so a future client mode can reuse the codec.
//! - server→client frames are NEVER masked; [`Frame::encode`] doesn't
//!   offer masking. [`encode_masked`] exists for the client side of
//!   tests.
//! - RSV bits must be 0 (we negotiate no extensions), §5.2.
//! - Payload lengths must use the minimal encoding, §5.2.
//! - Control frames: FIN set, payload ≤ 125, §5.5.
//!
//! Errors map to close codes via [`FrameError::close_code`]: protocol
//! violations → 1002, oversize → 1009, bad UTF-8 in text → 1007.

/// §5.2 opcodes. Reserved values are rejected at decode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Opcode {
    Continuation,
    Text,
    Binary,
    Close,
    Ping,
    Pong,
}

impl Opcode {
    fn from_u4(n: u8) -> Option<Self> {
        Some(match n {
            0x0 => Opcode::Continuation,
            0x1 => Opcode::Text,
            0x2 => Opcode::Binary,
            0x8 => Opcode::Close,
            0x9 => Opcode::Ping,
            0xA => Opcode::Pong,
            _ => return None,
        })
    }

    fn to_u4(self) -> u8 {
        match self {
            Opcode::Continuation => 0x0,
            Opcode::Text => 0x1,
            Opcode::Binary => 0x2,
            Opcode::Close => 0x8,
            Opcode::Ping => 0x9,
            Opcode::Pong => 0xA,
        }
    }

    pub fn is_control(self) -> bool {
        matches!(self, Opcode::Close | Opcode::Ping | Opcode::Pong)
    }
}

/// One decoded frame; payload already unmasked.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
    pub fin: bool,
    pub opcode: Opcode,
    pub payload: Vec<u8>,
}

impl Frame {
    pub fn new(opcode: Opcode, payload: impl Into<Vec<u8>>) -> Self {
        Frame { fin: true, opcode, payload: payload.into() }
    }

    /// Serialise unmasked (server→client, §5.1: a server MUST NOT mask).
    pub fn encode(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(self.payload.len() + 10);
        encode_head(&mut out, self.fin, self.opcode, self.payload.len(), None);
        out.extend_from_slice(&self.payload);
        out
    }
}

/// Why decoding (or assembly) failed. Fatal for the connection: the ws
/// close handshake should carry [`FrameError::close_code`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameError {
    /// §5.x violation; the str is a debugging breadcrumb, not protocol.
    Protocol(&'static str),
    /// Frame or assembled message exceeds the configured cap → 1009.
    TooLarge,
    /// A complete text message that is not valid UTF-8 → 1007.
    BadUtf8,
}

impl FrameError {
    pub fn close_code(self) -> u16 {
        match self {
            FrameError::Protocol(_) => 1002,
            FrameError::TooLarge => 1009,
            FrameError::BadUtf8 => 1007,
        }
    }
}

/// Try to decode one frame from the front of `buf`.
///
/// - `Ok(Some((frame, consumed)))` — drop `consumed` bytes off the front.
/// - `Ok(None)` — incomplete; read more. (Header-derived sizes are still
///   bounds-checked first, so a hostile 8-byte length can't make the
///   caller buffer toward it: oversize fails *before* the payload
///   arrives.)
/// - `Err(_)` — fatal; start the close handshake with the mapped code.
pub fn decode(
    buf: &[u8],
    require_masked: bool,
    max_payload: usize,
) -> Result<Option<(Frame, usize)>, FrameError> {
    if buf.len() < 2 {
        return Ok(None);
    }
    let b0 = buf[0];
    let b1 = buf[1];

    let fin = b0 & 0x80 != 0;
    if b0 & 0x70 != 0 {
        return Err(FrameError::Protocol("RSV bits set without extension"));
    }
    let opcode = Opcode::from_u4(b0 & 0x0F)
        .ok_or(FrameError::Protocol("reserved opcode"))?;

    let masked = b1 & 0x80 != 0;
    if require_masked && !masked {
        return Err(FrameError::Protocol("unmasked client frame"));
    }

    if opcode.is_control() {
        if !fin {
            return Err(FrameError::Protocol("fragmented control frame"));
        }
        if b1 & 0x7F > 125 {
            return Err(FrameError::Protocol("control payload > 125"));
        }
    }

    // Payload length: 7-bit, or 126 + u16, or 127 + u64 — minimal form
    // required (§5.2).
    let (len, mut pos): (u64, usize) = match b1 & 0x7F {
        126 => {
            if buf.len() < 4 {
                return Ok(None);
            }
            let l = u16::from_be_bytes([buf[2], buf[3]]) as u64;
            if l < 126 {
                return Err(FrameError::Protocol("non-minimal 16-bit length"));
            }
            (l, 4)
        }
        127 => {
            if buf.len() < 10 {
                return Ok(None);
            }
            let l = u64::from_be_bytes(buf[2..10].try_into().unwrap());
            if l & (1 << 63) != 0 {
                return Err(FrameError::Protocol("64-bit length MSB set"));
            }
            if l < 65536 {
                return Err(FrameError::Protocol("non-minimal 64-bit length"));
            }
            (l, 10)
        }
        n => (n as u64, 2),
    };

    // Cap check BEFORE waiting for the payload (see decode docs).
    if len > max_payload as u64 {
        return Err(FrameError::TooLarge);
    }
    let len = len as usize;

    let mask_key: Option<[u8; 4]> = if masked {
        if buf.len() < pos + 4 {
            return Ok(None);
        }
        let k = [buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]];
        pos += 4;
        Some(k)
    } else {
        None
    };

    if buf.len() < pos + len {
        return Ok(None);
    }

    let mut payload = buf[pos..pos + len].to_vec();
    if let Some(key) = mask_key {
        for (i, b) in payload.iter_mut().enumerate() {
            *b ^= key[i & 3];
        }
    }

    Ok(Some((Frame { fin, opcode, payload }, pos + len)))
}

fn encode_head(
    out: &mut Vec<u8>,
    fin: bool,
    opcode: Opcode,
    len: usize,
    mask: Option<[u8; 4]>,
) {
    let mask_bit = if mask.is_some() { 0x80 } else { 0 };
    out.push(if fin { 0x80 } else { 0 } | opcode.to_u4());
    if len <= 125 {
        out.push(mask_bit | len as u8);
    } else if len <= u16::MAX as usize {
        out.push(mask_bit | 126);
        out.extend_from_slice(&(len as u16).to_be_bytes());
    } else {
        out.push(mask_bit | 127);
        out.extend_from_slice(&(len as u64).to_be_bytes());
    }
    if let Some(k) = mask {
        out.extend_from_slice(&k);
    }
}

/// Client-side serialisation (masked). Servers never call this in
/// production — it exists so tests (and an eventual client mode) can
/// produce conformant client frames.
pub fn encode_masked(frame: &Frame, key: [u8; 4]) -> Vec<u8> {
    let mut out = Vec::with_capacity(frame.payload.len() + 14);
    encode_head(&mut out, frame.fin, frame.opcode, frame.payload.len(), Some(key));
    out.extend(frame.payload.iter().enumerate().map(|(i, b)| b ^ key[i & 3]));
    out
}

// ---------------------------------------------------------------------------
// Close payload (§5.5.1, §7.4)
// ---------------------------------------------------------------------------

/// Parse a close frame payload: empty is fine (`None`); otherwise a
/// 2-byte code + optional UTF-8 reason. A 1-byte payload, an invalid
/// wire code, or a non-UTF-8 reason are protocol errors.
pub fn parse_close_payload(payload: &[u8]) -> Result<Option<(u16, String)>, FrameError> {
    match payload.len() {
        0 => Ok(None),
        1 => Err(FrameError::Protocol("1-byte close payload")),
        _ => {
            let code = u16::from_be_bytes([payload[0], payload[1]]);
            if !close_code_valid_on_wire(code) {
                return Err(FrameError::Protocol("invalid close code"));
            }
            let reason = std::str::from_utf8(&payload[2..])
                .map_err(|_| FrameError::Protocol("close reason not UTF-8"))?;
            Ok(Some((code, reason.to_string())))
        }
    }
}

/// §7.4.1/.2: codes an endpoint may put on the wire. 1004 is reserved;
/// 1005/1006/1015 are signalling-only (MUST NOT appear in a close
/// frame); 1000–1011 otherwise fine; 3000–4999 are app/registry space.
fn close_code_valid_on_wire(code: u16) -> bool {
    matches!(code, 1000..=1003 | 1007..=1011 | 3000..=4999)
}

/// Build a close frame payload for `code` (+ optional reason, truncated
/// to fit the 125-byte control cap on a UTF-8 boundary).
pub fn close_payload(code: u16, reason: &str) -> Vec<u8> {
    let mut p = Vec::with_capacity(2 + reason.len().min(123));
    p.extend_from_slice(&code.to_be_bytes());
    let mut cut = reason.len().min(123);
    while !reason.is_char_boundary(cut) {
        cut -= 1;
    }
    p.extend_from_slice(&reason.as_bytes()[..cut]);
    p
}

// ---------------------------------------------------------------------------
// Message assembly (§5.4 fragmentation)
// ---------------------------------------------------------------------------

/// A complete application message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Message {
    Text(String),
    Binary(Vec<u8>),
}

/// Reassembles data frames into [`Message`]s. Feed it ONLY data frames
/// (Text/Binary/Continuation) — control frames interleave at the caller
/// (§5.4 allows them mid-fragmentation) and never enter the assembler.
///
/// Enforced here: continuation with nothing in progress; a new data
/// frame while a fragmented message is in progress; total message size
/// (1009); text UTF-8 validity, checked once on the complete message
/// (we deliver whole messages, so per-frame incremental validation buys
/// nothing).
#[derive(Debug)]
pub struct Assembler {
    buf: Vec<u8>,
    in_progress: Option<Opcode>, // Text or Binary
    max_message: usize,
}

impl Assembler {
    pub fn new(max_message: usize) -> Self {
        Assembler { buf: Vec::new(), in_progress: None, max_message }
    }

    pub fn push(&mut self, frame: Frame) -> Result<Option<Message>, FrameError> {
        match (frame.opcode, self.in_progress) {
            (Opcode::Text | Opcode::Binary, Some(_)) => {
                return Err(FrameError::Protocol("new data frame mid-fragmentation"));
            }
            (Opcode::Continuation, None) => {
                return Err(FrameError::Protocol("continuation without a message"));
            }
            (Opcode::Text | Opcode::Binary, None) if frame.fin => {
                // Unfragmented fast path: no buffer copy.
                if frame.payload.len() > self.max_message {
                    return Err(FrameError::TooLarge);
                }
                return Self::complete(frame.opcode, frame.payload).map(Some);
            }
            (Opcode::Text | Opcode::Binary, None) => {
                self.in_progress = Some(frame.opcode);
            }
            (Opcode::Continuation, Some(_)) => {}
            (op, _) if op.is_control() => {
                return Err(FrameError::Protocol("control frame fed to assembler"));
            }
            _ => unreachable!(),
        }

        if self.buf.len() + frame.payload.len() > self.max_message {
            self.reset();
            return Err(FrameError::TooLarge);
        }
        self.buf.extend_from_slice(&frame.payload);

        if frame.fin {
            let kind = self.in_progress.take().expect("checked above");
            let payload = std::mem::take(&mut self.buf);
            return Self::complete(kind, payload).map(Some);
        }
        Ok(None)
    }

    fn complete(kind: Opcode, payload: Vec<u8>) -> Result<Message, FrameError> {
        match kind {
            Opcode::Binary => Ok(Message::Binary(payload)),
            Opcode::Text => String::from_utf8(payload)
                .map(Message::Text)
                .map_err(|_| FrameError::BadUtf8),
            _ => unreachable!(),
        }
    }

    fn reset(&mut self) {
        self.buf.clear();
        self.in_progress = None;
    }
}

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

    const MAX: usize = 1 << 20;

    fn dec(buf: &[u8]) -> Result<Option<(Frame, usize)>, FrameError> {
        decode(buf, true, MAX)
    }

    // ----- §5.7 worked examples. -----

    #[test]
    fn rfc_masked_hello() {
        // "A single-frame masked text message" containing "Hello".
        let wire = [0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58];
        let (f, used) = dec(&wire).unwrap().unwrap();
        assert_eq!(used, wire.len());
        assert!(f.fin);
        assert_eq!(f.opcode, Opcode::Text);
        assert_eq!(f.payload, b"Hello");
    }

    #[test]
    fn rfc_unmasked_hello_rejected_then_allowed() {
        let wire = [0x81, 0x05, b'H', b'e', b'l', b'l', b'o'];
        assert_eq!(
            dec(&wire),
            Err(FrameError::Protocol("unmasked client frame"))
        );
        // Same bytes are fine when masking isn't required (client mode).
        let (f, _) = decode(&wire, false, MAX).unwrap().unwrap();
        assert_eq!(f.payload, b"Hello");
    }

    #[test]
    fn rfc_fragmented_hel_lo() {
        let f1 = [0x01, 0x03, b'H', b'e', b'l'];
        let f2 = [0x80, 0x02, b'l', b'o'];
        let (a, _) = decode(&f1, false, MAX).unwrap().unwrap();
        let (b, _) = decode(&f2, false, MAX).unwrap().unwrap();
        assert!(!a.fin);
        assert_eq!(a.opcode, Opcode::Text);
        assert_eq!(b.opcode, Opcode::Continuation);

        let mut asm = Assembler::new(MAX);
        assert_eq!(asm.push(a).unwrap(), None);
        assert_eq!(
            asm.push(b).unwrap(),
            Some(Message::Text("Hello".into()))
        );
    }

    #[test]
    fn rfc_ping_pong_hello() {
        let ping = [0x89, 0x05, b'H', b'e', b'l', b'l', b'o'];
        let (f, _) = decode(&ping, false, MAX).unwrap().unwrap();
        assert_eq!(f.opcode, Opcode::Ping);
        assert_eq!(f.payload, b"Hello");
        // Pong response carries the same payload back, unmasked.
        let pong = Frame::new(Opcode::Pong, f.payload.clone()).encode();
        assert_eq!(pong, [0x8A, 0x05, b'H', b'e', b'l', b'l', b'o']);
    }

    #[test]
    fn rfc_256_bytes_binary_16bit_len() {
        let payload = vec![0xAB; 256];
        let wire = Frame::new(Opcode::Binary, payload.clone()).encode();
        assert_eq!(&wire[..4], &[0x82, 0x7E, 0x01, 0x00]);
        let (f, used) = decode(&wire, false, MAX).unwrap().unwrap();
        assert_eq!(used, wire.len());
        assert_eq!(f.payload, payload);
    }

    #[test]
    fn rfc_64k_binary_64bit_len() {
        let payload = vec![7u8; 65536];
        let wire = Frame::new(Opcode::Binary, payload.clone()).encode();
        assert_eq!(&wire[..10], &[0x82, 0x7F, 0, 0, 0, 0, 0, 1, 0, 0]);
        let (f, _) = decode(&wire, false, MAX).unwrap().unwrap();
        assert_eq!(f.payload.len(), 65536);
    }

    // ----- round trips. -----

    #[test]
    fn masked_roundtrip_all_lengths() {
        // Cross the 125/126 and 65535/65536 encoding boundaries.
        for len in [0, 1, 125, 126, 127, 65535, 65536] {
            let f = Frame::new(Opcode::Binary, vec![0x5A; len]);
            let wire = encode_masked(&f, [0xDE, 0xAD, 0xBE, 0xEF]);
            let (g, used) = decode(&wire, true, 1 << 20).unwrap().unwrap();
            assert_eq!(used, wire.len(), "len {len}");
            assert_eq!(g, f, "len {len}");
        }
    }

    #[test]
    fn decode_is_incremental_at_every_prefix() {
        let f = Frame::new(Opcode::Text, b"incremental".to_vec());
        let wire = encode_masked(&f, [1, 2, 3, 4]);
        for cut in 0..wire.len() {
            assert_eq!(dec(&wire[..cut]).unwrap(), None, "prefix {cut}");
        }
        assert!(dec(&wire).unwrap().is_some());
    }

    #[test]
    fn decode_leaves_trailing_bytes() {
        let f = Frame::new(Opcode::Text, b"a".to_vec());
        let mut wire = encode_masked(&f, [9, 9, 9, 9]);
        let first_len = wire.len();
        wire.extend_from_slice(&[0x81]); // start of a second frame
        let (_, used) = dec(&wire).unwrap().unwrap();
        assert_eq!(used, first_len);
    }

    // ----- protocol violations. -----

    #[test]
    fn rsv_bits_rejected() {
        for rsv in [0x40, 0x20, 0x10] {
            let wire = [0x81 | rsv, 0x80, 0, 0, 0, 0];
            assert!(matches!(dec(&wire), Err(FrameError::Protocol(_))), "rsv {rsv:#x}");
        }
    }

    #[test]
    fn reserved_opcodes_rejected() {
        for op in [0x3, 0x4, 0x5, 0x6, 0x7, 0xB, 0xC, 0xD, 0xE, 0xF] {
            let wire = [0x80 | op, 0x80, 0, 0, 0, 0];
            assert!(matches!(dec(&wire), Err(FrameError::Protocol(_))), "op {op:#x}");
        }
    }

    #[test]
    fn fragmented_control_rejected() {
        let wire = [0x09, 0x80, 0, 0, 0, 0]; // ping, FIN=0
        assert_eq!(dec(&wire), Err(FrameError::Protocol("fragmented control frame")));
    }

    #[test]
    fn oversize_control_rejected() {
        let wire = [0x89, 0x80 | 126, 0x00, 0x7E]; // ping, 16-bit len 126
        assert_eq!(dec(&wire), Err(FrameError::Protocol("control payload > 125")));
    }

    #[test]
    fn non_minimal_lengths_rejected() {
        // 16-bit form carrying 5.
        let w16 = [0x81, 0x80 | 126, 0x00, 0x05];
        assert_eq!(dec(&w16), Err(FrameError::Protocol("non-minimal 16-bit length")));
        // 64-bit form carrying 5.
        let w64 = [0x81, 0x80 | 127, 0, 0, 0, 0, 0, 0, 0, 0x05];
        assert_eq!(dec(&w64), Err(FrameError::Protocol("non-minimal 64-bit length")));
    }

    #[test]
    fn msb_set_64bit_length_rejected() {
        let wire = [0x81, 0x80 | 127, 0x80, 0, 0, 0, 0, 0, 0, 0];
        assert_eq!(dec(&wire), Err(FrameError::Protocol("64-bit length MSB set")));
    }

    #[test]
    fn oversize_fails_before_payload_arrives() {
        // Header promises 1 MiB + 1 with a 1 MiB cap: must fail with just
        // the header in hand, not Ok(None).
        let mut wire = vec![0x82, 0x80 | 127];
        wire.extend_from_slice(&((MAX as u64) + 1).to_be_bytes());
        assert_eq!(dec(&wire), Err(FrameError::TooLarge));
    }

    // ----- close payloads. -----

    #[test]
    fn close_payload_parsing() {
        assert_eq!(parse_close_payload(&[]), Ok(None));
        assert_eq!(
            parse_close_payload(&[0x03]),
            Err(FrameError::Protocol("1-byte close payload"))
        );
        let normal = close_payload(1000, "bye");
        assert_eq!(parse_close_payload(&normal), Ok(Some((1000, "bye".into()))));
        // Signalling-only and reserved codes are invalid on the wire.
        for bad in [999u16, 1004, 1005, 1006, 1015, 1100, 2999, 5000] {
            assert!(
                parse_close_payload(&bad.to_be_bytes()).is_err(),
                "code {bad}"
            );
        }
        for good in [1000u16, 1001, 1002, 1003, 1007, 1011, 3000, 4999] {
            assert!(
                parse_close_payload(&good.to_be_bytes()).is_ok(),
                "code {good}"
            );
        }
    }

    #[test]
    fn close_payload_reason_truncates_on_char_boundary() {
        let reason = "é".repeat(100); // 200 bytes of 2-byte chars
        let p = close_payload(1000, &reason);
        assert!(p.len() <= 125);
        assert_eq!(p.len() % 2, 0, "must not split the 2-byte char");
        assert!(std::str::from_utf8(&p[2..]).is_ok());
    }

    // ----- assembler. -----

    #[test]
    fn assembler_rejects_orphan_continuation() {
        let mut a = Assembler::new(MAX);
        let f = Frame { fin: true, opcode: Opcode::Continuation, payload: vec![] };
        assert_eq!(
            a.push(f),
            Err(FrameError::Protocol("continuation without a message"))
        );
    }

    #[test]
    fn assembler_rejects_interleaved_data_message() {
        let mut a = Assembler::new(MAX);
        let start = Frame { fin: false, opcode: Opcode::Text, payload: b"a".to_vec() };
        a.push(start).unwrap();
        let intruder = Frame::new(Opcode::Binary, b"b".to_vec());
        assert_eq!(
            a.push(intruder),
            Err(FrameError::Protocol("new data frame mid-fragmentation"))
        );
    }

    #[test]
    fn assembler_message_size_cap_spans_fragments() {
        let mut a = Assembler::new(10);
        let f1 = Frame { fin: false, opcode: Opcode::Binary, payload: vec![0; 6] };
        let f2 = Frame { fin: true, opcode: Opcode::Continuation, payload: vec![0; 6] };
        assert_eq!(a.push(f1).unwrap(), None);
        assert_eq!(a.push(f2), Err(FrameError::TooLarge));
        // And the unfragmented fast path is capped too.
        assert_eq!(
            a.push(Frame::new(Opcode::Binary, vec![0; 11])),
            Err(FrameError::TooLarge)
        );
    }

    #[test]
    fn assembler_text_utf8_validated_on_completion() {
        let mut a = Assembler::new(MAX);
        // A 2-byte UTF-8 char split across the fragment boundary must
        // still assemble — validation is whole-message.
        let bytes = "héllo".as_bytes();
        let f1 = Frame { fin: false, opcode: Opcode::Text, payload: bytes[..2].to_vec() };
        let f2 = Frame { fin: true, opcode: Opcode::Continuation, payload: bytes[2..].to_vec() };
        assert_eq!(a.push(f1).unwrap(), None);
        assert_eq!(a.push(f2).unwrap(), Some(Message::Text("héllo".into())));

        // Invalid UTF-8 → BadUtf8 (close 1007).
        let bad = Frame::new(Opcode::Text, vec![0xFF, 0xFE]);
        assert_eq!(a.push(bad), Err(FrameError::BadUtf8));
        assert_eq!(FrameError::BadUtf8.close_code(), 1007);
    }

    #[test]
    fn assembler_reusable_after_message() {
        let mut a = Assembler::new(MAX);
        for _ in 0..3 {
            let f1 = Frame { fin: false, opcode: Opcode::Text, payload: b"He".to_vec() };
            let f2 = Frame { fin: true, opcode: Opcode::Continuation, payload: b"llo".to_vec() };
            assert_eq!(a.push(f1).unwrap(), None);
            assert_eq!(a.push(f2).unwrap(), Some(Message::Text("Hello".into())));
        }
    }
}