inventory.rs

//! MCP tool, prompt, and resource inventory.
//!
//! Thread-safe cache for MCP capabilities across all connected servers.

use dashmap::DashMap;

use crate::mcp::config::{Prompt, RawResource, Tool};

/// Cached tool with metadata
#[derive(Clone)]
pub struct CachedTool {
    pub server_name: String,
    pub tool: Tool,
}

/// Cached prompt with metadata
#[derive(Clone)]
pub struct CachedPrompt {
    pub server_name: String,
    pub prompt: Prompt,
}

/// Cached resource with metadata
#[derive(Clone)]
pub struct CachedResource {
    pub server_name: String,
    pub resource: RawResource,
}

/// Tool inventory with periodic refresh
///
/// Provides thread-safe caching of MCP tools, prompts, and resources.
/// Entries are refreshed periodically by background tasks.
pub struct ToolInventory {
    /// Map of tool_name -> cached tool
    tools: DashMap<String, CachedTool>,

    /// Map of prompt_name -> cached prompt
    prompts: DashMap<String, CachedPrompt>,

    /// Map of resource_uri -> cached resource
    resources: DashMap<String, CachedResource>,
}

impl ToolInventory {
    /// Create a new tool inventory
    pub fn new() -> Self {
        Self {
            tools: DashMap::new(),
            prompts: DashMap::new(),
            resources: DashMap::new(),
        }
    }
}

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

impl ToolInventory {
    // ============================================================================
    // Tool Methods
    // ============================================================================

    /// Get a tool if it exists
    pub fn get_tool(&self, tool_name: &str) -> Option<(String, Tool)> {
        self.tools
            .get(tool_name)
            .map(|entry| (entry.server_name.clone(), entry.tool.clone()))
    }

    /// Check if tool exists
    pub fn has_tool(&self, tool_name: &str) -> bool {
        self.tools.contains_key(tool_name)
    }

    /// Insert or update a tool
    pub fn insert_tool(&self, tool_name: String, server_name: String, tool: Tool) {
        self.tools
            .insert(tool_name, CachedTool { server_name, tool });
    }

    /// Get all tools
    pub fn list_tools(&self) -> Vec<(String, String, Tool)> {
        self.tools
            .iter()
            .map(|entry| {
                let (name, cached) = entry.pair();
                (
                    name.clone(),
                    cached.server_name.clone(),
                    cached.tool.clone(),
                )
            })
            .collect()
    }

    // ============================================================================
    // Prompt Methods
    // ============================================================================

    /// Get a prompt if it exists
    pub fn get_prompt(&self, prompt_name: &str) -> Option<(String, Prompt)> {
        self.prompts
            .get(prompt_name)
            .map(|entry| (entry.server_name.clone(), entry.prompt.clone()))
    }

    /// Check if prompt exists
    pub fn has_prompt(&self, prompt_name: &str) -> bool {
        self.prompts.contains_key(prompt_name)
    }

    /// Insert or update a prompt
    pub fn insert_prompt(&self, prompt_name: String, server_name: String, prompt: Prompt) {
        self.prompts.insert(
            prompt_name,
            CachedPrompt {
                server_name,
                prompt,
            },
        );
    }

    /// Get all prompts
    pub fn list_prompts(&self) -> Vec<(String, String, Prompt)> {
        self.prompts
            .iter()
            .map(|entry| {
                let (name, cached) = entry.pair();
                (
                    name.clone(),
                    cached.server_name.clone(),
                    cached.prompt.clone(),
                )
            })
            .collect()
    }

    // ============================================================================
    // Resource Methods
    // ============================================================================

    /// Get a resource if it exists
    pub fn get_resource(&self, resource_uri: &str) -> Option<(String, RawResource)> {
        self.resources
            .get(resource_uri)
            .map(|entry| (entry.server_name.clone(), entry.resource.clone()))
    }

    /// Check if resource exists
    pub fn has_resource(&self, resource_uri: &str) -> bool {
        self.resources.contains_key(resource_uri)
    }

    /// Insert or update a resource
    pub fn insert_resource(
        &self,
        resource_uri: String,
        server_name: String,
        resource: RawResource,
    ) {
        self.resources.insert(
            resource_uri,
            CachedResource {
                server_name,
                resource,
            },
        );
    }

    /// Get all resources
    pub fn list_resources(&self) -> Vec<(String, String, RawResource)> {
        self.resources
            .iter()
            .map(|entry| {
                let (uri, cached) = entry.pair();
                (
                    uri.clone(),
                    cached.server_name.clone(),
                    cached.resource.clone(),
                )
            })
            .collect()
    }

    // ============================================================================
    // Server Management Methods
    // ============================================================================

    /// Clear all cached items for a specific server (called when LRU evicts client)
    pub fn clear_server_tools(&self, server_name: &str) {
        self.tools
            .retain(|_, cached| cached.server_name != server_name);
        self.prompts
            .retain(|_, cached| cached.server_name != server_name);
        self.resources
            .retain(|_, cached| cached.server_name != server_name);
    }

    /// Get count of cached items
    pub fn counts(&self) -> (usize, usize, usize) {
        (self.tools.len(), self.prompts.len(), self.resources.len())
    }

    /// Clear all cached items
    pub fn clear_all(&self) {
        self.tools.clear();
        self.prompts.clear();
        self.resources.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcp::config::{Prompt, RawResource, Tool};

    // Helper to create a test tool
    fn create_test_tool(name: &str) -> Tool {
        use std::{borrow::Cow, sync::Arc};

        let schema_obj = serde_json::json!({
            "type": "object",
            "properties": {}
        });

        let schema_map = if let serde_json::Value::Object(m) = schema_obj {
            m
        } else {
            serde_json::Map::new()
        };

        Tool {
            name: Cow::Owned(name.to_string()),
            title: None,
            description: Some(Cow::Owned(format!("Test tool: {}", name))),
            input_schema: Arc::new(schema_map),
            output_schema: None,
            annotations: None,
            icons: None,
        }
    }

    // Helper to create a test prompt
    fn create_test_prompt(name: &str) -> Prompt {
        Prompt {
            name: name.to_string(),
            title: None,
            description: Some(format!("Test prompt: {}", name)),
            arguments: None,
            icons: None,
        }
    }

    // Helper to create a test resource
    fn create_test_resource(uri: &str) -> RawResource {
        RawResource {
            uri: uri.to_string(),
            name: uri.to_string(),
            title: None,
            description: Some(format!("Test resource: {}", uri)),
            mime_type: Some("text/plain".to_string()),
            size: None,
            icons: None,
        }
    }

    #[test]
    fn test_tool_insert_and_get() {
        let inventory = ToolInventory::new();
        let tool = create_test_tool("test_tool");

        inventory.insert_tool("test_tool".to_string(), "server1".to_string(), tool.clone());

        let result = inventory.get_tool("test_tool");
        assert!(result.is_some());

        let (server_name, retrieved_tool) = result.unwrap();
        assert_eq!(server_name, "server1");
        assert_eq!(retrieved_tool.name, "test_tool");
    }

    #[test]
    fn test_has_tool() {
        let inventory = ToolInventory::new();
        let tool = create_test_tool("check_tool");

        assert!(!inventory.has_tool("check_tool"));

        inventory.insert_tool("check_tool".to_string(), "server1".to_string(), tool);

        assert!(inventory.has_tool("check_tool"));
    }

    #[test]
    fn test_list_tools() {
        let inventory = ToolInventory::new();

        inventory.insert_tool(
            "tool1".to_string(),
            "server1".to_string(),
            create_test_tool("tool1"),
        );
        inventory.insert_tool(
            "tool2".to_string(),
            "server1".to_string(),
            create_test_tool("tool2"),
        );
        inventory.insert_tool(
            "tool3".to_string(),
            "server2".to_string(),
            create_test_tool("tool3"),
        );

        let tools = inventory.list_tools();
        assert_eq!(tools.len(), 3);
    }

    #[test]
    fn test_clear_server_tools() {
        let inventory = ToolInventory::new();

        inventory.insert_tool(
            "tool1".to_string(),
            "server1".to_string(),
            create_test_tool("tool1"),
        );
        inventory.insert_tool(
            "tool2".to_string(),
            "server2".to_string(),
            create_test_tool("tool2"),
        );

        assert_eq!(inventory.list_tools().len(), 2);

        inventory.clear_server_tools("server1");

        let tools = inventory.list_tools();
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0].0, "tool2");
    }

    #[test]
    fn test_prompt_operations() {
        let inventory = ToolInventory::new();
        let prompt = create_test_prompt("test_prompt");

        inventory.insert_prompt(
            "test_prompt".to_string(),
            "server1".to_string(),
            prompt.clone(),
        );

        assert!(inventory.has_prompt("test_prompt"));

        let result = inventory.get_prompt("test_prompt");
        assert!(result.is_some());

        let (server_name, retrieved_prompt) = result.unwrap();
        assert_eq!(server_name, "server1");
        assert_eq!(retrieved_prompt.name, "test_prompt");
    }

    #[test]
    fn test_resource_operations() {
        let inventory = ToolInventory::new();
        let resource = create_test_resource("file:///test.txt");

        inventory.insert_resource(
            "file:///test.txt".to_string(),
            "server1".to_string(),
            resource.clone(),
        );

        assert!(inventory.has_resource("file:///test.txt"));

        let result = inventory.get_resource("file:///test.txt");
        assert!(result.is_some());

        let (server_name, retrieved_resource) = result.unwrap();
        assert_eq!(server_name, "server1");
        assert_eq!(retrieved_resource.uri, "file:///test.txt");
    }

    #[tokio::test]
    async fn test_concurrent_access() {
        use std::sync::Arc;

        let inventory = Arc::new(ToolInventory::new());

        // Spawn multiple tasks that insert tools concurrently
        let mut handles = vec![];
        for i in 0..10 {
            let inv = Arc::clone(&inventory);
            let handle = tokio::spawn(async move {
                let tool = create_test_tool(&format!("tool_{}", i));
                inv.insert_tool(format!("tool_{}", i), format!("server_{}", i % 3), tool);
            });
            handles.push(handle);
        }

        // Wait for all tasks to complete
        for handle in handles {
            handle.await.unwrap();
        }

        // Should have 10 tools
        let (tools, _, _) = inventory.counts();
        assert_eq!(tools, 10);
    }

    #[test]
    fn test_clear_all() {
        let inventory = ToolInventory::new();

        inventory.insert_tool(
            "tool1".to_string(),
            "server1".to_string(),
            create_test_tool("tool1"),
        );
        inventory.insert_prompt(
            "prompt1".to_string(),
            "server1".to_string(),
            create_test_prompt("prompt1"),
        );
        inventory.insert_resource(
            "res1".to_string(),
            "server1".to_string(),
            create_test_resource("res1"),
        );

        let (tools, prompts, resources) = inventory.counts();
        assert_eq!(tools, 1);
        assert_eq!(prompts, 1);
        assert_eq!(resources, 1);

        inventory.clear_all();

        let (tools, prompts, resources) = inventory.counts();
        assert_eq!(tools, 0);
        assert_eq!(prompts, 0);
        assert_eq!(resources, 0);
    }
}