featherChat/warzone/crates/warzone-protocol/src/message.rs

use serde::{Deserialize, Serialize};

use crate::ratchet::RatchetHeader;
use crate::types::{Fingerprint, MessageId, SessionId};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum MessageType {
    Text,
    File,
    KeyExchange,
    Receipt,
}

/// An encrypted message on the wire.
#[derive(Clone, Serialize, Deserialize)]
pub struct WarzoneMessage {
    pub version: u8,
    pub id: MessageId,
    pub from: Fingerprint,
    pub to: Fingerprint,
    pub timestamp: i64,
    pub msg_type: MessageType,
    pub session_id: SessionId,
    pub ratchet_header: RatchetHeader,
    pub ciphertext: Vec<u8>,
    pub signature: Vec<u8>,
}

/// Plaintext message content (inside the encrypted envelope).
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum MessageContent {
    Text { body: String },
    File { filename: String, data: Vec<u8> },
    Receipt { message_id: MessageId },
}

/// Receipt type: delivered (received + decrypted) or read (user viewed).
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum ReceiptType {
    Delivered,
    Read,
}

/// Wire message format for transport between clients.
/// Used by both CLI and WASM — MUST be identical for interop.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum WireMessage {
    /// First message to a peer: X3DH key exchange + first ratchet message.
    KeyExchange {
        id: String,
        sender_fingerprint: String,
        sender_identity_encryption_key: [u8; 32],
        ephemeral_public: [u8; 32],
        used_one_time_pre_key_id: Option<u32>,
        ratchet_message: crate::ratchet::RatchetMessage,
    },
    /// Subsequent messages: ratchet-encrypted.
    Message {
        id: String,
        sender_fingerprint: String,
        ratchet_message: crate::ratchet::RatchetMessage,
    },
    /// Delivery / read receipt (plaintext, not encrypted).
    Receipt {
        sender_fingerprint: String,
        message_id: String,
        receipt_type: ReceiptType,
    },
    /// File transfer header: announces an incoming chunked file.
    FileHeader {
        id: String,
        sender_fingerprint: String,
        filename: String,
        file_size: u64,
        total_chunks: u32,
        sha256: String,
    },
    /// A single chunk of a file transfer (data is ratchet-encrypted).
    FileChunk {
        id: String,
        sender_fingerprint: String,
        filename: String,
        chunk_index: u32,
        total_chunks: u32,
        data: Vec<u8>,
    },
    /// Group message encrypted with sender key (O(1) instead of O(N)).
    GroupSenderKey {
        id: String,
        sender_fingerprint: String,
        group_name: String,
        generation: u32,
        counter: u32,
        ciphertext: Vec<u8>,
    },
    /// Sender key distribution: share your sender key with a group member.
    /// This is sent via 1:1 encrypted channel (wrapped in KeyExchange/Message).
    SenderKeyDistribution {
        sender_fingerprint: String,
        group_name: String,
        chain_key: [u8; 32],
        generation: u32,
    },
    /// Call signaling: SDP offers/answers, ICE candidates, call control.
    /// Routed through featherChat's E2E encrypted channel for WarzonePhone integration.
    CallSignal {
        id: String,
        sender_fingerprint: String,
        signal_type: CallSignalType,
        /// SDP offer/answer body, ICE candidate, or empty for hangup/reject.
        payload: String,
        /// Target peer (for 1:1) or group/room name (for group calls).
        target: String,
    },
}

/// Call signaling types for WarzonePhone integration.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum CallSignalType {
    /// Initiate a call (contains SDP offer or WZP connection params).
    Offer,
    /// Accept a call (contains SDP answer or WZP connection params).
    Answer,
    /// ICE candidate for NAT traversal.
    IceCandidate,
    /// Hang up / end call.
    Hangup,
    /// Reject incoming call.
    Reject,
    /// Call is ringing on the other side.
    Ringing,
    /// Peer is busy.
    Busy,
}

/// Current wire protocol version.
pub const WIRE_VERSION: u8 = 1;
/// Magic bytes to identify versioned envelope: "WZ"
pub const WIRE_MAGIC: [u8; 2] = [0x57, 0x5A];

/// Serialize a WireMessage with version envelope.
/// Format: [0x57][0x5A][version: u8][length: u32 BE][bincode payload]
pub fn serialize_envelope(msg: &WireMessage) -> Result<Vec<u8>, String> {
    let payload =
        bincode::serialize(msg).map_err(|e| format!("serialize: {}", e))?;
    let len = payload.len() as u32;
    let mut out = Vec::with_capacity(7 + payload.len());
    out.extend_from_slice(&WIRE_MAGIC);
    out.push(WIRE_VERSION);
    out.extend_from_slice(&len.to_be_bytes());
    out.extend_from_slice(&payload);
    Ok(out)
}

/// Deserialize a WireMessage, handling both envelope and legacy formats.
/// - Envelope: [0x57][0x5A][version][length][payload]
/// - Legacy: raw bincode (no envelope)
pub fn deserialize_envelope(data: &[u8]) -> Result<WireMessage, String> {
    if data.len() >= 7 && data[0] == WIRE_MAGIC[0] && data[1] == WIRE_MAGIC[1] {
        let version = data[2];
        let len =
            u32::from_be_bytes([data[3], data[4], data[5], data[6]]) as usize;
        if version > WIRE_VERSION {
            return Err(format!(
                "unsupported wire version {} (max {}). Please update your client.",
                version, WIRE_VERSION
            ));
        }
        if data.len() < 7 + len {
            return Err("truncated envelope".to_string());
        }
        bincode::deserialize(&data[7..7 + len])
            .map_err(|e| format!("v{} deserialize: {}", version, e))
    } else {
        // Legacy: raw bincode
        bincode::deserialize(data)
            .map_err(|e| format!("legacy deserialize: {}", e))
    }
}

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

    #[test]
    fn envelope_roundtrip() {
        let msg = WireMessage::Receipt {
            sender_fingerprint: "abc123".to_string(),
            message_id: "msg-001".to_string(),
            receipt_type: ReceiptType::Delivered,
        };
        let envelope = serialize_envelope(&msg).unwrap();
        assert_eq!(&envelope[..2], &WIRE_MAGIC);
        assert_eq!(envelope[2], WIRE_VERSION);

        let decoded = deserialize_envelope(&envelope).unwrap();
        match decoded {
            WireMessage::Receipt { message_id, .. } => {
                assert_eq!(message_id, "msg-001")
            }
            _ => panic!("wrong variant"),
        }
    }

    #[test]
    fn legacy_still_works() {
        let msg = WireMessage::Receipt {
            sender_fingerprint: "abc123".to_string(),
            message_id: "msg-002".to_string(),
            receipt_type: ReceiptType::Read,
        };
        let raw = bincode::serialize(&msg).unwrap();
        let decoded = deserialize_envelope(&raw).unwrap();
        match decoded {
            WireMessage::Receipt { message_id, .. } => {
                assert_eq!(message_id, "msg-002")
            }
            _ => panic!("wrong variant"),
        }
    }

    #[test]
    fn future_version_rejected() {
        let mut envelope = serialize_envelope(&WireMessage::Receipt {
            sender_fingerprint: "x".into(),
            message_id: "y".into(),
            receipt_type: ReceiptType::Delivered,
        })
        .unwrap();
        envelope[2] = 99; // fake future version
        let result = deserialize_envelope(&envelope);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("unsupported wire version"));
    }
}