fastmcp 0.0.0

use std::sync::Arc;
use std::time::Duration;

use dashmap::DashMap;
use tokio::sync::mpsc;
use tracing::{debug, error, info};

use crate::error::{Error, Result};
use crate::protocol::{ErrorMessage, Message, MessageType, Request, Response, SessionId};
use crate::tool::{BoxedTool, Tool, ToolContext};
use crate::transport::{Transport, TransportConfig, TransportMessage, create_transport};

/// Default timeout for tool execution
const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(30);

/// MCP Server
#[derive(Debug)]
pub struct McpServer {
    /// Registered tools
    tools: Arc<DashMap<String, BoxedTool>>,

    /// Active sessions
    sessions: Arc<DashMap<SessionId, Arc<ToolContext>>>,

    /// Default timeout for tool execution
    default_timeout: Duration,

    /// Transport for communication
    transport: Option<Box<dyn Transport>>,

    /// Debug mode
    debug_mode: bool,
}

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

impl McpServer {
    /// Create a new MCP server
    pub fn new() -> Self {
        Self {
            tools: Arc::new(DashMap::new()),
            sessions: Arc::new(DashMap::new()),
            default_timeout: DEFAULT_TOOL_TIMEOUT,
            transport: None,
            debug_mode: false,
        }
    }

    /// Create a new MCP server with a builder
    pub fn builder() -> McpServerBuilder {
        McpServerBuilder::new()
    }

    /// Register a tool
    pub fn register_tool<T: Tool + 'static>(&self, tool: T) {
        let name = tool.name().to_string();
        debug!("Registering tool: {}", name);
        self.tools.insert(name, Box::new(tool));
    }

    /// Set the transport
    pub fn set_transport(&mut self, transport: Box<dyn Transport>) {
        self.transport = Some(transport);
    }

    /// Start the MCP server
    pub async fn start(&mut self) -> Result<()> {
        info!("Starting MCP server");

        // Setup the transport
        let mut transport = self
            .transport
            .take() // 使用take而不是as_mut避免借用问题
            .ok_or_else(|| Error::Transport("No transport configured".to_string()))?;

        // Create channels for handling requests
        let (request_tx, mut request_rx) =
            mpsc::channel::<(Request, mpsc::Sender<TransportMessage>)>(100);

        // Share the tools and sessions maps using Arc
        let tools = self.tools.clone();
        let sessions = self.sessions.clone();
        let default_timeout = self.default_timeout;

        // Setup the request handler
        let request_tx_clone = request_tx.clone();
        let request_handler = Arc::new(move |request: Request| {
            // We'll create a channel to transfer the request and a response channel
            let (tx, _) = mpsc::channel(100);

            // Clone for the async block
            let _tx_clone = tx.clone();
            let request_tx = request_tx_clone.clone();

            // Spawn a task to handle this
            tokio::spawn(async move {
                // For each request, we create a channel for responses
                let (response_tx, _) = mpsc::channel(100);

                // Send the original request and the response channel
                if let Err(e) = request_tx.send((request, response_tx.clone())).await {
                    error!("Failed to send request: {}", e);
                }
            });

            tx
        });

        transport.set_request_handler(request_handler);

        // Spawn the transport
        let transport_handle = tokio::spawn(async move {
            if let Err(e) = transport.start().await {
                error!("Transport error: {}", e);
            }
        });

        // Process requests
        let tools_clone = tools.clone();
        let sessions_clone = sessions.clone();

        tokio::spawn(async move {
            while let Some((request, response_tx)) = request_rx.recv().await {
                let tools = tools_clone.clone();
                let sessions = sessions_clone.clone();

                // Process the request in a separate task
                tokio::spawn(async move {
                    if let Err(e) =
                        Self::handle_request(request, response_tx, tools, sessions, default_timeout)
                            .await
                    {
                        error!("Error handling request: {}", e);
                    }
                });
            }
        });

        // Wait for the transport to finish
        transport_handle
            .await
            .map_err(|e| Error::Transport(format!("Transport task failed: {e}")))?;

        Ok(())
    }

    /// Handle a single request
    async fn handle_request(
        request: Request, response_tx: mpsc::Sender<TransportMessage>,
        tools: Arc<DashMap<String, BoxedTool>>,
        sessions: Arc<DashMap<SessionId, Arc<ToolContext>>>, default_timeout: Duration,
    ) -> Result<()> {
        debug!("Handling request for tool: {}", request.tool);

        let tool_name = request.tool.clone();

        // Try to find the tool
        if !tools.contains_key(&tool_name) {
            // Send error response
            let error = ErrorMessage {
                message: Message {
                    id: request.message.id.clone(),
                    message_type: MessageType::Error,
                    session_id: request.message.session_id.clone(),
                    metadata: None,
                },
                code: 404,
                error: format!("Tool not found: {tool_name}"),
                details: None,
            };

            response_tx
                .send(TransportMessage::Error(error))
                .await
                .map_err(|_| Error::Transport("Failed to send error response".to_string()))?;

            return Ok(());
        }

        // Get or create a session
        let session_id = match &request.message.session_id {
            Some(id) => id.clone(),
            None => uuid::Uuid::new_v4().to_string(),
        };

        let context = if let Some(context) = sessions.get(&session_id) {
            context.clone()
        } else {
            // Create a new context
            let context = Arc::new(ToolContext::new(session_id.clone()));
            sessions.insert(session_id.clone(), context.clone());
            context
        };

        // 获取工具实例
        let tool_params = request.params.clone();
        let request_id = request.message.id.clone();

        // 获取调试模式状态
        let debug_mode = false; // 暂时设置为false，后续通过元数据传递

        if debug_mode {
            debug!("调试模式: 开始执行工具 {}", tool_name);
            debug!("调试模式: 请求ID {}", request_id);
            debug!("调试模式: 会话ID {}", session_id);
            debug!(
                "调试模式: 参数 {}",
                serde_json::to_string_pretty(&tool_params).unwrap_or_default()
            );
        }

        // 参数验证
        {
            let tool_ref = tools.get(&tool_name).unwrap();

            // 权限检查
            if let Err(permission_error) = tool_ref.check_permission(&context) {
                if debug_mode {
                    debug!("调试模式: 权限检查失败 {}", permission_error);
                }

                // 权限验证失败
                let error_message = ErrorMessage {
                    message: Message {
                        id: request_id.clone(),
                        message_type: MessageType::Error,
                        session_id: Some(session_id.clone()),
                        metadata: None,
                    },
                    code: 403,
                    error: format!("权限检查失败: {permission_error}"),
                    details: None,
                };

                response_tx
                    .send(TransportMessage::Error(error_message))
                    .await
                    .map_err(|_| Error::Transport("发送错误响应失败".to_string()))?;

                return Ok(());
            }

            // 资源限制检查
            if let Err(resource_error) = context.check_resource_limits() {
                if debug_mode {
                    debug!("调试模式: 资源限制检查失败 {}", resource_error);
                }

                // 资源限制验证失败
                let error_message = ErrorMessage {
                    message: Message {
                        id: request_id.clone(),
                        message_type: MessageType::Error,
                        session_id: Some(session_id.clone()),
                        metadata: None,
                    },
                    code: 429,
                    error: format!("资源限制检查失败: {resource_error}"),
                    details: None,
                };

                response_tx
                    .send(TransportMessage::Error(error_message))
                    .await
                    .map_err(|_| Error::Transport("发送错误响应失败".to_string()))?;

                return Ok(());
            }

            if let Err(validation_error) = tool_ref.validate_params(&tool_params) {
                // 参数验证失败
                let error_message = ErrorMessage {
                    message: Message {
                        id: request_id.clone(),
                        message_type: MessageType::Error,
                        session_id: Some(session_id.clone()),
                        metadata: None,
                    },
                    code: 400,
                    error: format!("参数验证失败: {validation_error}"),
                    details: None,
                };

                response_tx
                    .send(TransportMessage::Error(error_message))
                    .await
                    .map_err(|_| Error::Transport("发送错误响应失败".to_string()))?;

                return Ok(());
            }

            // 执行工具前的钩子
            if let Err(e) = tool_ref.before_execute(&context).await {
                // 前置钩子执行失败
                let error_message = ErrorMessage {
                    message: Message {
                        id: request_id.clone(),
                        message_type: MessageType::Error,
                        session_id: Some(session_id.clone()),
                        metadata: None,
                    },
                    code: 500,
                    error: format!("工具执行前处理失败: {e}"),
                    details: None,
                };

                response_tx
                    .send(TransportMessage::Error(error_message))
                    .await
                    .map_err(|_| Error::Transport("发送错误响应失败".to_string()))?;

                return Ok(());
            }
        }

        // 创建一个工具执行函数
        let tools_clone = tools.clone();
        let context_clone = context.clone();
        let tool_name_clone = tool_name.clone();
        let execute_tool = async move {
            // 记录API调用
            context_clone.record_api_call()?;

            let start_time = std::time::Instant::now();
            let tool_ref = tools_clone.get(&tool_name_clone).unwrap();

            // 执行工具
            let result = tool_ref.execute(tool_params, context_clone.clone()).await;

            // 记录CPU时间
            let cpu_time = start_time.elapsed().as_millis() as u64;
            context_clone.resource_usage().add_cpu_time(cpu_time);

            result
        };

        // 使用超时执行工具
        let execution_result = tokio::time::timeout(default_timeout, execute_tool).await;

        // 处理执行结果
        match execution_result {
            // 超时情况
            Err(_) => {
                // 超时错误
                let error_message = ErrorMessage {
                    message: Message {
                        id: request_id,
                        message_type: MessageType::Error,
                        session_id: Some(session_id.clone()),
                        metadata: None,
                    },
                    code: 408,
                    error: format!("工具 '{tool_name}' 执行超时"),
                    details: None,
                };

                response_tx
                    .send(TransportMessage::Error(error_message))
                    .await
                    .map_err(|_| Error::Transport("发送错误响应失败".to_string()))?;

                // 执行后钩子（超时情况）
                if let Some(tool_ref) = tools.get(&tool_name) {
                    if let Err(e) = tool_ref
                        .after_execute(&context, false, None, Some("执行超时"))
                        .await
                    {
                        error!("执行后钩子失败: {}", e);
                    }
                }

                return Ok(());
            }

            // 正常执行完成
            Ok(result) => {
                // 执行后钩子
                if let Some(tool_ref) = tools.get(&tool_name) {
                    let hook_result = match &result {
                        Ok(value) => {
                            tool_ref
                                .after_execute(&context, true, Some(value), None)
                                .await
                        }
                        Err(e) => {
                            tool_ref
                                .after_execute(&context, false, None, Some(&e.to_string()))
                                .await
                        }
                    };

                    if let Err(e) = hook_result {
                        error!("执行后钩子失败: {}", e);
                    }
                }

                // 处理工具执行结果
                match result {
                    Ok(result) => {
                        // Success response
                        let response = Response {
                            message: Message {
                                id: request_id,
                                message_type: MessageType::Response,
                                session_id: Some(session_id),
                                metadata: None,
                            },
                            data: result,
                            partial: None,
                        };

                        response_tx
                            .send(TransportMessage::Response(response))
                            .await
                            .map_err(|_| Error::Transport("发送响应失败".to_string()))?;
                    }
                    Err(error) => {
                        // Tool execution error
                        let error_message = ErrorMessage {
                            message: Message {
                                id: request_id,
                                message_type: MessageType::Error,
                                session_id: Some(session_id),
                                metadata: None,
                            },
                            code: 500,
                            error: error.to_string(),
                            details: None,
                        };

                        response_tx
                            .send(TransportMessage::Error(error_message))
                            .await
                            .map_err(|_| Error::Transport("发送错误响应失败".to_string()))?;
                    }
                };
            }
        }

        debug!("Tool execution completed");
        Ok(())
    }

    /// Enable debug mode
    pub fn enable_debug_mode(&mut self) {
        self.debug_mode = true;
        info!("调试模式已启用");
    }

    /// Disable debug mode
    pub fn disable_debug_mode(&mut self) {
        self.debug_mode = false;
        info!("调试模式已禁用");
    }

    /// Check if debug mode is enabled
    pub fn is_debug_mode(&self) -> bool {
        self.debug_mode
    }
}

/// Builder for MCP server
#[derive(Debug, Default)]
pub struct McpServerBuilder {
    /// Default timeout for tool execution
    default_timeout: Option<Duration>,

    /// Transport configuration
    transport_config: Option<TransportConfig>,

    /// Tools to register
    tools: Vec<BoxedTool>,

    /// Enable debug mode
    debug_mode: bool,
}

impl McpServerBuilder {
    /// Create a new builder
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the default timeout for tool execution
    pub fn with_default_timeout(mut self, timeout: Duration) -> Self {
        self.default_timeout = Some(timeout);
        self
    }

    /// Set the transport configuration
    pub fn with_transport(mut self, config: TransportConfig) -> Self {
        self.transport_config = Some(config);
        self
    }

    /// Add a tool to register
    pub fn with_tool<T: Tool + 'static>(mut self, tool: T) -> Self {
        self.tools.push(Box::new(tool));
        self
    }

    /// Enable debug mode
    pub fn with_debug_mode(mut self, enabled: bool) -> Self {
        self.debug_mode = enabled;
        self
    }

    /// Build the MCP server
    pub fn build(self) -> Result<McpServer> {
        let mut server = McpServer::new();

        // Set default timeout if provided
        if let Some(timeout) = self.default_timeout {
            server.default_timeout = timeout;
        }

        // Create and set transport if provided
        if let Some(config) = self.transport_config {
            let transport = create_transport(config)?;
            server.set_transport(transport);
        }

        // Register tools
        for tool in self.tools {
            let name = tool.name().to_string();
            server.tools.insert(name, tool);
        }

        // Set debug mode
        if self.debug_mode {
            server.enable_debug_mode();
        }

        Ok(server)
    }
}