robomotion 0.1.3

//! Message context for data flow between nodes.
//!
//! The Context struct provides access to the message payload flowing through nodes,
//! allowing nodes to read input values and set output values.

use serde_json::{json, Value};
use std::sync::Arc;
use parking_lot::RwLock;

/// Options for getting raw message data
pub type GetOption = Box<dyn Fn(Vec<u8>) -> crate::runtime::Result<Vec<u8>> + Send + Sync>;

/// Options for setting raw message data
pub type SetOption = Box<dyn Fn(Vec<u8>) -> crate::runtime::Result<Vec<u8>> + Send + Sync>;

/// Message context for data flow between nodes.
///
/// Context provides typed access to the message payload, allowing nodes to
/// read values from previous nodes and set values for subsequent nodes.
#[derive(Clone)]
pub struct Context {
    id: String,
    data: Arc<RwLock<Value>>,
}

impl Context {
    /// Create a new Context from raw JSON bytes.
    pub fn new(data: &[u8]) -> Self {
        let value: Value = if data.is_empty() {
            json!({})
        } else {
            serde_json::from_slice(data).unwrap_or_else(|_| json!({}))
        };

        let id = value
            .get("id")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string();

        Self {
            id,
            data: Arc::new(RwLock::new(value)),
        }
    }

    /// Get the message ID.
    pub fn get_id(&self) -> &str {
        &self.id
    }

    /// Set a value at the given JSON path.
    pub fn set(&self, path: &str, value: impl Into<Value>) -> crate::runtime::Result<()> {
        let path = convert_path(path);
        let mut data = self.data.write();
        set_nested_value(&mut data, &path, value.into());
        Ok(())
    }

    /// Get a value at the given JSON path.
    pub fn get(&self, path: &str) -> Option<Value> {
        let path = convert_path(path);
        let data = self.data.read();
        get_nested_value(&data, &path)
    }

    /// Get a string value at the given JSON path.
    pub fn get_string(&self, path: &str) -> String {
        self.get(path)
            .and_then(|v| v.as_str().map(|s| s.to_string()))
            .unwrap_or_default()
    }

    /// Get a boolean value at the given JSON path.
    pub fn get_bool(&self, path: &str) -> bool {
        self.get(path)
            .and_then(|v| v.as_bool())
            .unwrap_or(false)
    }

    /// Get an integer value at the given JSON path.
    pub fn get_int(&self, path: &str) -> i64 {
        self.get(path)
            .and_then(|v| v.as_i64())
            .unwrap_or(0)
    }

    /// Get a float value at the given JSON path.
    pub fn get_float(&self, path: &str) -> f64 {
        self.get(path)
            .and_then(|v| v.as_f64())
            .unwrap_or(0.0)
    }

    /// Get the raw JSON bytes.
    pub fn get_raw(&self) -> crate::runtime::Result<Vec<u8>> {
        let data = self.data.read();
        Ok(serde_json::to_vec(&*data)?)
    }

    /// Get the raw JSON bytes with options applied.
    pub fn get_raw_with_options(&self, options: &[GetOption]) -> crate::runtime::Result<Vec<u8>> {
        let mut raw = self.get_raw()?;
        for opt in options {
            raw = opt(raw)?;
        }
        Ok(raw)
    }

    /// Set the raw JSON bytes.
    pub fn set_raw(&self, data: Option<Vec<u8>>) -> crate::runtime::Result<()> {
        match data {
            Some(bytes) => {
                let value: Value = serde_json::from_slice(&bytes)?;
                *self.data.write() = value;
            }
            None => {
                *self.data.write() = Value::Null;
            }
        }
        Ok(())
    }

    /// Set the raw JSON bytes with options applied.
    pub fn set_raw_with_options(
        &self,
        data: Vec<u8>,
        options: &[SetOption],
    ) -> crate::runtime::Result<()> {
        let mut raw = data;
        for opt in options {
            raw = opt(raw)?;
        }
        self.set_raw(Some(raw))
    }

    /// Check if the context is empty.
    pub fn is_empty(&self) -> bool {
        let data = self.data.read();
        data.is_null() || (data.is_object() && data.as_object().map_or(true, |o| o.is_empty()))
    }
}

/// Convert a path with brackets to dot notation.
/// e.g., "foo[0].bar" -> "foo.0.bar"
fn convert_path(path: &str) -> String {
    path.replace('[', ".").replace(']', "")
}

/// Get a nested value from a JSON value using dot notation.
fn get_nested_value(value: &Value, path: &str) -> Option<Value> {
    let parts: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();
    let mut current = value;

    for part in parts {
        if let Ok(index) = part.parse::<usize>() {
            current = current.get(index)?;
        } else {
            current = current.get(part)?;
        }
    }

    Some(current.clone())
}

/// Set a nested value in a JSON value using dot notation.
fn set_nested_value(value: &mut Value, path: &str, new_value: Value) {
    let parts: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();

    if parts.is_empty() {
        return;
    }

    // Handle single part case
    if parts.len() == 1 {
        let part = parts[0];
        if let Ok(index) = part.parse::<usize>() {
            if !value.is_array() {
                *value = json!([]);
            }
            if let Some(arr) = value.as_array_mut() {
                while arr.len() <= index {
                    arr.push(Value::Null);
                }
                arr[index] = new_value;
            }
        } else {
            if !value.is_object() {
                *value = json!({});
            }
            value[part] = new_value;
        }
        return;
    }

    // Multi-part path: navigate and set recursively
    let first = parts[0];
    let rest = parts[1..].join(".");

    if let Ok(index) = first.parse::<usize>() {
        if !value.is_array() {
            *value = json!([]);
        }
        if let Some(arr) = value.as_array_mut() {
            while arr.len() <= index {
                arr.push(json!({}));
            }
            set_nested_value(&mut arr[index], &rest, new_value);
        }
    } else {
        if !value.is_object() {
            *value = json!({});
        }
        if value.get(first).is_none() {
            value[first] = json!({});
        }
        if let Some(child) = value.get_mut(first) {
            set_nested_value(child, &rest, new_value);
        }
    }
}

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

    #[test]
    fn test_context_new() {
        let ctx = Context::new(b"{}");
        assert!(ctx.is_empty());
    }

    #[test]
    fn test_context_set_get() {
        let ctx = Context::new(b"{}");
        ctx.set("foo", "bar").unwrap();
        assert_eq!(ctx.get_string("foo"), "bar");
    }

    #[test]
    fn test_context_nested() {
        let ctx = Context::new(b"{}");
        ctx.set("foo.bar.baz", 42).unwrap();
        assert_eq!(ctx.get_int("foo.bar.baz"), 42);
    }

    #[test]
    fn test_context_array() {
        let ctx = Context::new(b"{}");
        ctx.set("items.0", "first").unwrap();
        ctx.set("items.1", "second").unwrap();
        assert_eq!(ctx.get_string("items.0"), "first");
        assert_eq!(ctx.get_string("items.1"), "second");
    }
}