alou 0.1.0

use crate::types::*;
use crate::mcp_tool::DiscoveredMcpTool;
use crate::mcp_config::McpConfigLoader;
use crate::mcp_client::McpClientManager;
use crate::types::McpServerConfig;
use crate::workspace_context::{WorkspaceContext, BasicWorkspaceContext};
use crate::prompt_registry::PromptRegistry;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use anyhow::{Result, Context};
use serde_json::json;
use tokio::sync::RwLock;
use std::process::Command;
use std::time::Duration;
use tokio::time::timeout;

/// 工具注册统计信息
#[derive(Debug, Clone)]
pub struct ToolRegistrationStats {
    /// 总工具数量
    pub total_tools: usize,
    /// 按服务器分组的工具数量
    pub tools_by_server: HashMap<String, usize>,
    /// 成功注册的工具数量
    pub successful_registrations: usize,
    /// 失败注册的工具数量
    pub failed_registrations: usize,
    /// 所有工具名称列表
    pub tool_names: Vec<String>,
}

/// 服务器工具统计信息
#[derive(Debug, Clone)]
pub struct ServerToolStats {
    /// 服务器名称
    pub server_name: String,
    /// 工具数量
    pub tool_count: usize,
    /// 工具名称列表
    pub tool_names: Vec<String>,
    /// 注册是否成功
    pub registration_success: bool,
}

/// 工具注册状态
#[derive(Debug, Clone)]
pub enum RegistrationStatus {
    /// 没有工具注册
    NoToolsRegistered,
    /// 部分成功
    PartialSuccess(ToolRegistrationStats),
    /// 完全成功
    Success(ToolRegistrationStats),
}

/// 工具注册完整性报告
#[derive(Debug, Clone)]
pub struct RegistrationIntegrityReport {
    /// 统计信息
    pub stats: ToolRegistrationStats,
    /// 服务器统计信息
    pub server_stats: HashMap<String, ServerToolStats>,
    /// 发现的问题
    pub issues: Vec<String>,
    /// 警告信息
    pub warnings: Vec<String>,
    /// 整体状态
    pub overall_status: String,
}

/// 工具注册表
/// 管理所有可用工具的注册和发现
#[derive(Clone)]
pub struct ToolRegistry {
    tools: Arc<RwLock<HashMap<String, Box<dyn Tool + Send + Sync>>>>,
    mcp_manager: Arc<McpClientManager>,
}

impl ToolRegistry {
    /// 创建新的工具注册表
    pub fn new() -> Self {
        Self {
            tools: Arc::new(RwLock::new(HashMap::new())),
            mcp_manager: Arc::new(McpClientManager::new()),
        }
    }

    /// 注册MCP工具定义
    /// 
    /// # Arguments
    /// * `tool` - MCP工具对象
    pub async fn register_mcp_tool(&self, tool: Box<dyn Tool + Send + Sync>) -> Result<()> {
        // 只接受MCP工具
        if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
            let tool_name = tool.name().to_string();
            let mut tools = self.tools.write().await;
            
            if tools.contains_key(&tool_name) {
                // 转换为完全限定名称避免冲突
                let qualified_tool = mcp_tool.as_fully_qualified_tool();
                tools.insert(qualified_tool.name().to_string(), Box::new(qualified_tool));
            } else {
                tools.insert(tool_name, tool);
            }
            
            Ok(())
        } else {
            Err(anyhow::anyhow!("只支持注册MCP工具"))
        }
    }

    /// 移除所有MCP工具
    async fn remove_all_mcp_tools(&self) {
        let mut tools = self.tools.write().await;
        tools.clear(); // 清空所有工具，因为现在只管理MCP工具
    }

    /// 移除特定MCP服务器的所有工具
    /// 
    /// # Arguments
    /// * `server_name` - 要移除工具的服务器名称
    pub async fn remove_mcp_tools_by_server(&self, server_name: &str) {
        let mut tools = self.tools.write().await;
        tools.retain(|_, tool| {
            if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
                mcp_tool.server_name() != server_name
            } else {
                true
            }
        });
    }

    /// 发现所有MCP工具
    /// 可以多次调用以更新发现的工具
    /// 这将从MCP服务器发现工具
    pub async fn discover_all_mcp_tools(&self, debug_mode: bool) -> Result<()> {
        // 移除任何先前发现的MCP工具
        self.remove_all_mcp_tools().await;

        // 创建所需依赖的实例
        let workspace_context = Arc::new(BasicWorkspaceContext::new());
        let prompt_registry = Arc::new(PromptRegistry::new());

        // 从mcp.json加载MCP服务器配置
        let mcp_servers = self.load_mcp_servers().await?;
        
        if !mcp_servers.is_empty() {
            // 处理环境变量替换
            let processed_servers = self.process_environment_variables(mcp_servers).await?;
            
            // 发现MCP工具
            self.discover_mcp_tools_internal(processed_servers, debug_mode, workspace_context, prompt_registry).await?;
        }

        Ok(())
    }

    /// 从MCP服务器发现工具（别名方法）
    pub async fn discover_mcp_tools(&self, debug_mode: bool) -> Result<()> {
        self.discover_all_mcp_tools(debug_mode).await
    }

    /// 发现或重新发现单个MCP服务器的工具
    /// 
    /// # Arguments
    /// * `server_name` - 要发现工具的服务器名称
    /// * `debug_mode` - 是否启用调试模式
    pub async fn discover_tools_for_server(&self, server_name: &str, debug_mode: bool) -> Result<()> {
        // 移除此服务器先前发现的任何工具
        self.remove_mcp_tools_by_server(server_name).await;

        // 创建所需依赖的实例
        let workspace_context = Arc::new(BasicWorkspaceContext::new());
        let prompt_registry = Arc::new(PromptRegistry::new());

        // 从mcp.json加载MCP服务器配置
        let mcp_servers = self.load_mcp_servers().await?;
        
        // 过滤到仅指定的服务器
        let server_config = mcp_servers.get(server_name);
        let filtered_mcp_servers = if let Some(config) = server_config {
            let mut filtered = HashMap::new();
            filtered.insert(server_name.to_string(), config.clone());
            filtered
        } else {
            HashMap::new()
        };
        
        if !filtered_mcp_servers.is_empty() {
            // 处理环境变量替换
            let processed_servers = self.process_environment_variables(filtered_mcp_servers).await?;
            
            // 发现MCP工具
            self.discover_mcp_tools_internal(processed_servers, debug_mode, workspace_context, prompt_registry).await?;
        }

        Ok(())
    }

    /// 获取工具模式列表（FunctionDeclaration数组）
    /// 从ToolListUnion结构中提取声明
    /// 如果已配置，包括发现的（相对于注册的）工具
    /// 
    /// # Returns
    /// FunctionDeclaration数组
    pub async fn get_function_declarations(&self) -> Vec<serde_json::Value> {
        let tools = self.tools.read().await;
        let mut declarations = Vec::new();
        
        for tool in tools.values() {
            // 确保我们有正确的工具名称和参数模式
            if !tool.name().is_empty() {
                let declaration = serde_json::json!({
                    "type": "function",
                    "function": {
                        "name": tool.name(),
                        "description": tool.description(),
                        "parameters": tool.parameter_schema()
                    }
                });
                declarations.push(declaration);
            }
        }
        
        declarations
    }

    /// 基于工具名称列表获取过滤的工具模式列表
    /// 
    /// # Arguments
    /// * `tool_names` - 要包含的工具名称数组
    /// 
    /// # Returns
    /// 指定工具的FunctionDeclaration数组
    pub async fn get_function_declarations_filtered(&self, tool_names: &[String]) -> Vec<serde_json::Value> {
        let tools = self.tools.read().await;
        let mut declarations = Vec::new();
        
        for name in tool_names {
            if let Some(tool) = tools.get(name) {
                let declaration = serde_json::json!({
                    "type": "function",
                    "function": {
                        "name": tool.name(),
                        "description": tool.description(),
                        "parameters": tool.parameter_schema()
                    }
                });
                declarations.push(declaration);
            }
        }
        
        declarations
    }

    /// 返回所有已注册和发现的工具实例数组
    pub async fn get_all_tools(&self) -> Vec<Box<dyn Tool + Send + Sync>> {
        let tools = self.tools.read().await;
        let mut tool_list: Vec<Box<dyn Tool + Send + Sync>> = Vec::new();
        
        // 由于 dyn Tool 不能直接 clone，我们需要重新创建工具
        for (_name, tool) in tools.iter() {
            if let Some(cloned_tool) = self.clone_tool(tool).await {
                tool_list.push(cloned_tool);
            }
        }
        
        tool_list
    }

    /// 返回工具数量
    pub async fn get_tool_count(&self) -> usize {
        let tools = self.tools.read().await;
        tools.len()
    }

    /// 返回从特定MCP服务器注册的工具数组
    pub async fn get_tools_by_server(&self, server_name: &str) -> Vec<Box<dyn Tool + Send + Sync>> {
        let tools = self.tools.read().await;
        let mut server_tools = Vec::new();
        
        for tool in tools.values() {
            if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
                if mcp_tool.server_name() == server_name {
                    if let Some(cloned_tool) = self.clone_tool(tool).await {
                        server_tools.push(cloned_tool);
                    }
                }
            }
        }
        
        server_tools
    }

    /// 获取特定工具的定义
    pub async fn get_tool(&self, name: &str) -> Option<Box<dyn Tool + Send + Sync>> {
        let tools = self.tools.read().await;
        
        // 首先尝试精确匹配
        if let Some(tool) = tools.get(name) {
            return self.clone_tool(tool).await;
        }
        
        // 尝试其他可能的格式
        // 将连字符转换为下划线
        let underscore_name = name.replace('-', "_");
        if let Some(tool) = tools.get(&underscore_name) {
            return self.clone_tool(tool).await;
        }
        
        // 将下划线转换为连字符
        let dash_name = name.replace('_', "-");
        if let Some(tool) = tools.get(&dash_name) {
            return self.clone_tool(tool).await;
        }
        
        // 尝试查找MCP工具的完全限定名称
        for (key, tool) in tools.iter() {
            if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
                // 检查是否是完全限定名称
                if key == &format!("{}__{}", mcp_tool.server_name(), name) {
                    return self.clone_tool(tool).await;
                }
                // 检查原始名称
                if mcp_tool.server_tool_name() == name {
                    return self.clone_tool(tool).await;
                }
                // 检查原始名称的不同格式
                if mcp_tool.server_tool_name() == underscore_name {
                    return self.clone_tool(tool).await;
                }
                if mcp_tool.server_tool_name() == dash_name {
                    return self.clone_tool(tool).await;
                }
            }
        }
        
        None
    }
    
    /// 克隆MCP工具实例
    async fn clone_tool(&self, tool: &Box<dyn Tool + Send + Sync>) -> Option<Box<dyn Tool + Send + Sync>> {
        // 只处理MCP工具
        if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
            use crate::mcp_tool::{McpToolFactory, MockCallableTool};
            use std::sync::Arc;
            
            let mock_client = MockCallableTool::create_mock_client(mcp_tool.server_name().to_string());
            let cloned_tool = McpToolFactory::create_discovered_tool(
                mock_client,
                mcp_tool.server_name().to_string(),
                mcp_tool.server_tool_name().to_string(),
                mcp_tool.description().to_string(),
                mcp_tool.parameter_schema().clone(),
                Some(30000),
                Some(true),
            );
            
            Some(Box::new(cloned_tool))
        } else {
            // 如果不是MCP工具，返回None
            None
        }
    }

    /// 获取工具数量
    pub async fn tool_count(&self) -> usize {
        let tools = self.tools.read().await;
        tools.len()
    }

    /// 检查工具是否存在
    pub async fn has_tool(&self, name: &str) -> bool {
        self.get_tool(name).await.is_some()
    }

    /// 获取所有工具名称
    pub async fn get_tool_names(&self) -> Vec<String> {
        let tools = self.tools.read().await;
        tools.keys().cloned().collect()
    }

    /// 获取工具注册统计信息
    pub async fn get_registration_stats(&self) -> ToolRegistrationStats {
        let tools = self.tools.read().await;
        let mut stats = ToolRegistrationStats {
            total_tools: tools.len(),
            tools_by_server: HashMap::new(),
            successful_registrations: 0,
            failed_registrations: 0,
            tool_names: Vec::new(),
        };

        for (name, tool) in tools.iter() {
            stats.tool_names.push(name.clone());
            stats.successful_registrations += 1;

            if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
                let server_name = mcp_tool.server_name().to_string();
                *stats.tools_by_server.entry(server_name).or_insert(0) += 1;
            }
        }

        stats
    }

    /// 检测工具注册状态
    pub async fn check_registration_status(&self) -> RegistrationStatus {
        let stats = self.get_registration_stats().await;
        
        if stats.total_tools == 0 {
            RegistrationStatus::NoToolsRegistered
        } else if stats.tools_by_server.is_empty() {
            RegistrationStatus::PartialSuccess(stats)
        } else {
            RegistrationStatus::Success(stats)
        }
    }

    /// 获取按服务器分组的工具统计
    pub async fn get_tools_by_server_stats(&self) -> HashMap<String, ServerToolStats> {
        let tools = self.tools.read().await;
        let mut server_stats: HashMap<String, ServerToolStats> = HashMap::new();

        for (name, tool) in tools.iter() {
            if let Some(mcp_tool) = tool.as_any().downcast_ref::<DiscoveredMcpTool>() {
                let server_name = mcp_tool.server_name().to_string();
                let stats = server_stats.entry(server_name.clone()).or_insert(ServerToolStats {
                    server_name: server_name.clone(),
                    tool_count: 0,
                    tool_names: Vec::new(),
                    registration_success: true,
                });
                
                stats.tool_count += 1;
                stats.tool_names.push(name.clone());
            }
        }

        server_stats
    }

    /// 验证工具注册完整性
    pub async fn validate_registration_integrity(&self) -> RegistrationIntegrityReport {
        let stats = self.get_registration_stats().await;
        let server_stats = self.get_tools_by_server_stats().await;
        let server_statuses = self.get_all_server_statuses().await;

        let mut issues = Vec::new();
        let mut warnings = Vec::new();

        // 检查是否有工具注册
        if stats.total_tools == 0 {
            issues.push("没有发现任何工具".to_string());
        }

        // 检查服务器状态与工具注册的一致性
        for (server_name, server_status) in server_statuses.iter() {
            let tool_count = server_stats.get(server_name).map(|s| s.tool_count).unwrap_or(0);
            
            match server_status {
                crate::mcp_client::McpServerStatus::Connected => {
                    if tool_count == 0 {
                        warnings.push(format!("服务器 '{}' 已连接但没有发现工具", server_name));
                    }
                }
                crate::mcp_client::McpServerStatus::Connecting => {
                    warnings.push(format!("服务器 '{}' 仍在连接中", server_name));
                }
                crate::mcp_client::McpServerStatus::Disconnected => {
                    if tool_count > 0 {
                        warnings.push(format!("服务器 '{}' 已断开连接但仍有 {} 个工具", server_name, tool_count));
                    }
                }
            }
        }

        // 检查工具名称冲突
        let mut name_counts: HashMap<String, usize> = HashMap::new();
        for name in &stats.tool_names {
            *name_counts.entry(name.clone()).or_insert(0) += 1;
        }

        for (name, count) in name_counts {
            if count > 1 {
                issues.push(format!("工具名称 '{}' 存在冲突 ({} 个实例)", name, count));
            }
        }

        RegistrationIntegrityReport {
            stats,
            server_stats,
            issues: issues.clone(),
            warnings: warnings.clone(),
            overall_status: if issues.is_empty() {
                if warnings.is_empty() {
                    "健康".to_string()
                } else {
                    "警告".to_string()
                }
            } else {
                "错误".to_string()
            },
        }
    }

    /// 清空所有工具
    pub async fn clear_all_tools(&self) {
        let mut tools = self.tools.write().await;
        tools.clear();
    }

    /// 加载MCP服务器配置
    async fn load_mcp_servers(&self) -> Result<HashMap<String, McpServerConfig>> {
        // 尝试多个可能的路径
        let possible_paths = vec![
            PathBuf::from("mcp.json"),                    // 当前工作目录
                
            PathBuf::from("../mcp.json"),                 // 开发环境
        ];
        
        for path in &possible_paths {
            if path.exists() {
                println!("🔍 找到MCP配置文件: {}", path.display());
                if let Some(servers) = McpConfigLoader::load_mcp_config_with_root(".")? {
                    println!("✅ 成功加载MCP配置，包含 {} 个服务器", servers.len());
                    return Ok(servers);
                }
            }
        }
        
        println!("⚠️ 未找到MCP配置文件，尝试的路径:");
        for path in &possible_paths {
            println!("  - {}", path.display());
        }
        
        Ok(HashMap::new())
    }

    /// 处理MCP服务器配置中的环境变量
    async fn process_environment_variables(&self, mcp_servers: HashMap<String, McpServerConfig>) -> Result<HashMap<String, McpServerConfig>> {
        let mut processed_servers = HashMap::new();
        
        for (server_name, mut server_config) in mcp_servers {
            // 处理args数组中的环境变量
            if let Some(args) = &mut server_config.args {
                let mut new_args = Vec::new();
                
                for arg in args.iter() {
                    // 处理 ALOU_INSTALL_DIR
                    if arg.contains("${ALOU_INSTALL_DIR}") {
                        let install_dir = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
                        let resolved_path = self.resolve_fetch_path(arg, &install_dir, &server_name);
                        new_args.push(resolved_path);
                    }
                    // 处理 OS_FILESYSTEM_PATHS_ARRAY - 展开为多个参数
                    else if arg.contains("${OS_FILESYSTEM_PATHS_ARRAY}") {
                        let paths = self.get_os_filesystem_paths();
                        // 将逗号分隔的字符串转换为数组并展开
                        let path_array: Vec<String> = paths.split(',')
                            .map(|s| s.trim().to_string())
                            .filter(|s| !s.is_empty())
                            .collect();
                        
                        // 在Mac上添加调试信息
                        if cfg!(target_os = "macos") && !path_array.is_empty() {
                            tracing::debug!("Filesystem paths for Mac: {}", path_array.join(", "));
                        }
                        
                        new_args.extend(path_array);
                    }
                    else {
                        new_args.push(arg.clone());
                    }
                }
                
                server_config.args = Some(new_args);
            }
            
            processed_servers.insert(server_name, server_config);
        }
        
        Ok(processed_servers)
    }

    /// 解析fetch工具的路径
    fn resolve_fetch_path(&self, arg: &str, install_dir: &PathBuf, server_name: &str) -> String {
        // 智能路径解析：尝试多个可能的路径
        let possible_paths = vec![
            install_dir.join("dist/src/fetch.js"),           // 生产环境：npm包安装目录
            install_dir.join("../dist/src/fetch.js"),        // 开发环境：从src目录向上查找
            PathBuf::from("dist/src/fetch.js"),              // 当前工作目录
            PathBuf::from("src/fetch.ts"),                   // 开发环境：直接使用源码
        ];
        
        // 查找第一个存在的文件
        for test_path in possible_paths {
            if test_path.exists() {
                return test_path.to_string_lossy().to_string();
            }
        }
        
        // 如果找到了文件，使用绝对路径；否则使用原始路径
        tracing::warn!("Warning: Could not find fetch.js in any expected location for server {}", server_name);
        arg.replace("${ALOU_INSTALL_DIR}", &install_dir.to_string_lossy())
    }

    /// 获取跨平台的文件系统路径
    fn get_os_filesystem_paths(&self) -> String {
        if cfg!(target_os = "windows") {
            // Windows: 检测存在的盘符
            self.get_windows_drives()
        } else if cfg!(target_os = "macos") {
            // macOS: 只使用用户目录，避免根目录权限问题
            "/Users".to_string()
        } else if cfg!(target_os = "linux") {
            // Linux
            "/,/home".to_string()
        } else {
            "/".to_string()
        }
    }

    /// 获取Windows系统的可用盘符
    fn get_windows_drives(&self) -> String {
        let mut drives = Vec::new();
        
        // 检测常见的Windows盘符
        let common_drives = vec!['C', 'D', 'E', 'F'];
        for drive in common_drives {
            let drive_path = format!("{}:\\", drive);
            if PathBuf::from(&drive_path).exists() {
                drives.push(drive_path);
            }
        }
        
        // 如果没有找到任何盘符，至少返回C盘
        if drives.is_empty() {
            "C:\\".to_string()
        } else {
            drives.join(",")
        }
    }

    /// 内部MCP工具发现方法
    async fn discover_mcp_tools_internal(
        &self,
        mcp_servers: HashMap<String, McpServerConfig>,
        debug_mode: bool,
        workspace_context: Arc<dyn WorkspaceContext + Send + Sync>,
        prompt_registry: Arc<PromptRegistry>,
    ) -> Result<()> {
        if debug_mode {
            tracing::debug!("Discovering MCP tools from {} servers", mcp_servers.len());
        }
        
        if debug_mode {
            println!("🔧 使用MCP客户端发现工具");
        }
        
        // 发现所有MCP服务器的工具
        self.mcp_manager.discover_mcp_tools(
            mcp_servers,
            Arc::new(self.clone()),
            prompt_registry,
            debug_mode,
            workspace_context,
        ).await?;
        
        Ok(())
    }

    /// 并发调用多个工具
    /// 
    /// # Arguments
    /// * `tool_calls` - 工具调用列表
    /// 
    /// # Returns
    /// 工具调用结果列表
    pub async fn call_tools_concurrent(
        &self,
        tool_calls: Vec<(String, HashMap<String, serde_json::Value>)>,
    ) -> Result<Vec<ToolResultContent>> {
        let mut tasks = Vec::new();
        
        for (tool_name, params) in tool_calls {
            let registry = self.clone();
            let task = tokio::spawn(async move {
                if let Some(tool) = registry.get_tool(&tool_name).await {
                    tool.execute(params).await
                } else {
                    Err(anyhow::anyhow!("工具 '{}' 未找到", tool_name).into())
                }
            });
            tasks.push(task);
        }
        
        let mut results = Vec::new();
        for task in tasks {
            match task.await {
                Ok(Ok(result)) => results.push(result),
                Ok(Err(e)) => {
                    results.push(ToolResultContent {
                        content: format!("工具调用失败: {}", e),
                        mime_type: None,
                        llm_content: None,
                        return_display: None,
                    });
                }
                Err(e) => {
                    results.push(ToolResultContent {
                        content: format!("任务执行失败: {}", e),
                        mime_type: None,
                        llm_content: None,
                        return_display: None,
                    });
                }
            }
        }
        
        Ok(results)
    }

    /// 带超时的工具调用
    /// 
    /// # Arguments
    /// * `tool_name` - 工具名称
    /// * `params` - 工具参数
    /// * `timeout_duration` - 超时时间
    /// 
    /// # Returns
    /// 工具调用结果
    pub async fn call_tool_with_timeout(
        &self,
        tool_name: &str,
        params: HashMap<String, serde_json::Value>,
        timeout_duration: Duration,
    ) -> Result<ToolResultContent> {
        let registry = self.clone();
        let tool_name = tool_name.to_string();
        
        timeout(timeout_duration, async move {
            if let Some(tool) = registry.get_tool(&tool_name).await {
                tool.execute(params).await.map_err(|e| anyhow::anyhow!("工具执行失败: {}", e))
            } else {
                Err(anyhow::anyhow!("工具 '{}' 未找到", tool_name))
            }
        }).await.map_err(|e| anyhow::anyhow!("工具调用超时: {}", e))?
    }

    /// 获取MCP客户端管理器
    pub fn get_mcp_manager(&self) -> Arc<McpClientManager> {
        self.mcp_manager.clone()
    }

    /// 获取服务器状态
    pub async fn get_server_status(&self, server_name: &str) -> crate::mcp_client::McpServerStatus {
        self.mcp_manager.get_server_status(server_name).await
    }

    /// 获取所有服务器状态
    pub async fn get_all_server_statuses(&self) -> HashMap<String, crate::mcp_client::McpServerStatus> {
        self.mcp_manager.get_all_server_statuses().await
    }

    /// 获取发现状态
    pub async fn get_discovery_state(&self) -> crate::mcp_client::McpDiscoveryState {
        self.mcp_manager.get_discovery_state().await
    }

    /// 关闭所有MCP连接
    pub async fn close_all_mcp_connections(&self) -> Result<()> {
        self.mcp_manager.close_all().await
    }

    /// 重新发现特定服务器的工具
    pub async fn rediscover_server_tools(&self, server_name: &str, debug_mode: bool) -> Result<()> {
        // 先关闭该服务器的连接
        if let Some(client) = self.mcp_manager.get_client(server_name).await {
            if let Ok(mut client_guard) = client.try_write() {
                let _ = client_guard.close().await;
            }
        }
        
        // 重新发现工具
        self.discover_tools_for_server(server_name, debug_mode).await
    }

    /// 动态更新MCP服务器配置
    pub async fn update_server_config(
        &self,
        server_name: &str,
        new_config: McpServerConfig,
        debug_mode: bool,
    ) -> Result<()> {
        // 关闭现有连接
        if let Some(client) = self.mcp_manager.get_client(server_name).await {
            if let Ok(mut client_guard) = client.try_write() {
                let _ = client_guard.close().await;
            }
        }
        
        // 移除现有工具
        self.remove_mcp_tools_by_server(server_name).await;
        
        // 使用新配置重新连接
        let mut servers = HashMap::new();
        servers.insert(server_name.to_string(), new_config);
        
        let processed_servers = self.process_environment_variables(servers).await?;
        let workspace_context = Arc::new(BasicWorkspaceContext::new());
        let prompt_registry = Arc::new(PromptRegistry::new());
        
        self.discover_mcp_tools_internal(processed_servers, debug_mode, workspace_context, prompt_registry).await?;
        
        Ok(())
    }

    /// 打印工具注册状态报告
    pub async fn print_registration_report(&self) {
        let report = self.validate_registration_integrity().await;
        
        println!("📊 工具注册状态报告");
        println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        println!("🔍 整体状态: {}", report.overall_status);
        println!("📈 总工具数量: {}", report.stats.total_tools);
        println!("✅ 成功注册: {}", report.stats.successful_registrations);
        println!("❌ 失败注册: {}", report.stats.failed_registrations);
        
        if !report.stats.tools_by_server.is_empty() {
            println!("\n🛠️ 按服务器分组的工具:");
            for (server_name, count) in &report.stats.tools_by_server {
                println!("  📡 {}: {} 个工具", server_name, count);
            }
        }
        
        if !report.stats.tool_names.is_empty() {
            println!("\n📋 已注册的工具列表:");
            for (i, name) in report.stats.tool_names.iter().enumerate() {
                println!("  {}. {}", i + 1, name);
            }
        }
        
        if !report.warnings.is_empty() {
            println!("\n⚠️ 警告:");
            for warning in &report.warnings {
                println!("  • {}", warning);
            }
        }
        
        if !report.issues.is_empty() {
            println!("\n❌ 问题:");
            for issue in &report.issues {
                println!("  • {}", issue);
            }
        }
        
        println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    }

    /// 获取简化的注册状态摘要
    pub async fn get_registration_summary(&self) -> String {
        let stats = self.get_registration_stats().await;
        let server_stats = self.get_tools_by_server_stats().await;
        
        let mut summary = format!("工具总数: {}, 成功: {}, 失败: {}", 
            stats.total_tools, stats.successful_registrations, stats.failed_registrations);
        
        if !server_stats.is_empty() {
            summary.push_str("\n服务器工具分布:");
            for (server_name, server_stat) in server_stats {
                summary.push_str(&format!("\n  {}: {} 个工具", server_name, server_stat.tool_count));
            }
        }
        
        summary
    }
}

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

/// 工具注册表构建器
pub struct ToolRegistryBuilder {
    registry: ToolRegistry,
}

impl ToolRegistryBuilder {
    /// 创建新的构建器
    pub fn new() -> Self {
        Self {
            registry: ToolRegistry::new(),
        }
    }

    /// 添加MCP工具
    pub async fn add_mcp_tool(self, tool: Box<dyn Tool + Send + Sync>) -> Result<Self> {
        self.registry.register_mcp_tool(tool).await?;
        Ok(self)
    }

    /// 发现所有MCP工具
    pub async fn discover_all_mcp_tools(self, debug_mode: bool) -> Result<Self> {
        self.registry.discover_all_mcp_tools(debug_mode).await?;
        Ok(self)
    }

    /// 构建工具注册表
    pub fn build(self) -> ToolRegistry {
        self.registry
    }
}

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

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

    #[tokio::test]
    async fn test_mcp_tool_registry_creation() {
        let registry = ToolRegistry::new();
        assert_eq!(registry.tool_count().await, 0);
    }

    #[tokio::test]
    async fn test_register_mcp_tool() {
        let registry = ToolRegistry::new();
        let tool = McpToolFactory::create_mock_tool(
            "test_server".to_string(),
            "test_tool".to_string(),
            "Test MCP tool".to_string(),
            serde_json::json!({"type": "object"}),
        );
        
        registry.register_mcp_tool(Box::new(tool)).await.unwrap();
        assert_eq!(registry.tool_count().await, 1);
        assert!(registry.has_tool("test_tool").await);
    }

    #[tokio::test]
    async fn test_get_mcp_tool() {
        let registry = ToolRegistry::new();
        let tool = McpToolFactory::create_mock_tool(
            "test_server".to_string(),
            "test_tool".to_string(),
            "Test MCP tool".to_string(),
            serde_json::json!({"type": "object"}),
        );
        
        registry.register_mcp_tool(Box::new(tool)).await.unwrap();
        
        let retrieved_tool = registry.get_tool("test_tool").await;
        assert!(retrieved_tool.is_some());
        assert_eq!(retrieved_tool.unwrap().name(), "test_tool");
    }

    #[tokio::test]
    async fn test_get_mcp_function_declarations() {
        let registry = ToolRegistry::new();
        let tool = McpToolFactory::create_mock_tool(
            "test_server".to_string(),
            "test_tool".to_string(),
            "Test MCP tool".to_string(),
            serde_json::json!({"type": "object"}),
        );
        
        registry.register_mcp_tool(Box::new(tool)).await.unwrap();
        
        let declarations = registry.get_function_declarations().await;
        assert_eq!(declarations.len(), 1);
        assert_eq!(declarations[0]["function"]["name"], "test_tool");
    }

    #[tokio::test]
    async fn test_get_tools_by_server() {
        let registry = ToolRegistry::new();
        let tool = McpToolFactory::create_mock_tool(
            "test_server".to_string(),
            "test_tool".to_string(),
            "Test MCP tool".to_string(),
            serde_json::json!({"type": "object"}),
        );
        
        registry.register_mcp_tool(Box::new(tool)).await.unwrap();
        
        let server_tools = registry.get_tools_by_server("test_server").await;
        assert_eq!(server_tools.len(), 1);
        assert_eq!(server_tools[0].name(), "test_tool");
    }

    #[tokio::test]
    async fn test_reject_non_mcp_tool() {
        let registry = ToolRegistry::new();
        let tool = Box::new(crate::tools::BaseDeclarativeTool::new(
            "non_mcp_tool".to_string(),
            "Non MCP Tool".to_string(),
            "This is not an MCP tool".to_string(),
            crate::types::Kind::Other,
            serde_json::json!({}),
            true,
            false,
        ));
        
        let result = registry.register_mcp_tool(tool).await;
        assert!(result.is_err());
        assert_eq!(registry.tool_count().await, 0);
    }
}

impl ToolRegistry {
    /// 发现NPM MCP服务器（兼容性方法）
    pub async fn discover_npm_mcp_servers(&self) -> Result<Vec<McpServerConfig>> {
        let servers = self.load_mcp_servers().await?;
        Ok(servers.into_values().collect())
    }

    /// 导出MCP配置（兼容性方法）
    pub fn export_mcp_config(&self, configs: Vec<McpServerConfig>, output: &str) -> Result<()> {
        let config_json = serde_json::to_string_pretty(&configs)?;
        std::fs::write(output, config_json)?;
        Ok(())
    }

    /// 自动发现并导出（兼容性方法）
    pub async fn auto_discover_and_export(&self, output: &str) -> Result<()> {
        let configs = self.discover_npm_mcp_servers().await?;
        self.export_mcp_config(configs, output)
    }
}