zclaw_openfang/crates/zclaw-skills/src/orchestration/auto_compose.rs

//! Auto-compose skills
//!
//! Automatically compose skills into execution graphs based on
//! input/output schema matching and semantic compatibility.

use std::collections::{HashMap, HashSet};
use serde_json::Value;
use zclaw_types::{Result, SkillId};

use crate::registry::SkillRegistry;
use crate::SkillManifest;
use super::{SkillGraph, SkillNode, SkillEdge};

/// Auto-composer for automatic skill graph generation
pub struct AutoComposer<'a> {
    registry: &'a SkillRegistry,
}

impl<'a> AutoComposer<'a> {
    pub fn new(registry: &'a SkillRegistry) -> Self {
        Self { registry }
    }

    /// Compose multiple skills into an execution graph
    pub async fn compose(&self, skill_ids: &[SkillId]) -> Result<SkillGraph> {
        // 1. Load all skill manifests
        let manifests = self.load_manifests(skill_ids).await?;

        // 2. Analyze input/output schemas
        let analysis = self.analyze_skills(&manifests);

        // 3. Build dependency graph based on schema matching
        let edges = self.infer_edges(&manifests, &analysis);

        // 4. Create the skill graph
        let graph = self.build_graph(skill_ids, &manifests, edges);

        Ok(graph)
    }

    /// Load manifests for all skills
    async fn load_manifests(&self, skill_ids: &[SkillId]) -> Result<Vec<SkillManifest>> {
        let mut manifests = Vec::new();
        for id in skill_ids {
            if let Some(manifest) = self.registry.get_manifest(id).await {
                manifests.push(manifest);
            } else {
                return Err(zclaw_types::ZclawError::NotFound(
                    format!("Skill not found: {}", id)
                ));
            }
        }
        Ok(manifests)
    }

    /// Analyze skills for compatibility
    fn analyze_skills(&self, manifests: &[SkillManifest]) -> SkillAnalysis {
        let mut analysis = SkillAnalysis::default();

        for manifest in manifests {
            // Extract output types from schema
            if let Some(schema) = &manifest.output_schema {
                let types = self.extract_types_from_schema(schema);
                analysis.output_types.insert(manifest.id.clone(), types);
            }

            // Extract input types from schema
            if let Some(schema) = &manifest.input_schema {
                let types = self.extract_types_from_schema(schema);
                analysis.input_types.insert(manifest.id.clone(), types);
            }

            // Extract capabilities
            analysis.capabilities.insert(
                manifest.id.clone(),
                manifest.capabilities.clone(),
            );
        }

        analysis
    }

    /// Extract type names from JSON schema
    fn extract_types_from_schema(&self, schema: &Value) -> HashSet<String> {
        let mut types = HashSet::new();

        if let Some(obj) = schema.as_object() {
            // Get type field
            if let Some(type_val) = obj.get("type") {
                if let Some(type_str) = type_val.as_str() {
                    types.insert(type_str.to_string());
                } else if let Some(type_arr) = type_val.as_array() {
                    for t in type_arr {
                        if let Some(s) = t.as_str() {
                            types.insert(s.to_string());
                        }
                    }
                }
            }

            // Get properties
            if let Some(props) = obj.get("properties") {
                if let Some(props_obj) = props.as_object() {
                    for (name, prop) in props_obj {
                        types.insert(name.clone());
                        if let Some(prop_obj) = prop.as_object() {
                            if let Some(type_str) = prop_obj.get("type").and_then(|t| t.as_str()) {
                                types.insert(format!("{}:{}", name, type_str));
                            }
                        }
                    }
                }
            }
        }

        types
    }

    /// Infer edges based on schema matching
    fn infer_edges(
        &self,
        manifests: &[SkillManifest],
        analysis: &SkillAnalysis,
    ) -> Vec<(String, String)> {
        let mut edges = Vec::new();
        let mut used_outputs: HashMap<String, HashSet<String>> = HashMap::new();

        // Try to match outputs to inputs
        for (i, source) in manifests.iter().enumerate() {
            let source_outputs = analysis.output_types.get(&source.id).cloned().unwrap_or_default();

            for (j, target) in manifests.iter().enumerate() {
                if i == j {
                    continue;
                }

                let target_inputs = analysis.input_types.get(&target.id).cloned().unwrap_or_default();

                // Check for matching types
                let matches: Vec<_> = source_outputs
                    .intersection(&target_inputs)
                    .filter(|t| !t.starts_with("object") && !t.starts_with("array"))
                    .collect();

                if !matches.is_empty() {
                    // Check if this output hasn't been used yet
                    let used = used_outputs.entry(source.id.to_string()).or_default();
                    let new_matches: Vec<_> = matches
                        .into_iter()
                        .filter(|m| !used.contains(*m))
                        .collect();

                    if !new_matches.is_empty() {
                        edges.push((source.id.to_string(), target.id.to_string()));
                        for m in new_matches {
                            used.insert(m.clone());
                        }
                    }
                }
            }
        }

        // If no edges found, create a linear chain
        if edges.is_empty() && manifests.len() > 1 {
            for i in 0..manifests.len() - 1 {
                edges.push((
                    manifests[i].id.to_string(),
                    manifests[i + 1].id.to_string(),
                ));
            }
        }

        edges
    }

    /// Build the final skill graph
    fn build_graph(
        &self,
        skill_ids: &[SkillId],
        manifests: &[SkillManifest],
        edges: Vec<(String, String)>,
    ) -> SkillGraph {
        let nodes: Vec<SkillNode> = manifests
            .iter()
            .map(|m| SkillNode {
                id: m.id.to_string(),
                skill_id: m.id.clone(),
                description: m.description.clone(),
                input_mappings: HashMap::new(),
                retry: None,
                timeout_secs: None,
                when: None,
                skip_on_error: false,
            })
            .collect();

        let edges: Vec<SkillEdge> = edges
            .into_iter()
            .map(|(from, to)| SkillEdge {
                from_node: from,
                to_node: to,
                field_mapping: HashMap::new(),
                condition: None,
            })
            .collect();

        let graph_id = format!("auto-{}", uuid::Uuid::new_v4());

        SkillGraph {
            id: graph_id,
            name: format!("Auto-composed: {}", skill_ids.iter()
                .map(|id| id.to_string())
                .collect::<Vec<_>>()
                .join(" → ")),
            description: format!("Automatically composed from skills: {}",
                skill_ids.iter()
                    .map(|id| id.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")),
            nodes,
            edges,
            input_schema: None,
            output_mapping: HashMap::new(),
            on_error: Default::default(),
            timeout_secs: 300,
        }
    }

    /// Suggest skills that can be composed with a given skill
    pub async fn suggest_compatible_skills(
        &self,
        skill_id: &SkillId,
    ) -> Result<Vec<(SkillId, CompatibilityScore)>> {
        let manifest = self.registry.get_manifest(skill_id).await
            .ok_or_else(|| zclaw_types::ZclawError::NotFound(
                format!("Skill not found: {}", skill_id)
            ))?;

        let all_skills = self.registry.list().await;
        let mut suggestions = Vec::new();

        let output_types = manifest.output_schema
            .as_ref()
            .map(|s| self.extract_types_from_schema(s))
            .unwrap_or_default();

        for other in all_skills {
            if other.id == *skill_id {
                continue;
            }

            let input_types = other.input_schema
                .as_ref()
                .map(|s| self.extract_types_from_schema(s))
                .unwrap_or_default();

            // Calculate compatibility score
            let score = self.calculate_compatibility(&output_types, &input_types);

            if score > 0.0 {
                suggestions.push((other.id.clone(), CompatibilityScore {
                    skill_id: other.id.clone(),
                    score,
                    reason: format!("Output types match {} input types",
                        other.name),
                }));
            }
        }

        // Sort by score descending
        suggestions.sort_by(|a, b| b.1.score.partial_cmp(&a.1.score).unwrap());

        Ok(suggestions)
    }

    /// Calculate compatibility score between output and input types
    fn calculate_compatibility(
        &self,
        output_types: &HashSet<String>,
        input_types: &HashSet<String>,
    ) -> f32 {
        if output_types.is_empty() || input_types.is_empty() {
            return 0.0;
        }

        let intersection = output_types.intersection(input_types).count();
        let union = output_types.union(input_types).count();

        if union == 0 {
            0.0
        } else {
            intersection as f32 / union as f32
        }
    }
}

/// Skill analysis result
#[derive(Debug, Default)]
struct SkillAnalysis {
    /// Output types for each skill
    output_types: HashMap<SkillId, HashSet<String>>,
    /// Input types for each skill
    input_types: HashMap<SkillId, HashSet<String>>,
    /// Capabilities for each skill
    capabilities: HashMap<SkillId, Vec<String>>,
}

/// Compatibility score for skill composition
#[derive(Debug, Clone)]
pub struct CompatibilityScore {
    /// Skill ID
    pub skill_id: SkillId,
    /// Compatibility score (0.0 - 1.0)
    pub score: f32,
    /// Reason for the score
    pub reason: String,
}

/// Skill composition template
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CompositionTemplate {
    /// Template name
    pub name: String,
    /// Template description
    pub description: String,
    /// Skill slots to fill
    pub slots: Vec<CompositionSlot>,
    /// Fixed edges between slots
    pub edges: Vec<TemplateEdge>,
}

/// Slot in a composition template
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CompositionSlot {
    /// Slot identifier
    pub id: String,
    /// Required capabilities
    pub required_capabilities: Vec<String>,
    /// Expected input schema
    pub input_schema: Option<Value>,
    /// Expected output schema
    pub output_schema: Option<Value>,
}

/// Edge in a composition template
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TemplateEdge {
    /// Source slot
    pub from: String,
    /// Target slot
    pub to: String,
    /// Field mappings
    #[serde(default)]
    pub mapping: HashMap<String, String>,
}

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

    #[test]
    fn test_extract_types() {
        let registry: &'static SkillRegistry = Box::leak(Box::new(SkillRegistry::new()));
        let composer = AutoComposer {
            registry,
        };

        let schema = serde_json::json!({
            "type": "object",
            "properties": {
                "content": { "type": "string" },
                "count": { "type": "number" }
            }
        });

        let types = composer.extract_types_from_schema(&schema);
        assert!(types.contains("object"));
        assert!(types.contains("content"));
        assert!(types.contains("count"));
    }
}