coderlib 0.1.0

//! MCP registry for managing multiple MCP servers

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::{
    tools::ToolRegistry,
    mcp::{
        config::{McpConfig, McpServerConfig},
        client::McpClient,
        tool_adapter::McpToolAdapter,
        error::{McpError, McpResult},
        types::{McpServerHealth, McpConnectionInfo},
    },
};

/// Registry for managing multiple MCP servers and their tools
pub struct McpRegistry {
    /// Configuration for MCP
    config: McpConfig,
    /// Map of server name to MCP client
    clients: Arc<RwLock<HashMap<String, Arc<McpClient>>>>,
    /// Whether the registry has been started
    started: Arc<RwLock<bool>>,
}

impl McpRegistry {
    /// Create a new MCP registry
    pub fn new(config: McpConfig) -> Self {
        Self {
            config,
            clients: Arc::new(RwLock::new(HashMap::new())),
            started: Arc::new(RwLock::new(false)),
        }
    }
    
    /// Start all configured MCP servers
    pub async fn start_all_servers(&mut self) -> McpResult<()> {
        if !self.config.enabled {
            tracing::info!("MCP is disabled, skipping server startup");
            return Ok(());
        }
        
        tracing::info!("Starting MCP registry with {} servers", self.config.servers.len());
        
        let mut clients = self.clients.write().await;
        let mut started_count = 0;
        
        for (name, server_config) in &self.config.servers {
            if server_config.enabled && server_config.auto_start {
                match self.start_server(name, server_config, &mut clients).await {
                    Ok(_) => {
                        started_count += 1;
                        tracing::info!("Started MCP server: {}", name);
                    }
                    Err(e) => {
                        tracing::error!("Failed to start MCP server {}: {}", name, e);
                        // Continue with other servers even if one fails
                    }
                }
            } else {
                tracing::debug!("Skipping MCP server {} (disabled or not auto-start)", name);
            }
        }
        
        *self.started.write().await = true;
        tracing::info!("MCP registry started with {}/{} servers", started_count, self.config.servers.len());
        
        Ok(())
    }
    
    /// Start a specific MCP server
    async fn start_server(
        &self,
        name: &str,
        config: &McpServerConfig,
        clients: &mut HashMap<String, Arc<McpClient>>,
    ) -> McpResult<()> {
        let mut client = McpClient::new(name.to_string(), config.clone());
        
        client.start().await?;
        
        clients.insert(name.to_string(), Arc::new(client));
        
        Ok(())
    }
    
    /// Stop all MCP servers
    pub async fn stop_all_servers(&mut self) -> McpResult<()> {
        tracing::info!("Stopping all MCP servers");
        
        let mut clients = self.clients.write().await;
        let mut stopped_count = 0;
        
        for (name, client) in clients.drain() {
            // Get mutable reference to client
            if let Ok(mut client_mut) = Arc::try_unwrap(client) {
                match client_mut.stop().await {
                    Ok(_) => {
                        stopped_count += 1;
                        tracing::info!("Stopped MCP server: {}", name);
                    }
                    Err(e) => {
                        tracing::error!("Failed to stop MCP server {}: {}", name, e);
                    }
                }
            } else {
                tracing::warn!("Could not stop MCP server {} (still in use)", name);
            }
        }
        
        *self.started.write().await = false;
        tracing::info!("Stopped {} MCP servers", stopped_count);
        
        Ok(())
    }
    
    /// Register all MCP tools with the tool registry
    pub async fn register_tools_with_registry(&self, tool_registry: &mut ToolRegistry) -> McpResult<()> {
        if !self.config.enabled {
            tracing::debug!("MCP is disabled, skipping tool registration");
            return Ok(());
        }
        
        let clients = self.clients.read().await;
        let mut registered_count = 0;
        
        for (server_name, client) in clients.iter() {
            match self.register_server_tools(server_name, client, tool_registry).await {
                Ok(count) => {
                    registered_count += count;
                    tracing::info!("Registered {} tools from MCP server: {}", count, server_name);
                }
                Err(e) => {
                    tracing::error!("Failed to register tools from MCP server {}: {}", server_name, e);
                }
            }
        }
        
        tracing::info!("Registered {} MCP tools total", registered_count);
        Ok(())
    }
    
    /// Register tools from a specific server
    #[cfg(feature = "mcp")]
    async fn register_server_tools(
        &self,
        server_name: &str,
        client: &Arc<McpClient>,
        tool_registry: &mut ToolRegistry,
    ) -> McpResult<usize> {
        let tools = client.list_tools().await?;
        let mut count = 0;
        
        for tool in tools {
            let adapter = McpToolAdapter::new(
                server_name.to_string(),
                tool,
                client.clone()
            );
            
            tool_registry.register(adapter);
            count += 1;
        }
        
        Ok(count)
    }
    
    #[cfg(not(feature = "mcp"))]
    async fn register_server_tools(
        &self,
        _server_name: &str,
        _client: &Arc<McpClient>,
        _tool_registry: &mut ToolRegistry,
    ) -> McpResult<usize> {
        Ok(0)
    }
    
    /// Get a specific MCP client by server name
    pub async fn get_client(&self, server_name: &str) -> Option<Arc<McpClient>> {
        let clients = self.clients.read().await;
        clients.get(server_name).cloned()
    }
    
    /// List all server names
    pub async fn list_servers(&self) -> Vec<String> {
        let clients = self.clients.read().await;
        clients.keys().cloned().collect()
    }
    
    /// Get health information for all servers
    pub async fn get_all_health(&self) -> HashMap<String, McpServerHealth> {
        let clients = self.clients.read().await;
        let mut health_map = HashMap::new();
        
        for (name, client) in clients.iter() {
            let health = client.health().await;
            health_map.insert(name.clone(), health);
        }
        
        health_map
    }
    
    /// Get connection information for all servers
    pub async fn get_all_connections(&self) -> HashMap<String, Option<McpConnectionInfo>> {
        let clients = self.clients.read().await;
        let mut conn_map = HashMap::new();
        
        for (name, client) in clients.iter() {
            let conn_info = client.connection_info().await;
            conn_map.insert(name.clone(), conn_info);
        }
        
        conn_map
    }
    
    /// Check if the registry has been started
    pub async fn is_started(&self) -> bool {
        *self.started.read().await
    }
    
    /// Get the configuration
    pub fn config(&self) -> &McpConfig {
        &self.config
    }
    
    /// Update configuration (requires restart to take effect)
    pub fn update_config(&mut self, config: McpConfig) {
        self.config = config;
    }
    
    /// Restart a specific server
    pub async fn restart_server(&self, server_name: &str) -> McpResult<()> {
        let server_config = self.config.servers.get(server_name)
            .ok_or_else(|| McpError::server_not_found(server_name))?
            .clone();
        
        // Stop the existing server
        {
            let mut clients = self.clients.write().await;
            if let Some(client) = clients.remove(server_name) {
                if let Ok(mut client_mut) = Arc::try_unwrap(client) {
                    let _ = client_mut.stop().await;
                }
            }
        }
        
        // Start the server again
        {
            let mut clients = self.clients.write().await;
            self.start_server(server_name, &server_config, &mut clients).await?;
        }
        
        tracing::info!("Restarted MCP server: {}", server_name);
        Ok(())
    }
    
    /// Add a new server configuration and start it
    pub async fn add_server(&mut self, name: String, config: McpServerConfig) -> McpResult<()> {
        // Add to configuration
        self.config.servers.insert(name.clone(), config.clone());
        
        // Start the server if auto-start is enabled
        if config.enabled && config.auto_start {
            let mut clients = self.clients.write().await;
            self.start_server(&name, &config, &mut clients).await?;
        }
        
        Ok(())
    }
    
    /// Remove a server configuration and stop it
    pub async fn remove_server(&mut self, name: &str) -> McpResult<()> {
        // Stop the server
        {
            let mut clients = self.clients.write().await;
            if let Some(client) = clients.remove(name) {
                if let Ok(mut client_mut) = Arc::try_unwrap(client) {
                    let _ = client_mut.stop().await;
                }
            }
        }
        
        // Remove from configuration
        self.config.servers.remove(name);
        
        tracing::info!("Removed MCP server: {}", name);
        Ok(())
    }
}

impl Drop for McpRegistry {
    fn drop(&mut self) {
        // Note: We can't call async methods in Drop, so we just log
        tracing::debug!("MCP registry dropped");
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcp::config::McpTransportConfig;
    
    #[tokio::test]
    async fn test_registry_creation() {
        let config = McpConfig::default();
        let registry = McpRegistry::new(config);
        
        assert!(!registry.is_started().await);
        assert_eq!(registry.list_servers().await.len(), 0);
    }
    
    #[tokio::test]
    async fn test_disabled_registry() {
        let mut config = McpConfig::default();
        config.enabled = false;
        
        let mut registry = McpRegistry::new(config);
        let result = registry.start_all_servers().await;
        
        assert!(result.is_ok());
        assert!(!registry.is_started().await);
    }
}