bzzz-core 0.1.0

//! Agent Spec v1
//!
//! Agent Spec defines how a single Agent is exposed to Bzzz as a callable capability.
//! It defines only the external contract Bzzz needs:
//! - identity
//! - runtime mode
//! - input / output contract
//! - execution constraints
//!
//! It does NOT define:
//! - the Agent's internal prompt
//! - model choice
//! - tool implementation details

use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;

use serde::{Deserialize, Serialize};

use crate::{RunError, RuntimeKind};

/// Unique identifier for an Agent Spec
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AgentSpecId(pub String);

impl AgentSpecId {
    /// Create a new AgentSpecId
    pub fn new(name: impl Into<String>) -> Self {
        AgentSpecId(name.into())
    }

    /// Get the inner string
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// Version of the Agent Spec format
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum ApiVersion {
    /// Version 1
    #[serde(rename = "v1")]
    #[default]
    V1,
}

/// Runtime configuration for an Agent
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RuntimeConfig {
    /// Runtime kind (local, docker, http, cloud)
    pub kind: RuntimeKind,
    /// Optional runtime-specific configuration
    #[serde(default)]
    pub config: HashMap<String, String>,
}

impl RuntimeConfig {
    /// Create a new runtime config
    pub fn new(kind: RuntimeKind) -> Self {
        RuntimeConfig {
            kind,
            config: HashMap::new(),
        }
    }

    /// Add a configuration key-value pair
    pub fn with_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.config.insert(key.into(), value.into());
        self
    }
}

/// Input schema for an Agent
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InputSchema {
    /// JSON Schema for input validation
    pub schema: serde_json::Value,
    /// Whether input is required
    #[serde(default = "default_true")]
    pub required: bool,
}

fn default_true() -> bool {
    true
}

impl InputSchema {
    /// Create a new input schema
    pub fn new(schema: serde_json::Value) -> Self {
        InputSchema {
            schema,
            required: true,
        }
    }

    /// Mark input as optional
    pub fn optional(mut self) -> Self {
        self.required = false;
        self
    }
}

/// Output schema for an Agent
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct OutputSchema {
    /// JSON Schema for output validation
    pub schema: serde_json::Value,
}

impl OutputSchema {
    /// Create a new output schema
    pub fn new(schema: serde_json::Value) -> Self {
        OutputSchema { schema }
    }
}

/// Tool definition exposed by an Agent
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolDef {
    /// Tool name
    pub name: String,
    /// Tool description
    #[serde(default)]
    pub description: String,
    /// Input schema for the tool
    pub input_schema: InputSchema,
}

impl ToolDef {
    /// Create a new tool definition
    pub fn new(
        name: impl Into<String>,
        description: impl Into<String>,
        input_schema: InputSchema,
    ) -> Self {
        ToolDef {
            name: name.into(),
            description: description.into(),
            input_schema,
        }
    }
}

/// Interface definition for an Agent
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Interface {
    /// Input schema
    #[serde(default)]
    pub input: Option<InputSchema>,
    /// Output schema
    #[serde(default)]
    pub output: Option<OutputSchema>,
    /// Tools exposed by the Agent
    #[serde(default)]
    pub tools: Vec<ToolDef>,
}

impl Interface {
    /// Create a new interface
    pub fn new() -> Self {
        Interface {
            input: None,
            output: None,
            tools: Vec::new(),
        }
    }

    /// Set input schema
    pub fn with_input(mut self, input: InputSchema) -> Self {
        self.input = Some(input);
        self
    }

    /// Set output schema
    pub fn with_output(mut self, output: OutputSchema) -> Self {
        self.output = Some(output);
        self
    }

    /// Add a tool
    pub fn with_tool(mut self, tool: ToolDef) -> Self {
        self.tools.push(tool);
        self
    }
}

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

/// Execution constraints for an Agent
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Constraints {
    /// Maximum execution time
    #[serde(default, with = "serde_duration_opt")]
    pub timeout: Option<Duration>,
    /// Maximum memory in bytes
    #[serde(default)]
    pub max_memory_bytes: Option<u64>,
    /// Maximum retries on failure
    #[serde(default)]
    pub max_retries: u32,
    /// Retry delay
    #[serde(default, with = "serde_duration_opt")]
    pub retry_delay: Option<Duration>,
}

mod serde_duration_opt {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::time::Duration;

    pub fn serialize<S>(value: &Option<Duration>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        match value {
            Some(d) => d.as_secs().serialize(serializer),
            None => serializer.serialize_none(),
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<Duration>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let opt = Option::<u64>::deserialize(deserializer)?;
        Ok(opt.map(Duration::from_secs))
    }
}

impl Constraints {
    /// Create new constraints with defaults
    pub fn new() -> Self {
        Constraints {
            timeout: None,
            max_memory_bytes: None,
            max_retries: 0,
            retry_delay: None,
        }
    }

    /// Set timeout duration
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    /// Set maximum memory limit
    pub fn with_max_memory(mut self, bytes: u64) -> Self {
        self.max_memory_bytes = Some(bytes);
        self
    }

    /// Set retry configuration
    pub fn with_retries(mut self, max_retries: u32, delay: Duration) -> Self {
        self.max_retries = max_retries;
        self.retry_delay = Some(delay);
        self
    }
}

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

/// Agent Spec v1
///
/// Defines how a single Agent is exposed to Bzzz as a callable capability.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentSpec {
    /// API version
    #[serde(rename = "apiVersion")]
    pub api_version: ApiVersion,
    /// Unique identifier
    pub id: AgentSpecId,
    /// Runtime configuration
    pub runtime: RuntimeConfig,
    /// Interface definition
    #[serde(default)]
    pub interface: Interface,
    /// Execution constraints
    #[serde(default)]
    pub constraints: Constraints,
    /// Path to the spec file (set when loaded)
    #[serde(skip)]
    pub spec_path: Option<PathBuf>,
}

impl AgentSpec {
    /// Create a new Agent Spec with the given ID and runtime
    pub fn new(id: impl Into<String>, runtime: RuntimeKind) -> Self {
        AgentSpec {
            api_version: ApiVersion::V1,
            id: AgentSpecId::new(id),
            runtime: RuntimeConfig::new(runtime),
            interface: Interface::new(),
            constraints: Constraints::new(),
            spec_path: None,
        }
    }

    /// Set the interface
    pub fn with_interface(mut self, interface: Interface) -> Self {
        self.interface = interface;
        self
    }

    /// Set the constraints
    pub fn with_constraints(mut self, constraints: Constraints) -> Self {
        self.constraints = constraints;
        self
    }

    /// Validate the Agent Spec
    pub fn validate(&self) -> Result<(), RunError> {
        // Validate ID is not empty
        if self.id.0.is_empty() {
            return Err(RunError::InvalidConfig {
                message: "Agent Spec ID cannot be empty".into(),
            });
        }

        // Validate tools have names
        for tool in &self.interface.tools {
            if tool.name.is_empty() {
                return Err(RunError::InvalidConfig {
                    message: "Tool name cannot be empty".into(),
                });
            }
        }

        Ok(())
    }

    /// Load from a YAML file
    pub fn from_yaml_file(path: &PathBuf) -> Result<Self, RunError> {
        let content = std::fs::read_to_string(path).map_err(|e| RunError::InvalidConfig {
            message: format!("Failed to read {}: {}", path.display(), e),
        })?;

        let mut spec: AgentSpec =
            serde_yaml::from_str(&content).map_err(|e| RunError::InvalidConfig {
                message: format!("Failed to parse YAML: {}", e),
            })?;

        spec.spec_path = Some(path.clone());
        spec.validate()?;
        Ok(spec)
    }

    /// Save to a YAML file
    pub fn to_yaml_file(&self, path: &PathBuf) -> Result<(), RunError> {
        let content = serde_yaml::to_string(self).map_err(|e| RunError::InvalidConfig {
            message: format!("Failed to serialize: {}", e),
        })?;

        std::fs::write(path, content).map_err(|e| RunError::InvalidConfig {
            message: format!("Failed to write {}: {}", path.display(), e),
        })?;

        Ok(())
    }
}

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

    #[test]
    fn test_agent_spec_creation() {
        let spec = AgentSpec::new("my-agent", RuntimeKind::Local);
        assert_eq!(spec.id.as_str(), "my-agent");
        assert_eq!(spec.api_version, ApiVersion::V1);
        assert_eq!(spec.runtime.kind, RuntimeKind::Local);
    }

    #[test]
    fn test_agent_spec_validation() {
        let spec = AgentSpec::new("test", RuntimeKind::Local);
        assert!(spec.validate().is_ok());

        let invalid_spec = AgentSpec::new("", RuntimeKind::Local);
        assert!(invalid_spec.validate().is_err());
    }

    #[test]
    fn test_runtime_config() {
        let config =
            RuntimeConfig::new(RuntimeKind::Docker).with_config("image", "my-image:latest");

        assert_eq!(config.kind, RuntimeKind::Docker);
        assert_eq!(
            config.config.get("image"),
            Some(&"my-image:latest".to_string())
        );
    }

    #[test]
    fn test_interface() {
        let interface = Interface::new()
            .with_input(InputSchema::new(serde_json::json!({"type": "object"})))
            .with_output(OutputSchema::new(serde_json::json!({"type": "string"})))
            .with_tool(ToolDef::new(
                "my-tool",
                "A tool",
                InputSchema::new(serde_json::json!({})),
            ));

        assert!(interface.input.is_some());
        assert!(interface.output.is_some());
        assert_eq!(interface.tools.len(), 1);
    }

    #[test]
    fn test_constraints() {
        let constraints = Constraints::new()
            .with_timeout(Duration::from_secs(60))
            .with_max_memory(1024 * 1024 * 512) // 512MB
            .with_retries(3, Duration::from_secs(1));

        assert_eq!(constraints.timeout, Some(Duration::from_secs(60)));
        assert_eq!(constraints.max_memory_bytes, Some(1024 * 1024 * 512));
        assert_eq!(constraints.max_retries, 3);
    }

    #[test]
    fn test_yaml_roundtrip() {
        let spec = AgentSpec::new("test-agent", RuntimeKind::Local)
            .with_constraints(Constraints::new().with_timeout(Duration::from_secs(30)));

        let yaml = serde_yaml::to_string(&spec).unwrap();
        let parsed: AgentSpec = serde_yaml::from_str(&yaml).unwrap();

        assert_eq!(parsed.id, spec.id);
        assert_eq!(parsed.runtime.kind, spec.runtime.kind);
        assert_eq!(parsed.constraints.timeout, spec.constraints.timeout);
    }

    #[test]
    fn test_tool_validation() {
        let spec =
            AgentSpec::new("test", RuntimeKind::Local).with_interface(Interface::new().with_tool(
                ToolDef::new("", "empty name", InputSchema::new(serde_json::json!({}))),
            ));

        assert!(spec.validate().is_err());
    }
}