coderlib 0.1.0

//! Enhanced tool router system for CoderLib
//!
//! This module provides enhanced tool patterns inspired by rust-sdk,
//! including automatic tool discovery, structured parameter handling,
//! and JSON Schema support for better developer experience.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;
use std::collections::HashMap;

use crate::integration::HostIntegration;
use super::{Tool, ToolResponse, ToolError, Permission};

/// Wrapper for structured parameter extraction with JSON Schema support
#[derive(Debug, Clone)]
pub struct Parameters<T>(pub T);

impl<T> Parameters<T>
where
    T: for<'de> Deserialize<'de>,
{
    /// Extract parameters from JSON value with validation
    pub fn from_json(value: JsonValue) -> Result<Self, ToolError> {
        let params: T = serde_json::from_value(value)
            .map_err(|e| ToolError::InvalidParameters(format!("Parameter validation failed: {}", e)))?;
        Ok(Parameters(params))
    }
}

/// Result type for tool execution that matches MCP patterns
#[derive(Debug, Clone, Serialize)]
pub struct CallToolResult {
    pub content: Vec<Content>,
    pub is_error: bool,
}

/// Content type for tool responses
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum Content {
    #[serde(rename = "text")]
    Text { text: String },
    #[serde(rename = "image")]
    Image { data: String, mime_type: String },
    #[serde(rename = "resource")]
    Resource { resource: ResourceReference },
}

/// Resource reference for content
#[derive(Debug, Clone, Serialize)]
pub struct ResourceReference {
    pub uri: String,
    pub text: Option<String>,
}

impl CallToolResult {
    /// Create a successful result with text content
    pub fn success(content: Vec<Content>) -> Self {
        Self {
            content,
            is_error: false,
        }
    }

    /// Create an error result
    pub fn error(message: String) -> Self {
        Self {
            content: vec![Content::text(message)],
            is_error: true,
        }
    }
}

impl Content {
    /// Create text content
    pub fn text(text: String) -> Self {
        Self::Text { text }
    }

    /// Create image content
    pub fn image(data: String, mime_type: String) -> Self {
        Self::Image { data, mime_type }
    }

    /// Create resource content
    pub fn resource(uri: String, text: Option<String>) -> Self {
        Self::Resource {
            resource: ResourceReference { uri, text },
        }
    }
}

/// Enhanced tool trait that supports structured parameters and JSON Schema
#[async_trait]
pub trait EnhancedTool: Send + Sync {
    /// Execute the tool with structured parameters
    async fn execute_enhanced(
        &self,
        parameters: JsonValue,
        host: &dyn HostIntegration,
    ) -> Result<CallToolResult, ToolError>;

    /// Get the tool's name/identifier
    fn name(&self) -> &str;

    /// Get a description of what this tool does
    fn description(&self) -> &str;

    /// Get the JSON schema for the tool's parameters
    fn parameter_schema(&self) -> JsonValue;

    /// Get the permission required to execute this tool
    fn requires_permission(&self) -> Permission;

    /// Clone the tool (for use in collections)
    fn clone_enhanced(&self) -> Box<dyn EnhancedTool>;
}

/// Tool router for managing enhanced tools with automatic discovery
pub struct ToolRouter {
    tools: HashMap<String, Box<dyn EnhancedTool>>,
}

impl ToolRouter {
    /// Create a new tool router
    pub fn new() -> Self {
        Self {
            tools: HashMap::new(),
        }
    }

    /// Register an enhanced tool
    pub fn register<T: EnhancedTool + 'static>(&mut self, tool: T) {
        self.tools.insert(tool.name().to_string(), Box::new(tool));
    }

    /// Get a tool by name
    pub fn get(&self, name: &str) -> Option<&dyn EnhancedTool> {
        self.tools.get(name).map(|t| t.as_ref())
    }

    /// List all registered tools
    pub fn list_tools(&self) -> Vec<&str> {
        self.tools.keys().map(|s| s.as_str()).collect()
    }

    /// Get tool schemas for all registered tools
    pub fn get_tool_schemas(&self) -> HashMap<String, ToolSchema> {
        self.tools
            .iter()
            .map(|(name, tool)| {
                let schema = ToolSchema {
                    name: name.clone(),
                    description: tool.description().to_string(),
                    input_schema: tool.parameter_schema(),
                };
                (name.clone(), schema)
            })
            .collect()
    }

    /// Execute a tool by name with parameters
    pub async fn execute_tool(
        &self,
        name: &str,
        parameters: JsonValue,
        host: &dyn HostIntegration,
    ) -> Result<CallToolResult, ToolError> {
        let tool = self.get(name)
            .ok_or_else(|| ToolError::NotFound(name.to_string()))?;
        
        tool.execute_enhanced(parameters, host).await
    }
}

/// Tool schema for MCP compatibility
#[derive(Debug, Clone, Serialize)]
pub struct ToolSchema {
    pub name: String,
    pub description: String,
    #[serde(rename = "inputSchema")]
    pub input_schema: JsonValue,
}

/// Adapter to bridge enhanced tools with the legacy tool interface
pub struct EnhancedToolAdapter {
    enhanced_tool: Box<dyn EnhancedTool>,
}

impl EnhancedToolAdapter {
    /// Create a new adapter for an enhanced tool
    pub fn new<T: EnhancedTool + 'static>(tool: T) -> Self {
        Self {
            enhanced_tool: Box::new(tool),
        }
    }
}

#[async_trait]
impl Tool for EnhancedToolAdapter {
    async fn execute(
        &self,
        parameters: JsonValue,
        host: &dyn HostIntegration,
    ) -> Result<ToolResponse, ToolError> {
        let result = self.enhanced_tool.execute_enhanced(parameters, host).await?;
        
        // Convert CallToolResult to ToolResponse
        let content = result.content
            .into_iter()
            .map(|c| match c {
                Content::Text { text } => text,
                Content::Image { data, mime_type } => format!("Image: {} ({})", data, mime_type),
                Content::Resource { resource } => {
                    format!("Resource: {} ({})", resource.uri, resource.text.unwrap_or_default())
                }
            })
            .collect::<Vec<_>>()
            .join("\n");

        Ok(ToolResponse {
            content,
            success: !result.is_error,
            metadata: JsonValue::Null,
            affected_files: Vec::new(),
        })
    }

    fn requires_permission(&self) -> Permission {
        self.enhanced_tool.requires_permission()
    }

    fn description(&self) -> &str {
        self.enhanced_tool.description()
    }

    fn name(&self) -> &str {
        self.enhanced_tool.name()
    }

    fn parameter_schema(&self) -> JsonValue {
        self.enhanced_tool.parameter_schema()
    }

    fn clone_box(&self) -> Box<dyn Tool> {
        Box::new(EnhancedToolAdapter {
            enhanced_tool: self.enhanced_tool.clone_enhanced(),
        })
    }
}

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

/// Utility functions for creating JSON schemas
pub mod schema_utils {
    use super::*;
    #[cfg(feature = "mcp-server")]
    use schemars::{JsonSchema, schema_for};

    /// Generate JSON schema for a type that implements JsonSchema
    #[cfg(feature = "mcp-server")]
    pub fn generate_schema<T: JsonSchema>() -> JsonValue {
        let schema = schema_for!(T);
        serde_json::to_value(schema).unwrap_or(JsonValue::Null)
    }

    #[cfg(not(feature = "mcp-server"))]
    pub fn generate_schema<T>() -> JsonValue {
        serde_json::json!({})
    }

    /// Create a simple string parameter schema
    pub fn string_param(description: &str, required: bool) -> JsonValue {
        serde_json::json!({
            "type": "object",
            "properties": {
                "value": {
                    "type": "string",
                    "description": description
                }
            },
            "required": if required { vec!["value"] } else { vec![] }
        })
    }

    /// Create a simple number parameter schema
    pub fn number_param(description: &str, required: bool) -> JsonValue {
        serde_json::json!({
            "type": "object",
            "properties": {
                "value": {
                    "type": "number",
                    "description": description
                }
            },
            "required": if required { vec!["value"] } else { vec![] }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[cfg(feature = "schemars")]
    use schemars::JsonSchema;

    #[derive(Deserialize)]
    #[cfg_attr(feature = "schemars", derive(JsonSchema))]
    struct TestParams {
        message: String,
        count: Option<i32>,
    }

    struct TestTool;

    #[async_trait]
    impl EnhancedTool for TestTool {
        async fn execute_enhanced(
            &self,
            parameters: JsonValue,
            _host: &dyn HostIntegration,
        ) -> Result<CallToolResult, ToolError> {
            let Parameters(params): Parameters<TestParams> = Parameters::from_json(parameters)?;
            let response = format!("Message: {}, Count: {:?}", params.message, params.count);
            Ok(CallToolResult::success(vec![Content::text(response)]))
        }

        fn name(&self) -> &str {
            "test_tool"
        }

        fn description(&self) -> &str {
            "A test tool for demonstration"
        }

        fn parameter_schema(&self) -> JsonValue {
            schema_utils::generate_schema::<TestParams>()
        }

        fn requires_permission(&self) -> Permission {
            Permission::None
        }

        fn clone_enhanced(&self) -> Box<dyn EnhancedTool> {
            Box::new(TestTool)
        }
    }

    #[tokio::test]
    async fn test_enhanced_tool_router() {
        let mut router = ToolRouter::new();
        router.register(TestTool);

        assert_eq!(router.list_tools(), vec!["test_tool"]);
        assert!(router.get("test_tool").is_some());
        assert!(router.get("nonexistent").is_none());
    }

    #[test]
    fn test_parameters_extraction() {
        let json = serde_json::json!({
            "message": "hello",
            "count": 42
        });

        let params: Parameters<TestParams> = Parameters::from_json(json).unwrap();
        assert_eq!(params.0.message, "hello");
        assert_eq!(params.0.count, Some(42));
    }

    #[test]
    fn test_call_tool_result() {
        let result = CallToolResult::success(vec![Content::text("success".to_string())]);
        assert!(!result.is_error);
        assert_eq!(result.content.len(), 1);

        let error_result = CallToolResult::error("error message".to_string());
        assert!(error_result.is_error);
    }
}