//! UoT (UDP-over-TCP) framing layer

use bytes::{Buf, BufMut, Bytes, BytesMut};
use thiserror::Error;

/// Maximum UDP datagram size
pub const MAX_UDP_SIZE: usize = 65535;
/// Frame header: 2 bytes length + 1 byte flags
const HEADER_SIZE: usize = 3;

#[derive(Error, Debug)]
pub enum FrameError {
    #[error("frame too large: {0} bytes")]
    TooLarge(usize),
    #[error("incomplete frame")]
    Incomplete,
    #[error("invalid frame")]
    Invalid,
}

#[derive(Clone, Copy, Debug)]
pub struct FrameFlags(u8);

impl FrameFlags {
    pub const DATA: Self = Self(0x00);
    pub const CONTROL: Self = Self(0x01);
    pub const PADDING: Self = Self(0x02);

    pub fn is_data(self) -> bool {
        self.0 == 0x00
    }
}

/// Encapsulate UDP datagram into framed format with random padding
pub fn encode_frame(data: &[u8], flags: FrameFlags, padding: usize) -> Result<Bytes, FrameError> {
    let total_len = data.len() + padding;
    if total_len > MAX_UDP_SIZE {
        return Err(FrameError::TooLarge(total_len));
    }

    let mut buf = BytesMut::with_capacity(HEADER_SIZE + total_len);
    buf.put_u16(total_len as u16);
    buf.put_u8(flags.0);
    buf.put_slice(data);

    // Random padding to avoid fingerprinting
    if padding > 0 {
        let pad: Vec<u8> = (0..padding).map(|_| rand::random()).collect();
        buf.put_slice(&pad);
    }

    Ok(buf.freeze())
}

/// Decode frame from buffer, returns (payload, data_length, flags)
pub fn decode_frame(buf: &mut BytesMut) -> Result<Option<(Bytes, FrameFlags)>, FrameError> {
    if buf.len() < HEADER_SIZE {
        return Ok(None);
    }

    let len = u16::from_be_bytes([buf[0], buf[1]]) as usize;
    let flags = FrameFlags(buf[2]);

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

    buf.advance(HEADER_SIZE);
    let data = buf.split_to(len).freeze();

    Ok(Some((data, flags)))
}

/// MTU-aware fragmentation
pub struct Fragmenter {
    mtu: usize,
}

impl Fragmenter {
    pub fn new(mtu: usize) -> Self {
        Self { mtu }
    }

    pub fn fragment(&self, data: &[u8]) -> Vec<Bytes> {
        data.chunks(self.mtu)
            .map(|chunk| Bytes::copy_from_slice(chunk))
            .collect()
    }
}

/// Reassembly buffer for fragmented packets
pub struct Reassembler {
    buffer: BytesMut,
    expected_len: Option<usize>,
}

impl Reassembler {
    pub fn new() -> Self {
        Self {
            buffer: BytesMut::new(),
            expected_len: None,
        }
    }

    pub fn push(&mut self, fragment: &[u8], total_len: usize) {
        if self.expected_len.is_none() {
            self.expected_len = Some(total_len);
        }
        self.buffer.extend_from_slice(fragment);
    }

    pub fn is_complete(&self) -> bool {
        self.expected_len
            .map(|len| self.buffer.len() >= len)
            .unwrap_or(false)
    }

    pub fn take(&mut self) -> Option<Bytes> {
        if self.is_complete() {
            self.expected_len = None;
            Some(self.buffer.split().freeze())
        } else {
            None
        }
    }
}

impl Default for Reassembler {
    fn default() -> Self {
        Self::new()
    }
}
