imks/socket/parser.rs

use serde_json::Value;

use crate::socket::packet::{Packet, PacketError, PacketType};

pub fn encode(packet: &Packet) -> String {
    let type_char = packet.packet_type as u8 + b'0';
    let mut result = String::new();

    result.push(type_char as char);

    if packet.has_binary() {
        result.push_str(&packet.attachment_count().to_string());
        result.push('-');
    }

    if packet.namespace != "/" {
        result.push_str(&packet.namespace);
        result.push(',');
    }

    if let Some(id) = packet.id {
        result.push_str(&id.to_string());
    }

    if let Some(ref data) = packet.data {
        if packet.has_binary() {
            let data_with_placeholders = replace_binary_with_placeholders(data, packet.attachment_count());
            let encoded_data = serde_json::to_string(&data_with_placeholders)
                .unwrap_or_else(|e| {
                    tracing::error!("Failed to serialize socket packet data: {}", e);
                    "null".to_string()
                });
            result.push_str(&encoded_data);
        } else {
            let encoded_data = serde_json::to_string(data)
                .unwrap_or_else(|e| {
                    tracing::error!("Failed to serialize socket packet data: {}", e);
                    "null".to_string()
                });
            result.push_str(&encoded_data);
        }
    }

    result
}

pub fn encode_with_attachments(packet: &Packet) -> Vec<Vec<u8>> {
    let mut result = Vec::new();

    let encoded = encode(packet);
    result.push(encoded.into_bytes());

    for attachment in &packet.attachments {
        result.push(attachment.clone());
    }

    result
}

pub fn decode(input: &str) -> Result<Packet, PacketError> {
    if input.is_empty() {
        return Err(PacketError::Empty);
    }

    let mut chars = input.chars().peekable();

    let type_char = chars.next().ok_or(PacketError::Empty)?;
    let packet_type = PacketType::try_from(type_char)?;

    let attachment_count = if matches!(packet_type, PacketType::BinaryEvent | PacketType::BinaryAck) {
        let mut count_str = String::new();
        while let Some(&c) = chars.peek() {
            if c == '-' {
                chars.next();
                break;
            }
            if c.is_ascii_digit() {
                count_str.push(c);
                chars.next();
            } else {
                break;
            }
        }
        count_str.parse::<usize>().unwrap_or(0)
    } else {
        0
    };

    let remaining: String = chars.collect();

    let (namespace, rest) = if let Some(after_slash) = remaining.strip_prefix('/') {
        // Check if this is a custom namespace (has a comma separating namespace from data/id)
        // or if '/' is just the root namespace prefix followed immediately by data
        if let Some(comma_pos) = after_slash.find(',') {
            let ns = format!("/{}", &after_slash[..comma_pos]);
            let rest = after_slash[comma_pos + 1..].to_string();
            (ns, rest)
        } else if after_slash.starts_with('[')
            || after_slash.starts_with(|c: char| c.is_ascii_digit())
            || after_slash.is_empty()
        {
            // '/[' means '/' is the root namespace and '[' starts the data
            // '/<digits>' means root namespace followed by ack id
            // '/' alone means disconnect on root namespace
            ("/".to_string(), after_slash.to_string())
        } else {
            // Non-root namespace without data (e.g., disconnect on custom namespace)
            (remaining, String::new())
        }
    } else {
        ("/".to_string(), remaining)
    };

    let (id, data_str) = parse_id_and_data(&rest);

    let data = if data_str.is_empty() {
        None
    } else {
        Some(serde_json::from_str(&data_str)?)
    };

    Ok(Packet {
        packet_type,
        namespace,
        data,
        id,
        attachments: Vec::new(),
        // Store attachment_count for binary packets; actual attachments come via decode_with_attachments
        expected_attachments: if attachment_count > 0 { Some(attachment_count) } else { None },
    })
}

pub fn decode_with_attachments(
    main_packet: &str,
    attachments: Vec<Vec<u8>>,
) -> Result<Packet, PacketError> {
    let mut packet = decode(main_packet)?;

    let expected = packet.expected_attachments.unwrap_or(0);
    if expected != attachments.len() {
        return Err(PacketError::InvalidAttachmentCount);
    }

    packet.attachments = attachments;
    packet.expected_attachments = None;

    if packet.has_binary() {
        if let Some(ref data) = packet.data {
            packet.data = Some(replace_placeholders_with_binary(data, &packet.attachments));
        }
    }

    Ok(packet)
}

fn parse_id_and_data(input: &str) -> (Option<u64>, String) {
    let mut id_str = String::new();
    let mut chars = input.chars().peekable();

    while let Some(&c) = chars.peek() {
        if c.is_ascii_digit() {
            id_str.push(c);
            chars.next();
        } else {
            break;
        }
    }

    let id = if id_str.is_empty() {
        None
    } else {
        id_str.parse::<u64>().ok()
    };

    let data: String = chars.collect();

    (id, data)
}

/// Replace binary values in the data with { "_placeholder": true, "num": N } placeholders.
/// This is used when encoding binary events/acks for transmission over text-based transports.
fn replace_binary_with_placeholders(value: &Value, total_attachments: usize) -> Value {
    match value {
        Value::Array(arr) => {
            let mut placeholder_idx = total_attachments; // Start from known count
            let new_arr: Vec<Value> = arr
                .iter()
                .map(|v| replace_binary_with_placeholders_inner(v, &mut placeholder_idx))
                .collect();
            Value::Array(new_arr)
        }
        Value::Object(map) => {
            let mut placeholder_idx = total_attachments;
            let mut new_map = serde_json::Map::new();
            for (k, v) in map {
                new_map.insert(
                    k.clone(),
                    replace_binary_with_placeholders_inner(v, &mut placeholder_idx),
                );
            }
            Value::Object(new_map)
        }
        _ => value.clone(),
    }
}

fn replace_binary_with_placeholders_inner(value: &Value, placeholder_idx: &mut usize) -> Value {
    match value {
        Value::Array(arr) => {
            let new_arr: Vec<Value> = arr
                .iter()
                .map(|v| replace_binary_with_placeholders_inner(v, placeholder_idx))
                .collect();
            Value::Array(new_arr)
        }
        Value::Object(map) => {
            let mut new_map = serde_json::Map::new();
            for (k, v) in map {
                new_map.insert(
                    k.clone(),
                    replace_binary_with_placeholders_inner(v, placeholder_idx),
                );
            }
            Value::Object(new_map)
        }
        // Binary data would be represented as base64 strings in the initial data;
        // in the Socket.IO protocol, binary attachments are separate and referenced by placeholder.
        // This function handles the case where the data structure itself contains placeholder markers.
        _ => value.clone(),
    }
}

fn replace_placeholders_with_binary(value: &Value, attachments: &[Vec<u8>]) -> Value {
    match value {
        Value::Object(map) => {
            // Check if this is a placeholder object: { "_placeholder": true, "num": N }
            if let (Some(Value::Bool(true)), Some(Value::Number(num))) =
                (map.get("_placeholder"), map.get("num"))
            {
                if let Some(idx) = num.as_u64() {
                    if let Some(attachment) = attachments.get(idx as usize) {
                        return Value::String(base64::Engine::encode(
                            &base64::engine::general_purpose::STANDARD,
                            attachment,
                        ));
                    }
                }
            }

            let mut new_map = serde_json::Map::new();
            for (k, v) in map {
                new_map.insert(k.clone(), replace_placeholders_with_binary(v, attachments));
            }
            Value::Object(new_map)
        }
        Value::Array(arr) => Value::Array(
            arr.iter()
                .map(|v| replace_placeholders_with_binary(v, attachments))
                .collect(),
        ),
        _ => value.clone(),
    }
}

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

    #[test]
    fn test_encode_connect() {
        let packet = Packet::connect("/", None);
        let encoded = encode(&packet);
        assert_eq!(encoded, "0");

        let packet = Packet::connect("/admin", Some(json!({"sid": "abc"})));
        let encoded = encode(&packet);
        assert_eq!(encoded, "0/admin,{\"sid\":\"abc\"}");
    }

    #[test]
    fn test_encode_event() {
        let packet = Packet::event("/", json!(["foo"]), None);
        let encoded = encode(&packet);
        assert_eq!(encoded, "2[\"foo\"]");

        let packet = Packet::event("/admin", json!(["bar"]), None);
        let encoded = encode(&packet);
        assert_eq!(encoded, "2/admin,[\"bar\"]");
    }

    #[test]
    fn test_encode_event_with_ack() {
        let packet = Packet::event("/", json!(["foo"]), Some(12));
        let encoded = encode(&packet);
        assert_eq!(encoded, "212[\"foo\"]");
    }

    #[test]
    fn test_encode_ack() {
        let packet = Packet::ack("/", json!([]), 12);
        let encoded = encode(&packet);
        assert_eq!(encoded, "312[]");
    }

    #[test]
    fn test_encode_disconnect() {
        let packet = Packet::disconnect("/");
        let encoded = encode(&packet);
        assert_eq!(encoded, "1");

        let packet = Packet::disconnect("/admin");
        let encoded = encode(&packet);
        assert_eq!(encoded, "1/admin,");
    }

    #[test]
    fn test_encode_connect_error() {
        let packet = Packet::connect_error("/", "Not authorized");
        let encoded = encode(&packet);
        assert_eq!(encoded, "4{\"message\":\"Not authorized\"}");
    }

    #[test]
    fn test_decode_connect() {
        let packet = decode("0").unwrap();
        assert_eq!(packet.packet_type, PacketType::Connect);
        assert_eq!(packet.namespace, "/");
        assert!(packet.data.is_none());
    }

    #[test]
    fn test_decode_connect_with_namespace() {
        let packet = decode("0/admin,{\"sid\":\"abc\"}").unwrap();
        assert_eq!(packet.packet_type, PacketType::Connect);
        assert_eq!(packet.namespace, "/admin");
        assert!(packet.data.is_some());
    }

    #[test]
    fn test_decode_event() {
        let packet = decode("2[\"foo\"]").unwrap();
        assert_eq!(packet.packet_type, PacketType::Event);
        assert_eq!(packet.namespace, "/");
        assert_eq!(packet.data, Some(json!(["foo"])));
    }

    #[test]
    fn test_decode_event_with_namespace() {
        let packet = decode("2/admin,[\"bar\"]").unwrap();
        assert_eq!(packet.packet_type, PacketType::Event);
        assert_eq!(packet.namespace, "/admin");
        assert_eq!(packet.data, Some(json!(["bar"])));
    }

    #[test]
    fn test_decode_event_with_ack() {
        let packet = decode("212[\"foo\"]").unwrap();
        assert_eq!(packet.packet_type, PacketType::Event);
        assert_eq!(packet.id, Some(12));
        assert_eq!(packet.data, Some(json!(["foo"])));
    }

    #[test]
    fn test_decode_ack() {
        let packet = decode("312[]").unwrap();
        assert_eq!(packet.packet_type, PacketType::Ack);
        assert_eq!(packet.id, Some(12));
    }

    #[test]
    fn test_decode_disconnect() {
        let packet = decode("1").unwrap();
        assert_eq!(packet.packet_type, PacketType::Disconnect);
        assert_eq!(packet.namespace, "/");
    }

    #[test]
    fn test_decode_disconnect_with_namespace() {
        let packet = decode("1/admin,").unwrap();
        assert_eq!(packet.packet_type, PacketType::Disconnect);
        assert_eq!(packet.namespace, "/admin");
    }

    #[test]
    fn test_decode_binary_event_attachment_count() {
        let packet = decode("51-[\"baz\",{\"_placeholder\":true,\"num\":0}]").unwrap();
        assert_eq!(packet.packet_type, PacketType::BinaryEvent);
        assert_eq!(packet.expected_attachments, Some(1));
        assert_eq!(packet.namespace, "/");
    }
}