coderlib 0.1.0

//! MCP client implementation

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

use crate::mcp::{
    config::{McpServerConfig, McpTransportConfig},
    error::{McpError, McpResult},
    types::{McpServerHealth, McpServerStatus, McpConnectionInfo, McpServerCapabilities},
    simple::{SimpleTool, SimpleToolCall, SimpleToolResult, SimpleResource, SimpleServerInfo},
};

// Full protocol imports (when available)
#[cfg(all(feature = "mcp", feature = "mcp-full-protocol"))]
use rmcp::{ServiceExt, service::RoleClient, ClientHandler, transport::TokioChildProcess};
#[cfg(all(feature = "mcp", feature = "mcp-full-protocol"))]
use tokio::process::Command;

/// MCP client for connecting to and interacting with MCP servers
pub struct McpClient {
    /// Name of the server this client connects to
    server_name: String,
    /// Server configuration
    config: McpServerConfig,
    /// Running MCP service (when connected) - full protocol
    #[cfg(all(feature = "mcp", feature = "mcp-full-protocol"))]
    service: Option<rmcp::service::RunningService<rmcp::service::RoleClient, rmcp::model::ClientInfo>>,
    /// Simplified service placeholder
    #[cfg(all(feature = "mcp", not(feature = "mcp-full-protocol")))]
    service: Option<SimpleServerInfo>,
    #[cfg(not(feature = "mcp"))]
    service: Option<()>,
    /// Available tools from the server
    tools: Arc<RwLock<Vec<SimpleTool>>>,
    /// Available resources from the server
    resources: Arc<RwLock<Vec<SimpleResource>>>,
    /// Available prompts from the server (simplified as tools for now)
    prompts: Arc<RwLock<Vec<SimpleTool>>>,
    /// Server health information
    health: Arc<RwLock<McpServerHealth>>,
    /// Connection information
    connection_info: Arc<RwLock<Option<McpConnectionInfo>>>,
    /// Restart attempt counter
    restart_attempts: Arc<RwLock<u32>>,
}

impl McpClient {
    /// Create a new MCP client
    pub fn new(server_name: String, config: McpServerConfig) -> Self {
        Self {
            server_name: server_name.clone(),
            config,
            service: None,
            tools: Arc::new(RwLock::new(Vec::new())),
            resources: Arc::new(RwLock::new(Vec::new())),
            prompts: Arc::new(RwLock::new(Vec::new())),
            health: Arc::new(RwLock::new(McpServerHealth::new(server_name))),
            connection_info: Arc::new(RwLock::new(None)),
            restart_attempts: Arc::new(RwLock::new(0)),
        }
    }
    
    /// Start the MCP client and connect to the server
    pub async fn start(&mut self) -> McpResult<()> {
        tracing::info!("Starting MCP client for server: {}", self.server_name);
        
        // Update health status
        {
            let mut health = self.health.write().await;
            health.status = McpServerStatus::Starting;
        }
        
        let transport_config = self.config.transport.clone();
        match transport_config {
            McpTransportConfig::Stdio { command, args } => {
                self.start_stdio_client(&command, &args).await
            }
            McpTransportConfig::Sse { url, headers } => {
                self.start_sse_client(&url, &headers).await
            }
        }
    }
    
    /// Start a stdio-based MCP client
    #[cfg(feature = "mcp")]
    async fn start_stdio_client(&mut self, command: &str, args: &[String]) -> McpResult<()> {
        let mut cmd = Command::new(command);
        
        // Add arguments
        for arg in args {
            cmd.arg(arg);
        }
        
        // Set working directory if specified
        if let Some(working_dir) = &self.config.working_directory {
            cmd.current_dir(working_dir);
        }
        
        // Set environment variables
        let env_vars = self.get_environment_variables();
        for (key, value) in env_vars {
            cmd.env(key, value);
        }
        
        // Create transport and connect
        let transport = TokioChildProcess::new(cmd)
            .map_err(|e| McpError::transport(format!("Failed to create child process: {}", e)))?;
        
        let service = rmcp::model::ClientInfo::default().serve(transport).await
            .map_err(|e| McpError::connection(format!("Failed to connect to MCP server: {}", e)))?;
        
        // Store connection info
        {
            let mut conn_info = self.connection_info.write().await;
            *conn_info = Some(McpConnectionInfo {
                server_name: self.server_name.clone(),
                transport_type: "stdio".to_string(),
                connected_at: chrono::Utc::now(),
                reconnection_count: *self.restart_attempts.read().await,
                capabilities: Some(McpServerCapabilities::default()),
            });
        }
        
        // Discover capabilities
        self.discover_capabilities(&service).await?;

        self.service = Some(service);
        
        // Update health status
        {
            let mut health = self.health.write().await;
            health.status = McpServerStatus::Healthy;
        }
        
        tracing::info!("Successfully started stdio MCP client for: {}", self.server_name);
        Ok(())
    }
    
    #[cfg(not(feature = "mcp"))]
    async fn start_stdio_client(&mut self, _command: &str, _args: &[String]) -> McpResult<()> {
        Err(McpError::configuration("MCP feature not enabled"))
    }
    
    /// Start an SSE-based MCP client
    #[cfg(feature = "mcp")]
    async fn start_sse_client(&mut self, url: &str, _headers: &HashMap<String, String>) -> McpResult<()> {
        use rmcp::transport::SseClientTransport;

        let transport = SseClientTransport::start(url).await
            .map_err(|e| McpError::transport(format!("Failed to create SSE transport: {}", e)))?;

        let service = rmcp::model::ClientInfo::default().serve(transport).await
            .map_err(|e| McpError::connection(format!("Failed to connect to SSE MCP server: {}", e)))?;
        
        // Store connection info
        {
            let mut conn_info = self.connection_info.write().await;
            *conn_info = Some(McpConnectionInfo {
                server_name: self.server_name.clone(),
                transport_type: "sse".to_string(),
                connected_at: chrono::Utc::now(),
                reconnection_count: *self.restart_attempts.read().await,
                capabilities: Some(McpServerCapabilities::default()),
            });
        }
        
        // Discover capabilities
        self.discover_capabilities(&service).await?;

        self.service = Some(service);
        
        // Update health status
        {
            let mut health = self.health.write().await;
            health.status = McpServerStatus::Healthy;
        }
        
        tracing::info!("Successfully started SSE MCP client for: {}", self.server_name);
        Ok(())
    }
    
    #[cfg(not(feature = "mcp"))]
    async fn start_sse_client(&mut self, _url: &str, _headers: &HashMap<String, String>) -> McpResult<()> {
        Err(McpError::configuration("MCP feature not enabled"))
    }
    
    /// Discover server capabilities (tools, resources, prompts)
    #[cfg(feature = "mcp")]
    async fn discover_capabilities(&self, service: &rmcp::service::RunningService<RoleClient, rmcp::model::ClientInfo>) -> McpResult<()> {
        // Discover tools
        match service.list_all_tools().await {
            Ok(tools_result) => {
                let mut tools = self.tools.write().await;
                // Convert rmcp tools to simple tools
                *tools = tools_result.into_iter().map(|tool| {
                    SimpleTool {
                        name: tool.name.to_string(),
                        description: tool.description.map(|s| s.to_string()),
                        input_schema: serde_json::Value::Object((*tool.input_schema).clone()),
                    }
                }).collect();
                tracing::info!("Discovered {} tools from server: {}", tools.len(), self.server_name);
            }
            Err(e) => {
                tracing::warn!("Failed to discover tools from server {}: {}", self.server_name, e);
            }
        }
        
        // Discover resources
        match service.list_all_resources().await {
            Ok(resources_result) => {
                let mut resources = self.resources.write().await;
                // Convert rmcp resources to simple resources
                *resources = resources_result.into_iter().map(|resource| {
                    SimpleResource {
                        uri: resource.raw.uri,
                        name: Some(resource.raw.name),
                        description: resource.raw.description,
                        mime_type: resource.raw.mime_type,
                    }
                }).collect();
                tracing::info!("Discovered {} resources from server: {}", resources.len(), self.server_name);
            }
            Err(e) => {
                tracing::warn!("Failed to discover resources from server {}: {}", self.server_name, e);
            }
        }
        
        // Discover prompts
        match service.list_all_prompts().await {
            Ok(prompts_result) => {
                let mut prompts = self.prompts.write().await;
                // Convert rmcp prompts to simple tools (for compatibility)
                *prompts = prompts_result.into_iter().map(|prompt| {
                    SimpleTool {
                        name: prompt.name.to_string(),
                        description: prompt.description.map(|s| s.to_string()),
                        input_schema: serde_json::Value::Object(Default::default()), // Prompts don't have input schema
                    }
                }).collect();
                tracing::info!("Discovered {} prompts from server: {}", prompts.len(), self.server_name);
            }
            Err(e) => {
                tracing::warn!("Failed to discover prompts from server {}: {}", self.server_name, e);
            }
        }
        
        Ok(())
    }
    
    #[cfg(not(feature = "mcp"))]
    async fn discover_capabilities(&self, _service: &()) -> McpResult<()> {
        Ok(())
    }
    
    /// Get environment variables for the server process
    fn get_environment_variables(&self) -> HashMap<String, String> {
        let mut env = std::env::vars().collect::<HashMap<_, _>>();
        
        // Add server-specific environment variables
        for (key, value) in &self.config.env {
            env.insert(key.clone(), value.clone());
        }
        
        env
    }
    
    /// Stop the MCP client
    pub async fn stop(&mut self) -> McpResult<()> {
        tracing::info!("Stopping MCP client for server: {}", self.server_name);
        
        // Update health status
        {
            let mut health = self.health.write().await;
            health.status = McpServerStatus::Stopping;
        }
        
        #[cfg(feature = "mcp")]
        if let Some(service) = self.service.take() {
            if let Err(e) = service.cancel().await {
                tracing::warn!("Error stopping MCP service for {}: {}", self.server_name, e);
            }
        }
        
        // Clear connection info
        {
            let mut conn_info = self.connection_info.write().await;
            *conn_info = None;
        }
        
        // Update health status
        {
            let mut health = self.health.write().await;
            health.status = McpServerStatus::Disconnected;
        }
        
        tracing::info!("Successfully stopped MCP client for: {}", self.server_name);
        Ok(())
    }
    
    /// Check if the client is connected
    pub async fn is_connected(&self) -> bool {
        self.service.is_some() && {
            let health = self.health.read().await;
            health.is_available()
        }
    }
    
    /// Get server health information
    pub async fn health(&self) -> McpServerHealth {
        self.health.read().await.clone()
    }
    
    /// Get connection information
    pub async fn connection_info(&self) -> Option<McpConnectionInfo> {
        self.connection_info.read().await.clone()
    }
    
    /// Get the server name
    pub fn server_name(&self) -> &str {
        &self.server_name
    }
    
    /// Get the server configuration
    pub fn config(&self) -> &McpServerConfig {
        &self.config
    }

    /// Call an MCP tool
    #[cfg(feature = "mcp")]
    pub async fn call_tool(&self, request: rmcp::model::CallToolRequestParam) -> McpResult<rmcp::model::CallToolResult> {
        let service = self.service.as_ref()
            .ok_or_else(|| McpError::connection("Not connected to MCP server"))?;

        service.call_tool(request).await
            .map_err(|e| McpError::tool_execution(e.to_string()))
    }

    #[cfg(not(feature = "mcp"))]
    pub async fn call_tool(&self, _request: ()) -> McpResult<()> {
        Err(McpError::configuration("MCP feature not enabled"))
    }

    /// List available tools
    #[cfg(feature = "mcp")]
    pub async fn list_tools(&self) -> McpResult<Vec<rmcp::model::Tool>> {
        let simple_tools = self.tools.read().await;
        let rmcp_tools = simple_tools.iter().map(|tool| {
            rmcp::model::Tool {
                name: tool.name.clone().into(),
                description: tool.description.clone().map(|s| s.into()),
                input_schema: std::sync::Arc::new(
                    tool.input_schema.as_object().cloned().unwrap_or_default()
                ),
                annotations: None,
            }
        }).collect();
        Ok(rmcp_tools)
    }

    #[cfg(not(feature = "mcp"))]
    pub async fn list_tools(&self) -> McpResult<Vec<()>> {
        Ok(vec![])
    }

    /// List available resources
    #[cfg(feature = "mcp")]
    pub async fn list_resources(&self) -> McpResult<Vec<rmcp::model::Resource>> {
        let simple_resources = self.resources.read().await;
        let rmcp_resources = simple_resources.iter().map(|resource| {
            rmcp::model::Resource {
                raw: rmcp::model::RawResource {
                    uri: resource.uri.clone(),
                    name: resource.name.clone().unwrap_or_default(),
                    description: resource.description.clone(),
                    mime_type: resource.mime_type.clone(),
                    size: None,
                },
                annotations: None,
            }
        }).collect();
        Ok(rmcp_resources)
    }

    #[cfg(not(feature = "mcp"))]
    pub async fn list_resources(&self) -> McpResult<Vec<()>> {
        Ok(vec![])
    }

    /// List available prompts
    #[cfg(feature = "mcp")]
    pub async fn list_prompts(&self) -> McpResult<Vec<rmcp::model::Prompt>> {
        let simple_prompts = self.prompts.read().await;
        let rmcp_prompts = simple_prompts.iter().map(|prompt| {
            rmcp::model::Prompt {
                name: prompt.name.clone(),
                description: prompt.description.clone(),
                arguments: None, // SimpleTool doesn't have arguments field
            }
        }).collect();
        Ok(rmcp_prompts)
    }

    #[cfg(not(feature = "mcp"))]
    pub async fn list_prompts(&self) -> McpResult<Vec<()>> {
        Ok(vec![])
    }
}