enki-runtime 0.1.4

//! Agent configuration from TOML files.

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

/// Configuration for a single LLM agent.
///
/// This struct can be loaded from a TOML file and used to create an `LlmAgent`.
///
/// # Example
///
/// ```toml
/// name = "my_agent"
/// model = "ollama::gemma3:latest"
/// system_prompt = "You are a helpful assistant."
/// temperature = 0.7
/// max_tokens = 1024
/// ```
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentConfig {
    /// Unique name for the agent
    pub name: String,

    /// Model in "provider::model" format (e.g., "ollama::gemma3:latest")
    pub model: String,

    /// System prompt for the agent
    #[serde(default = "default_system_prompt")]
    pub system_prompt: String,

    /// API key for the LLM provider (optional, can use env vars)
    #[serde(default)]
    pub api_key: Option<String>,

    /// Temperature for generation (0.0 - 2.0)
    #[serde(default)]
    pub temperature: Option<f32>,

    /// Maximum tokens to generate
    #[serde(default)]
    pub max_tokens: Option<u32>,

    /// Base URL override for the LLM provider
    #[serde(default)]
    pub base_url: Option<String>,

    /// Top-p (nucleus) sampling parameter
    #[serde(default)]
    pub top_p: Option<f32>,

    /// Top-k sampling parameter
    #[serde(default)]
    pub top_k: Option<u32>,

    /// Request timeout in seconds
    #[serde(default)]
    pub timeout_seconds: Option<u64>,

    /// Enable reasoning mode (for supported providers)
    #[serde(default)]
    pub reasoning: Option<bool>,

    /// Reasoning effort level
    #[serde(default)]
    pub reasoning_effort: Option<String>,

    /// Memory configuration for this agent (optional, overrides mesh default)
    #[serde(default)]
    pub memory: Option<MemoryConfig>,

    /// MCP tool configurations for this agent.
    /// Each entry defines an external MCP server that provides tools.
    #[serde(default)]
    pub mcp_tools: Vec<McpToolConfig>,
}

fn default_system_prompt() -> String {
    "You are a helpful AI assistant.".to_string()
}

impl Default for AgentConfig {
    fn default() -> Self {
        Self {
            name: String::new(),
            model: String::new(),
            system_prompt: default_system_prompt(),
            api_key: None,
            temperature: None,
            max_tokens: None,
            base_url: None,
            top_p: None,
            top_k: None,
            timeout_seconds: None,
            reasoning: None,
            reasoning_effort: None,
            memory: None,
            mcp_tools: Vec::new(),
        }
    }
}

impl AgentConfig {
    /// Create a new AgentConfig with required fields.
    pub fn new(name: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            model: model.into(),
            ..Default::default()
        }
    }

    /// Parse an AgentConfig from a TOML string.
    ///
    /// # Example
    ///
    /// ```rust
    /// use enki_runtime::config::AgentConfig;
    ///
    /// let toml = r#"
    /// name = "my_agent"
    /// model = "ollama::llama2"
    /// "#;
    ///
    /// let config = AgentConfig::from_toml(toml).unwrap();
    /// assert_eq!(config.name, "my_agent");
    /// ```
    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 an AgentConfig from a TOML file.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use enki_runtime::config::AgentConfig;
    ///
    /// let config = AgentConfig::from_file("agent.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)
    }

    /// Set the system prompt.
    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = prompt.into();
        self
    }

    /// Set the API key.
    pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
        self.api_key = Some(api_key.into());
        self
    }

    /// Set the temperature.
    pub fn with_temperature(mut self, temperature: f32) -> Self {
        self.temperature = Some(temperature);
        self
    }

    /// Set max tokens.
    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    /// Set memory configuration for this agent.
    pub fn with_memory(mut self, memory: MemoryConfig) -> Self {
        self.memory = Some(memory);
        self
    }

    /// Get the effective memory configuration for this agent.
    ///
    /// Returns the agent's own memory config if set, otherwise returns the fallback
    /// (typically the mesh-level memory config).
    ///
    /// # Arguments
    ///
    /// * `fallback` - The mesh-level memory config to use if agent has no memory config
    ///
    /// # Example
    ///
    /// ```rust
    /// use enki_runtime::config::{AgentConfig, MemoryConfig, MeshConfig};
    ///
    /// let mesh_memory = MemoryConfig::in_memory().with_max_entries(1000);
    /// let agent = AgentConfig::new("agent1", "ollama::llama2");
    ///
    /// // Agent without memory uses mesh fallback
    /// let effective = agent.get_effective_memory_config(Some(&mesh_memory));
    /// assert!(effective.is_some());
    ///
    /// // Agent with its own memory uses that
    /// let agent_with_memory = agent.with_memory(MemoryConfig::in_memory().with_max_entries(500));
    /// let effective = agent_with_memory.get_effective_memory_config(Some(&mesh_memory));
    /// assert_eq!(effective.unwrap().max_entries, Some(500));
    /// ```
    pub fn get_effective_memory_config<'a>(
        &'a self,
        fallback: Option<&'a MemoryConfig>,
    ) -> Option<&'a MemoryConfig> {
        self.memory.as_ref().or(fallback)
    }
}

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

    #[test]
    fn test_parse_basic_config() {
        let toml = r#"
            name = "test_agent"
            model = "ollama::llama2"
        "#;

        let config = AgentConfig::from_toml(toml).unwrap();
        assert_eq!(config.name, "test_agent");
        assert_eq!(config.model, "ollama::llama2");
        assert_eq!(config.system_prompt, "You are a helpful AI assistant.");
    }

    #[test]
    fn test_parse_full_config() {
        let toml = r#"
            name = "researcher"
            model = "openai::gpt-4"
            system_prompt = "You are a research assistant."
            api_key = "sk-test"
            temperature = 0.7
            max_tokens = 2048
        "#;

        let config = AgentConfig::from_toml(toml).unwrap();
        assert_eq!(config.name, "researcher");
        assert_eq!(config.model, "openai::gpt-4");
        assert_eq!(config.system_prompt, "You are a research assistant.");
        assert_eq!(config.api_key, Some("sk-test".to_string()));
        assert_eq!(config.temperature, Some(0.7));
        assert_eq!(config.max_tokens, Some(2048));
    }

    #[test]
    fn test_invalid_toml() {
        let toml = "this is not valid toml [[[";
        let result = AgentConfig::from_toml(toml);
        assert!(result.is_err());
    }

    #[test]
    fn test_builder_methods() {
        let config = AgentConfig::new("agent", "ollama::llama2")
            .with_system_prompt("Custom prompt")
            .with_temperature(0.5)
            .with_max_tokens(1024);

        assert_eq!(config.name, "agent");
        assert_eq!(config.system_prompt, "Custom prompt");
        assert_eq!(config.temperature, Some(0.5));
        assert_eq!(config.max_tokens, Some(1024));
    }

    #[test]
    fn test_with_memory() {
        let memory = MemoryConfig::in_memory().with_max_entries(500);
        let config = AgentConfig::new("agent", "ollama::llama2").with_memory(memory);

        assert!(config.memory.is_some());
        assert_eq!(config.memory.as_ref().unwrap().max_entries, Some(500));
    }

    #[test]
    fn test_get_effective_memory_config_with_agent_memory() {
        let mesh_memory = MemoryConfig::in_memory().with_max_entries(1000);
        let agent_memory = MemoryConfig::in_memory().with_max_entries(500);
        let config = AgentConfig::new("agent", "ollama::llama2").with_memory(agent_memory);

        // Agent's own memory takes precedence
        let effective = config.get_effective_memory_config(Some(&mesh_memory));
        assert!(effective.is_some());
        assert_eq!(effective.unwrap().max_entries, Some(500));
    }

    #[test]
    fn test_get_effective_memory_config_fallback_to_mesh() {
        let mesh_memory = MemoryConfig::in_memory().with_max_entries(1000);
        let config = AgentConfig::new("agent", "ollama::llama2");

        // Falls back to mesh memory when agent has no memory config
        let effective = config.get_effective_memory_config(Some(&mesh_memory));
        assert!(effective.is_some());
        assert_eq!(effective.unwrap().max_entries, Some(1000));
    }

    #[test]
    fn test_get_effective_memory_config_no_memory() {
        let config = AgentConfig::new("agent", "ollama::llama2");

        // Returns None when no memory config anywhere
        let effective = config.get_effective_memory_config(None);
        assert!(effective.is_none());
    }
}