enki-runtime 0.1.4

//! Mesh configuration from TOML files.

use super::{AgentConfig, MemoryConfig};
use crate::core::error::{Error, Result};
use serde::{Deserialize, Serialize};
use std::path::Path;

/// Configuration for a multi-agent mesh.
///
/// This struct allows defining multiple agents in a single TOML file.
///
/// # Example
///
/// ```toml
/// name = "research_mesh"
///
/// [memory]
/// backend = "sqlite"
/// path = "./data/knowledge.db"
///
/// [[agents]]
/// name = "researcher"
/// model = "ollama::gemma3:latest"
/// system_prompt = "You are a research assistant."
/// temperature = 0.7
///
/// [[agents]]
/// name = "summarizer"
/// model = "ollama::gemma3:latest"
/// system_prompt = "You are a summarization expert."
/// temperature = 0.5
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MeshConfig {
    /// Name of the mesh
    pub name: String,

    /// Version of the mesh configuration
    #[serde(default)]
    pub version: Option<String>,

    /// Name of the agent to use as coordinator (must exist in agents list)
    /// Takes precedence over auto_coordinator if both are set
    #[serde(default)]
    pub coordinator: Option<String>,

    /// If true, mesh will auto-generate a coordinator agent with orchestration capabilities
    /// Only used if coordinator is not explicitly set
    #[serde(default)]
    pub auto_coordinator: bool,

    /// Default model for auto-generated coordinator (defaults to first agent's model)
    #[serde(default)]
    pub coordinator_model: Option<String>,

    /// Shared memory configuration for all agents (can be overridden per-agent)
    #[serde(default)]
    pub memory: Option<MemoryConfig>,

    /// List of agent configurations
    #[serde(default)]
    pub agents: Vec<AgentConfig>,
}

impl Default for MeshConfig {
    fn default() -> Self {
        Self {
            name: "mesh".to_string(),
            version: None,
            coordinator: None,
            auto_coordinator: false,
            coordinator_model: None,
            memory: None,
            agents: Vec::new(),
        }
    }
}

impl MeshConfig {
    /// Create a new MeshConfig with a name.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            version: None,
            coordinator: None,
            auto_coordinator: false,
            coordinator_model: None,
            memory: None,
            agents: Vec::new(),
        }
    }

    /// Parse a MeshConfig from a TOML string.
    ///
    /// # Example
    ///
    /// ```rust
    /// use enki_runtime::config::MeshConfig;
    ///
    /// let toml = r#"
    /// name = "my_mesh"
    ///
    /// [[agents]]
    /// name = "agent1"
    /// model = "ollama::llama2"
    /// "#;
    ///
    /// let config = MeshConfig::from_toml(toml).unwrap();
    /// assert_eq!(config.name, "my_mesh");
    /// assert_eq!(config.agents.len(), 1);
    /// ```
    pub fn from_toml(toml_str: &str) -> Result<Self> {
        toml::from_str(toml_str)
            .map_err(|e| Error::ConfigError(format!("Failed to parse TOML: {}", e)))
    }

    /// Load a MeshConfig from a TOML file.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use enki_runtime::config::MeshConfig;
    ///
    /// let config = MeshConfig::from_file("agents.toml").unwrap();
    /// ```
    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
        let content = std::fs::read_to_string(path.as_ref()).map_err(|e| {
            Error::ConfigError(format!(
                "Failed to read file '{}': {}",
                path.as_ref().display(),
                e
            ))
        })?;
        Self::from_toml(&content)
    }

    /// Add an agent configuration.
    pub fn add_agent(mut self, agent: AgentConfig) -> Self {
        self.agents.push(agent);
        self
    }

    /// Get agent config by name.
    pub fn get_agent(&self, name: &str) -> Option<&AgentConfig> {
        self.agents.iter().find(|a| a.name == name)
    }

    /// Get the coordinator agent configuration.
    ///
    /// Priority order:
    /// 1. Explicitly set `coordinator` field (must exist in agents list)
    /// 2. First agent in the list (default behavior)
    ///
    /// Returns None if no agents are defined.
    pub fn get_coordinator(&self) -> Option<&AgentConfig> {
        // If coordinator is explicitly set, find it
        if let Some(ref coordinator_name) = self.coordinator {
            return self.get_agent(coordinator_name);
        }

        // Check if there's an auto-generated coordinator
        if self.auto_coordinator {
            if let Some(coord) = self.get_agent("coordinator") {
                return Some(coord);
            }
        }

        // Fallback to first agent
        self.agents.first()
    }

    /// Get the name of the coordinator agent.
    pub fn get_coordinator_name(&self) -> Option<&str> {
        self.get_coordinator().map(|a| a.name.as_str())
    }

    /// Ensure a coordinator agent exists if auto_coordinator is enabled.
    ///
    /// This creates a built-in coordinator agent with orchestration capabilities
    /// that knows about all other agents in the mesh.
    ///
    /// Call this after loading the config and before using the mesh.
    pub fn ensure_coordinator(&mut self) {
        // Skip if coordinator is explicitly set
        if self.coordinator.is_some() {
            return;
        }

        // Skip if auto_coordinator is not enabled
        if !self.auto_coordinator {
            return;
        }

        // Skip if coordinator already exists
        if self.get_agent("coordinator").is_some() {
            return;
        }

        // Determine the model to use for coordinator
        let model = self
            .coordinator_model
            .clone()
            .or_else(|| self.agents.first().map(|a| a.model.clone()))
            .unwrap_or_else(|| "ollama::gemma3:latest".to_string());

        // Build list of available agents for the system prompt
        let agent_list: Vec<String> = self
            .agents
            .iter()
            .map(|a| {
                let desc = if a.system_prompt.len() > 100 {
                    format!("{}...", &a.system_prompt[..100])
                } else {
                    a.system_prompt.clone()
                };
                format!("- {}: {}", a.name, desc)
            })
            .collect();

        let system_prompt = format!(
            r#"You are a coordinator agent responsible for orchestrating tasks across multiple specialized agents.

Your role is to:
1. Understand user requests and break them down into subtasks
2. Delegate subtasks to the most appropriate specialized agent
3. Synthesize responses from multiple agents when needed
4. Provide cohesive, helpful responses to the user

Available agents in this mesh:
{}

When delegating tasks, consider each agent's specialization and choose the best fit.
If a task requires multiple agents, coordinate their efforts and combine their outputs."#,
            agent_list.join("\n")
        );

        let coordinator = AgentConfig::new("coordinator", model)
            .with_system_prompt(system_prompt)
            .with_temperature(0.7);

        // Insert coordinator at the beginning
        self.agents.insert(0, coordinator);
    }

    /// Set the coordinator agent by name.
    pub fn with_coordinator(mut self, name: impl Into<String>) -> Self {
        self.coordinator = Some(name.into());
        self
    }

    /// Enable auto-coordinator with optional model.
    pub fn with_auto_coordinator(mut self, model: Option<String>) -> Self {
        self.auto_coordinator = true;
        self.coordinator_model = model;
        self
    }
}

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

    #[test]
    fn test_parse_mesh_config() {
        let toml = r#"
            name = "test_mesh"

            [[agents]]
            name = "agent1"
            model = "ollama::llama2"
            system_prompt = "First agent"

            [[agents]]
            name = "agent2"
            model = "openai::gpt-4"
            temperature = 0.8
        "#;

        let config = MeshConfig::from_toml(toml).unwrap();
        assert_eq!(config.name, "test_mesh");
        assert_eq!(config.agents.len(), 2);
        assert_eq!(config.agents[0].name, "agent1");
        assert_eq!(config.agents[1].name, "agent2");
        assert_eq!(config.agents[1].temperature, Some(0.8));
    }

    #[test]
    fn test_get_agent() {
        let config = MeshConfig::new("test")
            .add_agent(AgentConfig::new("agent1", "ollama::llama2"))
            .add_agent(AgentConfig::new("agent2", "ollama::llama2"));

        assert!(config.get_agent("agent1").is_some());
        assert!(config.get_agent("agent2").is_some());
        assert!(config.get_agent("nonexistent").is_none());
    }

    #[test]
    fn test_empty_agents() {
        let toml = r#"
            name = "empty_mesh"
        "#;

        let config = MeshConfig::from_toml(toml).unwrap();
        assert_eq!(config.name, "empty_mesh");
        assert!(config.agents.is_empty());
    }

    #[test]
    fn test_get_coordinator_default_first_agent() {
        let config = MeshConfig::new("test")
            .add_agent(AgentConfig::new("agent1", "ollama::llama2"))
            .add_agent(AgentConfig::new("agent2", "ollama::llama2"));

        // Default: first agent is coordinator
        let coordinator = config.get_coordinator();
        assert!(coordinator.is_some());
        assert_eq!(coordinator.unwrap().name, "agent1");
    }

    #[test]
    fn test_get_coordinator_designated() {
        let config = MeshConfig::new("test")
            .add_agent(AgentConfig::new("agent1", "ollama::llama2"))
            .add_agent(AgentConfig::new("agent2", "ollama::llama2"))
            .with_coordinator("agent2");

        // Designated coordinator takes precedence
        let coordinator = config.get_coordinator();
        assert!(coordinator.is_some());
        assert_eq!(coordinator.unwrap().name, "agent2");
    }

    #[test]
    fn test_ensure_coordinator_auto() {
        let mut config = MeshConfig::new("test")
            .add_agent(AgentConfig::new("worker1", "ollama::llama2"))
            .add_agent(AgentConfig::new("worker2", "ollama::llama2"))
            .with_auto_coordinator(None);

        // Before ensure_coordinator
        assert_eq!(config.agents.len(), 2);

        // After ensure_coordinator
        config.ensure_coordinator();
        assert_eq!(config.agents.len(), 3);

        // Coordinator should be first
        assert_eq!(config.agents[0].name, "coordinator");

        // get_coordinator should return the auto-generated coordinator
        let coordinator = config.get_coordinator();
        assert!(coordinator.is_some());
        assert_eq!(coordinator.unwrap().name, "coordinator");

        // System prompt should mention available agents
        assert!(config.agents[0].system_prompt.contains("worker1"));
        assert!(config.agents[0].system_prompt.contains("worker2"));
    }

    #[test]
    fn test_ensure_coordinator_skips_if_designated() {
        let mut config = MeshConfig::new("test")
            .add_agent(AgentConfig::new("agent1", "ollama::llama2"))
            .add_agent(AgentConfig::new("agent2", "ollama::llama2"))
            .with_coordinator("agent1")
            .with_auto_coordinator(None);

        config.auto_coordinator = true; // Also set auto_coordinator

        // Should not create auto coordinator when designated exists
        config.ensure_coordinator();
        assert_eq!(config.agents.len(), 2); // No new agent added
    }

    #[test]
    fn test_parse_coordinator_from_toml() {
        let toml = r#"
            name = "test_mesh"
            coordinator = "agent2"

            [[agents]]
            name = "agent1"
            model = "ollama::llama2"

            [[agents]]
            name = "agent2"
            model = "ollama::llama2"
        "#;

        let config = MeshConfig::from_toml(toml).unwrap();
        assert_eq!(config.coordinator, Some("agent2".to_string()));

        let coordinator = config.get_coordinator();
        assert_eq!(coordinator.unwrap().name, "agent2");
    }

    #[test]
    fn test_parse_auto_coordinator_from_toml() {
        let toml = r#"
            name = "test_mesh"
            auto_coordinator = true
            coordinator_model = "ollama::gemma3:latest"

            [[agents]]
            name = "worker1"
            model = "ollama::llama2"
        "#;

        let mut config = MeshConfig::from_toml(toml).unwrap();
        assert!(config.auto_coordinator);
        assert_eq!(
            config.coordinator_model,
            Some("ollama::gemma3:latest".to_string())
        );

        config.ensure_coordinator();
        assert_eq!(config.agents.len(), 2);
        assert_eq!(config.agents[0].name, "coordinator");
        assert_eq!(config.agents[0].model, "ollama::gemma3:latest");
    }
}