protocol-core 0.3.12

//! Peer-to-peer repair request/response framing.
//!
//! All peer transport messages share a 1-byte frame type prefix followed by
//! a dCBOR body. The consumer dispatches on the first byte.

use dcbor::prelude::*;

use crate::ProtocolError;

const KEY_REQUEST_ID: u64 = 0;
const KEY_NAMED_PATH: u64 = 1;
const KEY_OBJECT_BYTES: u64 = 2;
const KEY_ATTACHMENT_ID: u64 = 0;
const KEY_CHUNK_INDEX: u64 = 1;
const KEY_TOTAL_CHUNKS: u64 = 2;
const KEY_DATA: u64 = 3;
const KEY_REASON: u64 = 1;

/// Frame type prefix byte for mesh peer transport messages.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum FrameType {
    RepairRequest = 0x01,
    RepairResponse = 0x02,
    RepairNotFound = 0x03,
    SyncManifest = 0x10,
    SyncChunk = 0x11,
    SyncAck = 0x12,
    AttachmentChunk = 0x20,
    AttachmentChunkAck = 0x21,
    AttachmentAbort = 0x22,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PeerRepairRequest {
    pub request_id: [u8; 16],
    pub named_path: String,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PeerRepairResponse {
    pub request_id: [u8; 16],
    pub object_bytes: Vec<u8>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PeerRepairNotFound {
    pub request_id: [u8; 16],
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttachmentChunk {
    pub attachment_id: [u8; 16],
    pub chunk_index: u32,
    pub total_chunks: u32,
    pub data: Vec<u8>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttachmentChunkAck {
    pub attachment_id: [u8; 16],
    pub chunk_index: u32,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttachmentAbort {
    pub attachment_id: [u8; 16],
    pub reason: String,
}

/// Parse the frame type from the first byte. Returns `None` for unknown types.
pub fn frame_type(bytes: &[u8]) -> Option<FrameType> {
    let first = *bytes.first()?;
    match first {
        0x01 => Some(FrameType::RepairRequest),
        0x02 => Some(FrameType::RepairResponse),
        0x03 => Some(FrameType::RepairNotFound),
        0x10 => Some(FrameType::SyncManifest),
        0x11 => Some(FrameType::SyncChunk),
        0x12 => Some(FrameType::SyncAck),
        0x20 => Some(FrameType::AttachmentChunk),
        0x21 => Some(FrameType::AttachmentChunkAck),
        0x22 => Some(FrameType::AttachmentAbort),
        _ => None,
    }
}

fn prepend_frame(frame: FrameType, body: Vec<u8>) -> Vec<u8> {
    let mut out = Vec::with_capacity(1 + body.len());
    out.push(frame as u8);
    out.extend_from_slice(&body);
    out
}

fn strip_frame(expected: FrameType, bytes: &[u8]) -> Result<&[u8], ProtocolError> {
    match bytes.first() {
        None => Err(ProtocolError::InvalidEncoding("empty frame".to_string())),
        Some(&b) if b != expected as u8 => Err(ProtocolError::InvalidEncoding(format!(
            "expected frame type 0x{:02x}, got 0x{:02x}",
            expected as u8, b
        ))),
        _ => Ok(&bytes[1..]),
    }
}

fn parse_map(bytes: &[u8]) -> Result<Map, ProtocolError> {
    let cbor =
        CBOR::try_from_data(bytes).map_err(|e| ProtocolError::InvalidEncoding(e.to_string()))?;
    cbor.try_into_map()
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))
}

fn extract_fixed<const N: usize>(
    map: &Map,
    key: u64,
    field: &str,
) -> Result<[u8; N], ProtocolError> {
    let cbor: CBOR = map
        .extract(key)
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))?;
    let bytes = cbor
        .try_into_byte_string()
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))?;
    bytes
        .try_into()
        .map_err(|_| ProtocolError::InvalidEnvelope(format!("{field} must be {N} bytes")))
}

fn extract_bytes(map: &Map, key: u64) -> Result<Vec<u8>, ProtocolError> {
    let cbor: CBOR = map
        .extract(key)
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))?;
    cbor.try_into_byte_string()
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))
}

fn extract_u64(map: &Map, key: u64) -> Result<u64, ProtocolError> {
    map.extract(key)
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))
}

fn extract_u32(map: &Map, key: u64, field: &str) -> Result<u32, ProtocolError> {
    let value = extract_u64(map, key)?;
    u32::try_from(value)
        .map_err(|_| ProtocolError::InvalidEnvelope(format!("{field} exceeds u32 range: {value}")))
}

fn extract_text(map: &Map, key: u64) -> Result<String, ProtocolError> {
    let cbor: CBOR = map
        .extract(key)
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))?;
    cbor.try_into_text()
        .map_err(|e| ProtocolError::InvalidEnvelope(e.to_string()))
}

pub fn encode_repair_request(req: &PeerRepairRequest) -> Vec<u8> {
    let mut map = Map::new();
    map.insert(KEY_REQUEST_ID, CBOR::to_byte_string(req.request_id));
    map.insert(KEY_NAMED_PATH, req.named_path.clone());
    prepend_frame(FrameType::RepairRequest, CBOR::from(map).to_cbor_data())
}

pub fn decode_repair_request(bytes: &[u8]) -> Result<PeerRepairRequest, ProtocolError> {
    let body = strip_frame(FrameType::RepairRequest, bytes)?;
    let map = parse_map(body)?;
    Ok(PeerRepairRequest {
        request_id: extract_fixed::<16>(&map, KEY_REQUEST_ID, "request_id")?,
        named_path: extract_text(&map, KEY_NAMED_PATH)?,
    })
}

pub fn encode_repair_response(resp: &PeerRepairResponse) -> Vec<u8> {
    let mut map = Map::new();
    map.insert(KEY_REQUEST_ID, CBOR::to_byte_string(resp.request_id));
    map.insert(KEY_OBJECT_BYTES, CBOR::to_byte_string(&resp.object_bytes));
    prepend_frame(FrameType::RepairResponse, CBOR::from(map).to_cbor_data())
}

pub fn decode_repair_response(bytes: &[u8]) -> Result<PeerRepairResponse, ProtocolError> {
    let body = strip_frame(FrameType::RepairResponse, bytes)?;
    let map = parse_map(body)?;
    Ok(PeerRepairResponse {
        request_id: extract_fixed::<16>(&map, KEY_REQUEST_ID, "request_id")?,
        object_bytes: extract_bytes(&map, KEY_OBJECT_BYTES)?,
    })
}

pub fn encode_repair_not_found(nf: &PeerRepairNotFound) -> Vec<u8> {
    let mut map = Map::new();
    map.insert(KEY_REQUEST_ID, CBOR::to_byte_string(nf.request_id));
    prepend_frame(FrameType::RepairNotFound, CBOR::from(map).to_cbor_data())
}

pub fn decode_repair_not_found(bytes: &[u8]) -> Result<PeerRepairNotFound, ProtocolError> {
    let body = strip_frame(FrameType::RepairNotFound, bytes)?;
    let map = parse_map(body)?;
    Ok(PeerRepairNotFound {
        request_id: extract_fixed::<16>(&map, KEY_REQUEST_ID, "request_id")?,
    })
}

pub fn encode_attachment_chunk(chunk: &AttachmentChunk) -> Vec<u8> {
    let mut map = Map::new();
    map.insert(KEY_ATTACHMENT_ID, CBOR::to_byte_string(chunk.attachment_id));
    map.insert(KEY_CHUNK_INDEX, u64::from(chunk.chunk_index));
    map.insert(KEY_TOTAL_CHUNKS, u64::from(chunk.total_chunks));
    map.insert(KEY_DATA, CBOR::to_byte_string(&chunk.data));
    prepend_frame(FrameType::AttachmentChunk, CBOR::from(map).to_cbor_data())
}

pub fn decode_attachment_chunk(bytes: &[u8]) -> Result<AttachmentChunk, ProtocolError> {
    let body = strip_frame(FrameType::AttachmentChunk, bytes)?;
    let map = parse_map(body)?;
    Ok(AttachmentChunk {
        attachment_id: extract_fixed::<16>(&map, KEY_ATTACHMENT_ID, "attachment_id")?,
        chunk_index: extract_u32(&map, KEY_CHUNK_INDEX, "chunk_index")?,
        total_chunks: extract_u32(&map, KEY_TOTAL_CHUNKS, "total_chunks")?,
        data: extract_bytes(&map, KEY_DATA)?,
    })
}

pub fn encode_attachment_chunk_ack(ack: &AttachmentChunkAck) -> Vec<u8> {
    let mut map = Map::new();
    map.insert(KEY_ATTACHMENT_ID, CBOR::to_byte_string(ack.attachment_id));
    map.insert(KEY_CHUNK_INDEX, u64::from(ack.chunk_index));
    prepend_frame(
        FrameType::AttachmentChunkAck,
        CBOR::from(map).to_cbor_data(),
    )
}

pub fn decode_attachment_chunk_ack(bytes: &[u8]) -> Result<AttachmentChunkAck, ProtocolError> {
    let body = strip_frame(FrameType::AttachmentChunkAck, bytes)?;
    let map = parse_map(body)?;
    Ok(AttachmentChunkAck {
        attachment_id: extract_fixed::<16>(&map, KEY_ATTACHMENT_ID, "attachment_id")?,
        chunk_index: extract_u32(&map, KEY_CHUNK_INDEX, "chunk_index")?,
    })
}

pub fn encode_attachment_abort(abort: &AttachmentAbort) -> Vec<u8> {
    let mut map = Map::new();
    map.insert(KEY_ATTACHMENT_ID, CBOR::to_byte_string(abort.attachment_id));
    map.insert(KEY_REASON, abort.reason.clone());
    prepend_frame(FrameType::AttachmentAbort, CBOR::from(map).to_cbor_data())
}

pub fn decode_attachment_abort(bytes: &[u8]) -> Result<AttachmentAbort, ProtocolError> {
    let body = strip_frame(FrameType::AttachmentAbort, bytes)?;
    let map = parse_map(body)?;
    Ok(AttachmentAbort {
        attachment_id: extract_fixed::<16>(&map, KEY_ATTACHMENT_ID, "attachment_id")?,
        reason: extract_text(&map, KEY_REASON)?,
    })
}

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

    #[test]
    fn test_frame_type_identifies_all_variants() {
        let cases: &[(u8, FrameType)] = &[
            (0x01, FrameType::RepairRequest),
            (0x02, FrameType::RepairResponse),
            (0x03, FrameType::RepairNotFound),
            (0x10, FrameType::SyncManifest),
            (0x11, FrameType::SyncChunk),
            (0x12, FrameType::SyncAck),
            (0x20, FrameType::AttachmentChunk),
            (0x21, FrameType::AttachmentChunkAck),
            (0x22, FrameType::AttachmentAbort),
        ];
        for &(byte, expected) in cases {
            assert_eq!(
                frame_type(&[byte, 0x00]),
                Some(expected),
                "byte 0x{byte:02x}"
            );
        }
    }

    #[test]
    fn test_frame_type_returns_none_for_unknown_byte() {
        assert_eq!(frame_type(&[0xFF]), None);
        assert_eq!(frame_type(&[0x00]), None);
        assert_eq!(frame_type(&[0x04]), None);
    }

    #[test]
    fn test_frame_type_returns_none_for_empty_slice() {
        assert_eq!(frame_type(&[]), None);
    }

    #[test]
    fn test_repair_request_round_trip() {
        let req = PeerRepairRequest {
            request_id: [0xAB; 16],
            named_path: "mesh/room/deadbeef/object/cafebabe/1".to_string(),
        };
        let encoded = encode_repair_request(&req);
        assert_eq!(encoded[0], FrameType::RepairRequest as u8);
        let decoded = decode_repair_request(&encoded).expect("decode must succeed");
        assert_eq!(decoded, req);
    }

    #[test]
    fn test_repair_response_round_trip() {
        let resp = PeerRepairResponse {
            request_id: [0x01; 16],
            object_bytes: vec![0xDE, 0xAD, 0xBE, 0xEF],
        };
        let encoded = encode_repair_response(&resp);
        assert_eq!(encoded[0], FrameType::RepairResponse as u8);
        let decoded = decode_repair_response(&encoded).expect("decode must succeed");
        assert_eq!(decoded, resp);
    }

    #[test]
    fn test_repair_not_found_round_trip() {
        let nf = PeerRepairNotFound {
            request_id: [0xFF; 16],
        };
        let encoded = encode_repair_not_found(&nf);
        assert_eq!(encoded[0], FrameType::RepairNotFound as u8);
        let decoded = decode_repair_not_found(&encoded).expect("decode must succeed");
        assert_eq!(decoded, nf);
    }

    #[test]
    fn test_attachment_chunk_round_trip() {
        let chunk = AttachmentChunk {
            attachment_id: [0xCC; 16],
            chunk_index: 3,
            total_chunks: 10,
            data: vec![1, 2, 3, 4, 5],
        };
        let encoded = encode_attachment_chunk(&chunk);
        assert_eq!(encoded[0], FrameType::AttachmentChunk as u8);
        let decoded = decode_attachment_chunk(&encoded).expect("decode must succeed");
        assert_eq!(decoded, chunk);
    }

    #[test]
    fn test_attachment_chunk_ack_round_trip() {
        let ack = AttachmentChunkAck {
            attachment_id: [0xDD; 16],
            chunk_index: 7,
        };
        let encoded = encode_attachment_chunk_ack(&ack);
        assert_eq!(encoded[0], FrameType::AttachmentChunkAck as u8);
        let decoded = decode_attachment_chunk_ack(&encoded).expect("decode must succeed");
        assert_eq!(decoded, ack);
    }

    #[test]
    fn test_attachment_abort_round_trip() {
        let abort = AttachmentAbort {
            attachment_id: [0xEE; 16],
            reason: "transfer cancelled by sender".to_string(),
        };
        let encoded = encode_attachment_abort(&abort);
        assert_eq!(encoded[0], FrameType::AttachmentAbort as u8);
        let decoded = decode_attachment_abort(&encoded).expect("decode must succeed");
        assert_eq!(decoded, abort);
    }
}