imks/engine/codec.rs

use base64::{engine::general_purpose::STANDARD as BASE64, Engine};

use crate::engine::packet::{Packet, PacketData, PacketError, PacketType};

const RECORD_SEPARATOR: char = '\x1e';

pub fn encode_packet(packet: &Packet) -> String {
    let type_char = packet.packet_type as u8 + b'0';
    let type_str = type_char as char;

    match &packet.data {
        PacketData::Text(s) => format!("{type_str}{s}"),
        PacketData::Binary(b) => format!("{type_str}b{}", BASE64.encode(b)),
        PacketData::Empty => type_str.to_string(),
    }
}

pub fn encode_packet_binary_ws(packet: &Packet) -> Vec<u8> {
    let type_byte = packet.packet_type as u8 + b'0';

    match &packet.data {
        PacketData::Text(s) => {
            let mut buf = Vec::with_capacity(1 + s.len());
            buf.push(type_byte);
            buf.extend_from_slice(s.as_bytes());
            buf
        }
        PacketData::Binary(b) => {
            let mut buf = Vec::with_capacity(1 + b.len());
            buf.push(type_byte);
            buf.extend_from_slice(b);
            buf
        }
        PacketData::Empty => vec![type_byte],
    }
}

pub fn decode_packet(input: &str) -> Result<Packet, PacketError> {
    let mut chars = input.chars();
    let type_char = chars.next().ok_or(PacketError::Empty)?;
    let packet_type = PacketType::try_from(type_char)?;
    let rest: String = chars.collect();

    if rest.is_empty() {
        return Ok(Packet {
            packet_type,
            data: PacketData::Empty,
        });
    }

    if let Some(b64_data) = rest.strip_prefix('b') {
        let decoded = BASE64.decode(b64_data)?;
        return Ok(Packet {
            packet_type,
            data: PacketData::Binary(decoded),
        });
    }

    Ok(Packet {
        packet_type,
        data: PacketData::Text(rest),
    })
}

/// Decode a WebSocket binary frame into a Packet.
/// Handles both text-encoded packets (UTF-8 payload after type byte)
/// and binary packets (raw binary payload after type byte).
pub fn decode_packet_ws(input: &[u8]) -> Result<Packet, PacketError> {
    if input.is_empty() {
        return Err(PacketError::Empty);
    }

    let type_byte = input[0];
    let packet_type = PacketType::try_from(type_byte.wrapping_sub(b'0'))?;
    let rest = &input[1..];

    if rest.is_empty() {
        return Ok(Packet {
            packet_type,
            data: PacketData::Empty,
        });
    }

    // Try UTF-8 first; if it fails, treat as binary data
    match String::from_utf8(rest.to_vec()) {
        Ok(text) => Ok(Packet {
            packet_type,
            data: PacketData::Text(text),
        }),
        Err(_) => Ok(Packet {
            packet_type,
            data: PacketData::Binary(rest.to_vec()),
        }),
    }
}

pub fn encode_payload(packets: &[Packet]) -> String {
    packets
        .iter()
        .map(encode_packet)
        .collect::<Vec<_>>()
        .join(&RECORD_SEPARATOR.to_string())
}

pub fn decode_payload(input: &str) -> Result<Vec<Packet>, PacketError> {
    input
        .split(RECORD_SEPARATOR)
        .filter(|s| !s.is_empty())
        .map(decode_packet)
        .collect()
}

pub fn encode_webtransport_header(payload_len: usize, is_binary: bool) -> Vec<u8> {
    let binary_bit: u8 = if is_binary { 0x80 } else { 0x00 };

    if payload_len <= 125 {
        vec![binary_bit | (payload_len as u8)]
    } else if payload_len <= 65535 {
        let mut header = vec![binary_bit | 126];
        header.extend_from_slice(&(payload_len as u16).to_be_bytes());
        header
    } else {
        let mut header = vec![binary_bit | 127];
        header.extend_from_slice(&(payload_len as u64).to_be_bytes());
        header
    }
}

pub fn decode_webtransport_header(header: &[u8]) -> Option<(usize, bool)> {
    if header.is_empty() {
        return None;
    }

    let first = header[0];
    let is_binary = (first & 0x80) != 0;
    let len_indicator = first & 0x7f;

    if len_indicator <= 125 {
        Some((len_indicator as usize, is_binary))
    } else if len_indicator == 126 {
        if header.len() < 3 {
            return None;
        }
        let len = u16::from_be_bytes([header[1], header[2]]) as usize;
        Some((len, is_binary))
    } else {
        if header.len() < 9 {
            return None;
        }
        let len = u64::from_be_bytes([
            header[1], header[2], header[3], header[4], header[5], header[6], header[7], header[8],
        ]) as usize;
        Some((len, is_binary))
    }
}

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

    #[test]
    fn test_encode_decode_text_packet() {
        let packet = Packet::message_text("hello");
        let encoded = encode_packet(&packet);
        assert_eq!(encoded, "4hello");

        let decoded = decode_packet(&encoded).unwrap();
        assert_eq!(decoded.packet_type, PacketType::Message);
        assert_eq!(decoded.data, PacketData::Text("hello".to_string()));
    }

    #[test]
    fn test_encode_decode_binary_packet() {
        let packet = Packet::message_binary(vec![1, 2, 3, 4]);
        let encoded = encode_packet(&packet);
        assert_eq!(encoded, "4bAQIDBA==");

        let decoded = decode_packet(&encoded).unwrap();
        assert_eq!(decoded.packet_type, PacketType::Message);
        assert_eq!(decoded.data, PacketData::Binary(vec![1, 2, 3, 4]));
    }

    #[test]
    fn test_encode_decode_payload() {
        let packets = vec![
            Packet::message_text("hello"),
            Packet::ping(""),
            Packet::message_text("world"),
        ];
        let encoded = encode_payload(&packets);
        assert_eq!(encoded, "4hello\x1e2\x1e4world");

        let decoded = decode_payload(&encoded).unwrap();
        assert_eq!(decoded.len(), 3);
        assert_eq!(decoded[0].packet_type, PacketType::Message);
        assert_eq!(decoded[1].packet_type, PacketType::Ping);
        assert_eq!(decoded[2].packet_type, PacketType::Message);
    }

    #[test]
    fn test_webtransport_header() {
        let header = encode_webtransport_header(6, false);
        assert_eq!(header, vec![0x06]);
        let (len, is_binary) = decode_webtransport_header(&header).unwrap();
        assert_eq!(len, 6);
        assert!(!is_binary);

        let header = encode_webtransport_header(200, true);
        assert_eq!(header.len(), 3);
        let (len, is_binary) = decode_webtransport_header(&header).unwrap();
        assert_eq!(len, 200);
        assert!(is_binary);
    }

    #[test]
    fn test_decode_packet_ws_text() {
        let input = b"4hello";
        let decoded = decode_packet_ws(input).unwrap();
        assert_eq!(decoded.packet_type, PacketType::Message);
        assert_eq!(decoded.data, PacketData::Text("hello".to_string()));
    }

    #[test]
    fn test_decode_packet_ws_binary() {
        // Type byte 4 (Message) + raw binary payload (non-UTF-8)
        let input: Vec<u8> = vec![b'4', 0x80, 0xFF, 0x00, 0x01];
        let decoded = decode_packet_ws(&input).unwrap();
        assert_eq!(decoded.packet_type, PacketType::Message);
        assert_eq!(decoded.data, PacketData::Binary(vec![0x80, 0xFF, 0x00, 0x01]));
    }

    #[test]
    fn test_decode_packet_ws_empty() {
        let input = b"4";
        let decoded = decode_packet_ws(input).unwrap();
        assert_eq!(decoded.packet_type, PacketType::Message);
        assert_eq!(decoded.data, PacketData::Empty);
    }
}