robomotion 0.1.3

//! Large Message Object (LMO) support for payloads > 256KB.

use crate::runtime::{client, Result, RobomotionError};
use crate::utils;
use data_encoding::BASE32;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::path::PathBuf;
use tokio::fs;
use uuid::Uuid;

/// Magic number to identify LMO objects.
pub const LMO_MAGIC: i64 = 0x1343B7E;

/// Size limit for LMO (256KB).
pub const LMO_LIMIT: usize = 256 << 10;

/// LMO version.
pub const LMO_VERSION: i32 = 0x01;

/// Number of bytes to include in the head preview.
pub const LMO_HEAD: usize = 100;

/// Large Message Object for handling payloads > 256KB.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LargeMessageObject {
    pub magic: i64,
    pub version: i32,
    pub id: String,
    pub head: String,
    pub size: i64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub data: Option<Value>,
}

impl LargeMessageObject {
    /// Get the underlying data value.
    pub fn value(&self) -> Option<&Value> {
        self.data.as_ref()
    }
}

/// Generate a new unique ID for LMO.
pub fn new_id() -> String {
    let uuid_bytes = Uuid::new_v4().as_bytes().to_vec();
    let encoded = BASE32.encode(&uuid_bytes);
    // Truncate to 26 chars (removes padding)
    encoded.chars().take(26).collect::<String>().to_lowercase()
}

/// Check if a value is an LMO.
pub fn is_lmo(value: &Value) -> bool {
    if let Some(obj) = value.as_object() {
        if let Some(magic) = obj.get("magic") {
            if let Some(magic_num) = magic.as_i64() {
                return magic_num == LMO_MAGIC;
            }
            if let Some(magic_float) = magic.as_f64() {
                return magic_float as i64 == LMO_MAGIC;
            }
        }
    }
    false
}

/// Serialize a value to LMO if it exceeds the size limit.
///
/// Returns None if the value doesn't need LMO serialization.
pub async fn serialize_lmo<T: Serialize>(value: &T) -> Result<Option<LargeMessageObject>> {
    if !client::is_lmo_capable().await {
        return Ok(None);
    }

    let data = serde_json::to_vec(value)?;
    let data_len = data.len();

    if data_len < LMO_LIMIT {
        return Ok(None);
    }

    let id = new_id();
    let head = String::from_utf8_lossy(&data[..std::cmp::min(LMO_HEAD, data_len)]).to_string();

    let lmo = LargeMessageObject {
        magic: LMO_MAGIC,
        version: LMO_VERSION,
        id: id.clone(),
        head,
        size: data_len as i64,
        data: Some(serde_json::from_slice(&data)?),
    };

    // Save to file
    let robot_id = client::get_robot_id().await.unwrap_or_else(|_| "unknown".to_string());
    let dir = get_lmo_dir(&robot_id);
    fs::create_dir_all(&dir).await?;

    let file_path = dir.join(format!("{}.lmo", id));
    let lmo_json = serde_json::to_vec(&lmo)?;
    fs::write(&file_path, &lmo_json).await?;

    // Return LMO without data (data is stored in file)
    Ok(Some(LargeMessageObject {
        magic: LMO_MAGIC,
        version: LMO_VERSION,
        id,
        head: lmo.head,
        size: lmo.size,
        data: None,
    }))
}

/// Deserialize an LMO from its ID.
pub async fn deserialize_lmo(id: &str) -> Result<LargeMessageObject> {
    let robot_id = client::get_robot_id().await.unwrap_or_else(|_| "unknown".to_string());
    let file_path = get_lmo_dir(&robot_id).join(format!("{}.lmo", id));

    let data = fs::read(&file_path).await?;
    let lmo: LargeMessageObject = serde_json::from_slice(&data)?;

    Ok(lmo)
}

/// Deserialize an LMO from a map value.
pub async fn deserialize_lmo_from_map(
    map: &serde_json::Map<String, Value>,
) -> Result<LargeMessageObject> {
    let id = map
        .get("id")
        .and_then(|v| v.as_str())
        .ok_or_else(|| RobomotionError::Lmo("Missing LMO id".to_string()))?;

    deserialize_lmo(id).await
}

/// Delete an LMO file by ID.
pub async fn delete_lmo_by_id(id: &str) -> Result<()> {
    let robot_id = client::get_robot_id().await.unwrap_or_else(|_| "unknown".to_string());
    let file_path = get_lmo_dir(&robot_id).join(format!("{}.lmo", id));

    if file_path.exists() {
        fs::remove_file(&file_path).await?;
    }

    Ok(())
}

/// Pack a message, converting large values to LMOs.
pub async fn pack_message(msg: &mut serde_json::Map<String, Value>) -> Result<()> {
    if !client::is_lmo_capable().await {
        return Ok(());
    }

    let keys: Vec<String> = msg.keys().cloned().collect();
    for key in keys {
        if let Some(value) = msg.remove(&key) {
            if let Some(lmo) = serialize_lmo(&value).await? {
                msg.insert(key, serde_json::to_value(&lmo)?);
            } else {
                msg.insert(key, value);
            }
        }
    }

    Ok(())
}

/// Pack message bytes.
pub async fn pack_message_bytes(data: &[u8]) -> Result<Vec<u8>> {
    if !client::is_lmo_capable().await || data.len() < LMO_LIMIT {
        return Ok(data.to_vec());
    }

    let mut msg: serde_json::Map<String, Value> = serde_json::from_slice(data)?;
    pack_message(&mut msg).await?;
    Ok(serde_json::to_vec(&msg)?)
}

/// Unpack a message, resolving LMOs to their actual values.
pub async fn unpack_message(
    data: &[u8],
    msg: &mut serde_json::Map<String, Value>,
) -> Result<()> {
    if !client::is_lmo_capable().await {
        return Ok(());
    }

    let parsed: serde_json::Map<String, Value> = serde_json::from_slice(data)?;

    for (key, value) in parsed {
        if is_lmo(&value) {
            if let Some(obj) = value.as_object() {
                let lmo = deserialize_lmo_from_map(obj).await?;
                if let Some(data) = lmo.data {
                    msg.insert(key, data);
                    continue;
                }
            }
        }
        msg.insert(key, value);
    }

    Ok(())
}

/// Unpack message bytes.
pub async fn unpack_message_bytes(data: &[u8]) -> Result<Vec<u8>> {
    let mut msg = serde_json::Map::new();
    unpack_message(data, &mut msg).await?;
    Ok(serde_json::to_vec(&msg)?)
}

/// Get the LMO storage directory.
fn get_lmo_dir(robot_id: &str) -> PathBuf {
    let temp_path = utils::get_temp_path();
    temp_path.join("robots").join(robot_id)
}