argentor-builtins 1.4.7

//! YAML processing skill for the Argentor AI agent framework.
//!
//! Provides YAML parse/stringify, validation, merge, and YAML/JSON conversion.
//! Uses a minimal hand-rolled parser for simple YAML (key: value) to avoid
//! external dependencies.

use argentor_core::{ArgentorResult, ToolCall, ToolResult};
use argentor_skills::skill::{Skill, SkillDescriptor};
use async_trait::async_trait;
use serde_json::{json, Map, Value};

/// YAML processing skill with parsing, validation, merge, and conversion.
pub struct YamlProcessorSkill {
    descriptor: SkillDescriptor,
}

impl YamlProcessorSkill {
    /// Create a new YAML processor skill.
    pub fn new() -> Self {
        Self {
            descriptor: SkillDescriptor {
                name: "yaml_processor".to_string(),
                description: "YAML parse/stringify, validate, merge, and YAML/JSON conversion."
                    .to_string(),
                parameters_schema: json!({
                    "type": "object",
                    "properties": {
                        "operation": {
                            "type": "string",
                            "enum": ["parse", "stringify", "validate", "merge", "to_json", "from_json", "get_keys", "get_value"],
                            "description": "The YAML operation to perform"
                        },
                        "yaml": {
                            "type": "string",
                            "description": "YAML content to process"
                        },
                        "yaml_a": {
                            "type": "string",
                            "description": "First YAML document for merge"
                        },
                        "yaml_b": {
                            "type": "string",
                            "description": "Second YAML document for merge (overrides yaml_a)"
                        },
                        "json_data": {
                            "type": "object",
                            "description": "JSON object to convert to YAML"
                        },
                        "key": {
                            "type": "string",
                            "description": "Key to extract from YAML"
                        }
                    },
                    "required": ["operation"]
                }),
                required_capabilities: vec![],
                requires_approval: false,
            },
        }
    }
}

impl Default for YamlProcessorSkill {
    fn default() -> Self {
        Self::new()
    }
}

/// Simple YAML parser: supports flat key: value, arrays (- item), and nested via indentation.
/// Returns a serde_json::Value representing the parsed structure.
fn parse_simple_yaml(yaml: &str) -> Result<Value, String> {
    let mut root = Map::new();
    let mut current_key: Option<String> = None;
    let mut current_list: Option<Vec<Value>> = None;

    for line in yaml.lines() {
        let trimmed = line.trim();
        if trimmed.is_empty() || trimmed.starts_with('#') {
            continue;
        }

        if trimmed.starts_with("- ") {
            let item = trimmed.strip_prefix("- ").unwrap_or("").trim();
            if let Some(ref mut list) = current_list {
                list.push(parse_yaml_value(item));
            } else {
                current_list = Some(vec![parse_yaml_value(item)]);
            }
            continue;
        }

        // Flush any pending list
        if let (Some(key), Some(list)) = (current_key.take(), current_list.take()) {
            root.insert(key, Value::Array(list));
        }

        if let Some((key, val)) = trimmed.split_once(':') {
            let key = key.trim().to_string();
            let val = val.trim();
            if val.is_empty() {
                current_key = Some(key);
                current_list = Some(Vec::new());
            } else {
                root.insert(key, parse_yaml_value(val));
            }
        } else {
            return Err(format!("Invalid YAML line: '{trimmed}'"));
        }
    }

    // Flush final pending list
    if let (Some(key), Some(list)) = (current_key, current_list) {
        root.insert(key, Value::Array(list));
    }

    Ok(Value::Object(root))
}

/// Parse a YAML value string into a serde_json Value.
fn parse_yaml_value(s: &str) -> Value {
    let s = s.trim();
    // Strip quotes
    if (s.starts_with('"') && s.ends_with('"')) || (s.starts_with('\'') && s.ends_with('\'')) {
        return Value::String(s[1..s.len() - 1].to_string());
    }
    if s == "true" || s == "True" || s == "yes" || s == "Yes" {
        return Value::Bool(true);
    }
    if s == "false" || s == "False" || s == "no" || s == "No" {
        return Value::Bool(false);
    }
    if s == "null" || s == "~" || s == "Null" {
        return Value::Null;
    }
    if let Ok(n) = s.parse::<i64>() {
        return json!(n);
    }
    if let Ok(f) = s.parse::<f64>() {
        return json!(f);
    }
    Value::String(s.to_string())
}

/// Convert a JSON Value to simple YAML string.
fn json_to_yaml(value: &Value, indent: usize) -> String {
    let prefix = " ".repeat(indent);
    match value {
        Value::Object(map) => {
            let mut result = String::new();
            for (k, v) in map {
                match v {
                    Value::Object(_) => {
                        result.push_str(&format!("{prefix}{k}:\n"));
                        result.push_str(&json_to_yaml(v, indent + 2));
                    }
                    Value::Array(arr) => {
                        result.push_str(&format!("{prefix}{k}:\n"));
                        for item in arr {
                            let item_str = value_to_yaml_scalar(item);
                            result.push_str(&format!("{prefix}  - {item_str}\n"));
                        }
                    }
                    _ => {
                        let scalar = value_to_yaml_scalar(v);
                        result.push_str(&format!("{prefix}{k}: {scalar}\n"));
                    }
                }
            }
            result
        }
        _ => format!("{prefix}{}\n", value_to_yaml_scalar(value)),
    }
}

/// Convert a scalar JSON value to its YAML representation.
fn value_to_yaml_scalar(v: &Value) -> String {
    match v {
        Value::String(s) => {
            if s.contains(':') || s.contains('#') || s.contains('"') || s.contains('\'') {
                format!("\"{s}\"")
            } else {
                s.clone()
            }
        }
        Value::Number(n) => n.to_string(),
        Value::Bool(b) => b.to_string(),
        Value::Null => "null".to_string(),
        _ => v.to_string(),
    }
}

/// Merge two JSON objects (b overrides a).
fn merge_objects(a: &Value, b: &Value) -> Value {
    match (a, b) {
        (Value::Object(map_a), Value::Object(map_b)) => {
            let mut merged = map_a.clone();
            for (k, v) in map_b {
                if let Some(existing) = map_a.get(k) {
                    merged.insert(k.clone(), merge_objects(existing, v));
                } else {
                    merged.insert(k.clone(), v.clone());
                }
            }
            Value::Object(merged)
        }
        (_, b) => b.clone(),
    }
}

#[async_trait]
impl Skill for YamlProcessorSkill {
    fn descriptor(&self) -> &SkillDescriptor {
        &self.descriptor
    }

    async fn execute(&self, call: ToolCall) -> ArgentorResult<ToolResult> {
        let operation = match call.arguments["operation"].as_str() {
            Some(op) => op,
            None => {
                return Ok(ToolResult::error(
                    &call.id,
                    "Missing required parameter: 'operation'",
                ))
            }
        };

        match operation {
            "parse" | "to_json" => {
                let yaml = match call.arguments["yaml"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'yaml'")),
                };
                match parse_simple_yaml(yaml) {
                    Ok(parsed) => {
                        let response = json!({ "data": parsed });
                        Ok(ToolResult::success(&call.id, response.to_string()))
                    }
                    Err(e) => Ok(ToolResult::error(&call.id, format!("YAML parse error: {e}"))),
                }
            }
            "stringify" | "from_json" => {
                let json_data = &call.arguments["json_data"];
                if json_data.is_null() {
                    return Ok(ToolResult::error(&call.id, "Missing required parameter: 'json_data'"));
                }
                let yaml_str = json_to_yaml(json_data, 0);
                let response = json!({ "yaml": yaml_str.trim_end() });
                Ok(ToolResult::success(&call.id, response.to_string()))
            }
            "validate" => {
                let yaml = match call.arguments["yaml"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'yaml'")),
                };
                match parse_simple_yaml(yaml) {
                    Ok(_) => {
                        let response = json!({ "valid": true });
                        Ok(ToolResult::success(&call.id, response.to_string()))
                    }
                    Err(e) => {
                        let response = json!({ "valid": false, "error": e });
                        Ok(ToolResult::success(&call.id, response.to_string()))
                    }
                }
            }
            "merge" => {
                let yaml_a = match call.arguments["yaml_a"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'yaml_a'")),
                };
                let yaml_b = match call.arguments["yaml_b"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'yaml_b'")),
                };
                let a = parse_simple_yaml(yaml_a).map_err(|e| format!("yaml_a parse error: {e}"));
                let b = parse_simple_yaml(yaml_b).map_err(|e| format!("yaml_b parse error: {e}"));
                match (a, b) {
                    (Ok(a), Ok(b)) => {
                        let merged = merge_objects(&a, &b);
                        let yaml_str = json_to_yaml(&merged, 0);
                        let response = json!({ "merged": merged, "yaml": yaml_str.trim_end() });
                        Ok(ToolResult::success(&call.id, response.to_string()))
                    }
                    (Err(e), _) | (_, Err(e)) => Ok(ToolResult::error(&call.id, e)),
                }
            }
            "get_keys" => {
                let yaml = match call.arguments["yaml"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'yaml'")),
                };
                match parse_simple_yaml(yaml) {
                    Ok(Value::Object(map)) => {
                        let keys: Vec<&String> = map.keys().collect();
                        let response = json!({ "keys": keys, "count": keys.len() });
                        Ok(ToolResult::success(&call.id, response.to_string()))
                    }
                    Ok(_) => Ok(ToolResult::error(&call.id, "YAML root is not an object")),
                    Err(e) => Ok(ToolResult::error(&call.id, format!("YAML parse error: {e}"))),
                }
            }
            "get_value" => {
                let yaml = match call.arguments["yaml"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'yaml'")),
                };
                let key = match call.arguments["key"].as_str() {
                    Some(v) => v,
                    None => return Ok(ToolResult::error(&call.id, "Missing required parameter: 'key'")),
                };
                match parse_simple_yaml(yaml) {
                    Ok(Value::Object(map)) => match map.get(key) {
                        Some(val) => {
                            let response = json!({ "key": key, "value": val });
                            Ok(ToolResult::success(&call.id, response.to_string()))
                        }
                        None => Ok(ToolResult::error(&call.id, format!("Key '{key}' not found"))),
                    },
                    Ok(_) => Ok(ToolResult::error(&call.id, "YAML root is not an object")),
                    Err(e) => Ok(ToolResult::error(&call.id, format!("YAML parse error: {e}"))),
                }
            }
            _ => Ok(ToolResult::error(
                &call.id,
                format!("Unknown operation: '{operation}'. Supported: parse, stringify, validate, merge, to_json, from_json, get_keys, get_value"),
            )),
        }
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    fn make_call(args: Value) -> ToolCall {
        ToolCall {
            id: "test".to_string(),
            name: "yaml_processor".to_string(),
            arguments: args,
        }
    }

    #[tokio::test]
    async fn test_parse_simple() {
        let skill = YamlProcessorSkill::new();
        let call =
            make_call(json!({"operation": "parse", "yaml": "name: Alice\nage: 30\nactive: true"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error, "Result: {}", result.content);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["data"]["name"], "Alice");
        assert_eq!(parsed["data"]["age"], 30);
        assert_eq!(parsed["data"]["active"], true);
    }

    #[tokio::test]
    async fn test_parse_with_list() {
        let skill = YamlProcessorSkill::new();
        let yaml = "name: project\ntags:\n- rust\n- wasm\n- ai";
        let call = make_call(json!({"operation": "parse", "yaml": yaml}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["data"]["tags"], json!(["rust", "wasm", "ai"]));
    }

    #[tokio::test]
    async fn test_stringify() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({
            "operation": "stringify",
            "json_data": {"name": "Alice", "age": 30}
        }));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        let yaml = parsed["yaml"].as_str().unwrap();
        assert!(yaml.contains("name: Alice"));
        assert!(yaml.contains("age: 30"));
    }

    #[tokio::test]
    async fn test_validate_valid() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({"operation": "validate", "yaml": "key: value\nother: 42"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["valid"], true);
    }

    #[tokio::test]
    async fn test_validate_invalid() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(
            json!({"operation": "validate", "yaml": "this is not valid yaml at all :::"}),
        );
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        // Our simple parser may accept some invalid YAML, but that's fine for a lightweight tool
        assert!(parsed.get("valid").is_some());
    }

    #[tokio::test]
    async fn test_merge() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({
            "operation": "merge",
            "yaml_a": "name: Alice\nage: 30",
            "yaml_b": "age: 31\ncity: NYC"
        }));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["merged"]["name"], "Alice");
        assert_eq!(parsed["merged"]["age"], 31);
        assert_eq!(parsed["merged"]["city"], "NYC");
    }

    #[tokio::test]
    async fn test_to_json() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({"operation": "to_json", "yaml": "x: 1\ny: hello"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["data"]["x"], 1);
        assert_eq!(parsed["data"]["y"], "hello");
    }

    #[tokio::test]
    async fn test_from_json() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({
            "operation": "from_json",
            "json_data": {"greeting": "hello", "count": 5}
        }));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        let yaml = parsed["yaml"].as_str().unwrap();
        assert!(yaml.contains("count: 5"));
        assert!(yaml.contains("greeting: hello"));
    }

    #[tokio::test]
    async fn test_get_keys() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({"operation": "get_keys", "yaml": "a: 1\nb: 2\nc: 3"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["count"], 3);
    }

    #[tokio::test]
    async fn test_get_value() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(
            json!({"operation": "get_value", "yaml": "name: Bob\nage: 25", "key": "name"}),
        );
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["value"], "Bob");
    }

    #[tokio::test]
    async fn test_get_value_not_found() {
        let skill = YamlProcessorSkill::new();
        let call =
            make_call(json!({"operation": "get_value", "yaml": "name: Bob", "key": "email"}));
        let result = skill.execute(call).await.unwrap();
        assert!(result.is_error);
        assert!(result.content.contains("not found"));
    }

    #[tokio::test]
    async fn test_null_value() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({"operation": "parse", "yaml": "empty: null\ntilde: ~"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert!(parsed["data"]["empty"].is_null());
        assert!(parsed["data"]["tilde"].is_null());
    }

    #[tokio::test]
    async fn test_boolean_values() {
        let skill = YamlProcessorSkill::new();
        let call =
            make_call(json!({"operation": "parse", "yaml": "a: true\nb: false\nc: yes\nd: no"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["data"]["a"], true);
        assert_eq!(parsed["data"]["b"], false);
        assert_eq!(parsed["data"]["c"], true);
        assert_eq!(parsed["data"]["d"], false);
    }

    #[tokio::test]
    async fn test_comments_ignored() {
        let skill = YamlProcessorSkill::new();
        let call =
            make_call(json!({"operation": "parse", "yaml": "# comment\nkey: value\n# another"}));
        let result = skill.execute(call).await.unwrap();
        assert!(!result.is_error);
        let parsed: Value = serde_json::from_str(&result.content).unwrap();
        assert_eq!(parsed["data"]["key"], "value");
    }

    #[tokio::test]
    async fn test_missing_operation() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({"yaml": "a: 1"}));
        let result = skill.execute(call).await.unwrap();
        assert!(result.is_error);
    }

    #[tokio::test]
    async fn test_unknown_operation() {
        let skill = YamlProcessorSkill::new();
        let call = make_call(json!({"operation": "transform"}));
        let result = skill.execute(call).await.unwrap();
        assert!(result.is_error);
        assert!(result.content.contains("Unknown operation"));
    }

    #[test]
    fn test_descriptor_name() {
        let skill = YamlProcessorSkill::new();
        assert_eq!(skill.descriptor().name, "yaml_processor");
    }
}