p2p-foundation 0.1.6

//! Model Context Protocol (MCP) Server Implementation
//!
//! This module provides a fully-featured MCP server that integrates with the P2P network,
//! enabling AI agents to discover and call tools across the distributed network.
//! 
//! The implementation includes:
//! - MCP message routing over P2P transport
//! - Tool registration and discovery through DHT
//! - Security and authentication for remote calls
//! - Service health monitoring and load balancing

pub mod security;

use crate::dht::{Key, DHT};
use crate::{PeerId, Result, P2PError};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime, Instant};
use tokio::sync::{RwLock, mpsc, oneshot};
use tokio::time::timeout;
use tracing::{debug, info};
use rand;

pub use security::*;

/// MCP protocol version
pub const MCP_VERSION: &str = "2024-11-05";

/// Maximum message size for MCP calls (1MB)
pub const MAX_MESSAGE_SIZE: usize = 1024 * 1024;

/// Default timeout for MCP calls
pub const DEFAULT_CALL_TIMEOUT: Duration = Duration::from_secs(30);

/// MCP protocol identifier for P2P messaging
pub const MCP_PROTOCOL: &str = "/p2p-foundation/mcp/1.0.0";

/// Service discovery refresh interval
pub const SERVICE_DISCOVERY_INTERVAL: Duration = Duration::from_secs(60);

/// MCP message types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MCPMessage {
    /// Initialize MCP session
    Initialize {
        /// MCP protocol version being used
        protocol_version: String,
        /// Client capabilities for this session
        capabilities: MCPCapabilities,
        /// Information about the connecting client
        client_info: MCPClientInfo,
    },
    /// Initialize response
    InitializeResult {
        /// MCP protocol version the server supports
        protocol_version: String,
        /// Server capabilities for this session
        capabilities: MCPCapabilities,
        /// Information about the MCP server
        server_info: MCPServerInfo,
    },
    /// List available tools
    ListTools {
        /// Pagination cursor for large tool lists
        cursor: Option<String>,
    },
    /// List tools response
    ListToolsResult {
        /// Available tools on this server
        tools: Vec<MCPTool>,
        /// Next pagination cursor if more tools available
        next_cursor: Option<String>,
    },
    /// Call a tool
    CallTool {
        /// Name of the tool to call
        name: String,
        /// Arguments to pass to the tool
        arguments: Value,
    },
    /// Tool call response
    CallToolResult {
        /// Content returned by the tool
        content: Vec<MCPContent>,
        /// Whether the call resulted in an error
        is_error: bool,
    },
    /// List available prompts
    ListPrompts {
        /// Pagination cursor for large prompt lists
        cursor: Option<String>,
    },
    /// List prompts response
    ListPromptsResult {
        /// Available prompts on this server
        prompts: Vec<MCPPrompt>,
        /// Next pagination cursor if more prompts available
        next_cursor: Option<String>,
    },
    /// Get a prompt
    GetPrompt {
        /// Name of the prompt to retrieve
        name: String,
        /// Arguments to customize the prompt
        arguments: Option<Value>,
    },
    /// Get prompt response
    GetPromptResult {
        /// Description of the prompt
        description: Option<String>,
        /// Prompt messages/content
        messages: Vec<MCPPromptMessage>,
    },
    /// List available resources
    ListResources {
        /// Pagination cursor for large resource lists
        cursor: Option<String>,
    },
    /// List resources response
    ListResourcesResult {
        /// Available resources on this server
        resources: Vec<MCPResource>,
        /// Next pagination cursor if more resources available
        next_cursor: Option<String>,
    },
    /// Read a resource
    ReadResource {
        /// URI of the resource to read
        uri: String,
    },
    /// Read resource response
    ReadResourceResult {
        /// Contents of the requested resource
        contents: Vec<MCPResourceContent>,
    },
    /// Subscribe to resource
    SubscribeResource {
        /// URI of the resource to subscribe to
        uri: String,
    },
    /// Unsubscribe from resource
    UnsubscribeResource {
        /// URI of the resource to unsubscribe from
        uri: String,
    },
    /// Resource updated notification
    ResourceUpdated {
        /// URI of the resource that was updated
        uri: String,
    },
    /// List logs
    ListLogs {
        /// Pagination cursor for large log lists
        cursor: Option<String>,
    },
    /// List logs response
    ListLogsResult {
        /// Log entries available on this server
        logs: Vec<MCPLogEntry>,
        /// Next pagination cursor if more logs available
        next_cursor: Option<String>,
    },
    /// Set log level
    SetLogLevel {
        /// Log level to set for the server
        level: MCPLogLevel,
    },
    /// Error response
    Error {
        /// Error code identifying the type of error
        code: i32,
        /// Human-readable error message
        message: String,
        /// Optional additional error data
        data: Option<Value>,
    },
}

/// MCP capabilities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPCapabilities {
    /// Experimental capabilities
    pub experimental: Option<Value>,
    /// Sampling capability
    pub sampling: Option<Value>,
    /// Tools capability
    pub tools: Option<MCPToolsCapability>,
    /// Prompts capability
    pub prompts: Option<MCPPromptsCapability>,
    /// Resources capability
    pub resources: Option<MCPResourcesCapability>,
    /// Logging capability
    pub logging: Option<MCPLoggingCapability>,
}

/// Tools capability configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPToolsCapability {
    /// Whether tools are supported
    pub list_changed: Option<bool>,
}

/// Prompts capability configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPPromptsCapability {
    /// Whether prompts are supported
    pub list_changed: Option<bool>,
}

/// Resources capability configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPResourcesCapability {
    /// Whether resources are supported
    pub subscribe: Option<bool>,
    /// Whether resource listing is supported
    pub list_changed: Option<bool>,
}

/// Logging capability configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPLoggingCapability {
    /// Available log levels
    pub levels: Option<Vec<MCPLogLevel>>,
}

/// MCP client information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPClientInfo {
    /// Client name
    pub name: String,
    /// Client version
    pub version: String,
}

/// MCP server information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPServerInfo {
    /// Server name
    pub name: String,
    /// Server version
    pub version: String,
}

/// MCP tool definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPTool {
    /// Tool name
    pub name: String,
    /// Tool description
    pub description: String,
    /// Input schema (JSON Schema)
    pub input_schema: Value,
}

/// MCP tool implementation
pub struct Tool {
    /// Tool definition
    pub definition: MCPTool,
    /// Tool handler function
    pub handler: Box<dyn ToolHandler + Send + Sync>,
    /// Tool metadata
    pub metadata: ToolMetadata,
}

/// Tool metadata for tracking and optimization
#[derive(Debug, Clone)]
pub struct ToolMetadata {
    /// Tool creation time
    pub created_at: SystemTime,
    /// Last call time
    pub last_called: Option<SystemTime>,
    /// Total number of calls
    pub call_count: u64,
    /// Average execution time
    pub avg_execution_time: Duration,
    /// Tool health status
    pub health_status: ToolHealthStatus,
    /// Tags for categorization
    pub tags: Vec<String>,
}

/// Tool health status
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ToolHealthStatus {
    /// Tool is healthy and responsive
    Healthy,
    /// Tool is experiencing issues
    Degraded,
    /// Tool is not responding
    Unhealthy,
    /// Tool is disabled
    Disabled,
}

/// Tool handler trait
pub trait ToolHandler {
    /// Execute the tool with given arguments
    fn execute(&self, arguments: Value) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Value>> + Send + '_>>;
    
    /// Validate tool arguments
    fn validate(&self, arguments: &Value) -> Result<()> {
        // Default implementation - no validation
        let _ = arguments;
        Ok(())
    }
    
    /// Get tool resource requirements
    fn get_requirements(&self) -> ToolRequirements {
        ToolRequirements::default()
    }
}

/// Tool resource requirements
#[derive(Debug, Clone)]
pub struct ToolRequirements {
    /// Maximum memory usage in bytes
    pub max_memory: Option<u64>,
    /// Maximum execution time allowed for tool calls
    pub max_execution_time: Option<Duration>,
    /// Required capabilities that must be available
    pub required_capabilities: Vec<String>,
    /// Whether this tool requires network access
    pub requires_network: bool,
    /// Whether this tool requires file system access
    pub requires_filesystem: bool,
}

impl Default for ToolRequirements {
    fn default() -> Self {
        Self {
            max_memory: Some(100 * 1024 * 1024), // 100MB default
            max_execution_time: Some(Duration::from_secs(30)),
            required_capabilities: Vec::new(),
            requires_network: false,
            requires_filesystem: false,
        }
    }
}

/// MCP content types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MCPContent {
    /// Text content
    Text {
        /// The text content
        text: String,
    },
    /// Image content
    Image {
        /// Base64-encoded image data
        data: String,
        /// MIME type of the image
        mime_type: String,
    },
    /// Resource content
    Resource {
        /// Reference to an MCP resource
        resource: MCPResourceReference,
    },
}

/// MCP resource reference
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPResourceReference {
    /// Resource URI
    pub uri: String,
    /// Resource type
    pub type_: Option<String>,
}

/// MCP prompt definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPPrompt {
    /// Prompt name
    pub name: String,
    /// Prompt description
    pub description: Option<String>,
    /// Prompt arguments schema
    pub arguments: Option<Value>,
}

/// MCP prompt message
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPPromptMessage {
    /// Message role
    pub role: MCPRole,
    /// Message content
    pub content: MCPContent,
}

/// MCP message roles
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MCPRole {
    /// User message
    User,
    /// Assistant message
    Assistant,
    /// System message
    System,
}

/// MCP resource definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPResource {
    /// Resource URI
    pub uri: String,
    /// Resource name
    pub name: String,
    /// Resource description
    pub description: Option<String>,
    /// Resource MIME type
    pub mime_type: Option<String>,
}

/// MCP resource content
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPResourceContent {
    /// Content URI
    pub uri: String,
    /// Content MIME type
    pub mime_type: String,
    /// Content data
    pub text: Option<String>,
    /// Binary content (base64 encoded)
    pub blob: Option<String>,
}

/// MCP log levels
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum MCPLogLevel {
    /// Debug level
    Debug,
    /// Info level
    Info,
    /// Notice level
    Notice,
    /// Warning level
    Warning,
    /// Error level
    Error,
    /// Critical level
    Critical,
    /// Alert level
    Alert,
    /// Emergency level
    Emergency,
}

/// MCP log entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPLogEntry {
    /// Log level
    pub level: MCPLogLevel,
    /// Log message
    pub data: Value,
    /// Logger name
    pub logger: Option<String>,
}

/// MCP service descriptor for discovery
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPService {
    /// Service identifier
    pub service_id: String,
    /// Node providing the service
    pub node_id: PeerId,
    /// Available tools
    pub tools: Vec<String>,
    /// Service capabilities
    pub capabilities: MCPCapabilities,
    /// Service metadata
    pub metadata: MCPServiceMetadata,
    /// Service registration time
    pub registered_at: SystemTime,
    /// Service endpoint information
    pub endpoint: MCPEndpoint,
}

/// MCP service metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPServiceMetadata {
    /// Service name
    pub name: String,
    /// Service version
    pub version: String,
    /// Service description
    pub description: Option<String>,
    /// Service tags
    pub tags: Vec<String>,
    /// Service health status
    pub health_status: ServiceHealthStatus,
    /// Service load metrics
    pub load_metrics: ServiceLoadMetrics,
}

/// Service health status
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ServiceHealthStatus {
    /// Service is healthy
    Healthy,
    /// Service is degraded
    Degraded,
    /// Service is unhealthy
    Unhealthy,
    /// Service is maintenance mode
    Maintenance,
}

/// Service load metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceLoadMetrics {
    /// Current active requests
    pub active_requests: u32,
    /// Requests per second
    pub requests_per_second: f64,
    /// Average response time in milliseconds
    pub avg_response_time_ms: f64,
    /// Error rate (0.0-1.0)
    pub error_rate: f64,
    /// CPU usage percentage
    pub cpu_usage: f64,
    /// Memory usage in bytes
    pub memory_usage: u64,
}

/// MCP endpoint information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPEndpoint {
    /// Endpoint protocol (p2p, http, etc.)
    pub protocol: String,
    /// Endpoint address
    pub address: String,
    /// Endpoint port
    pub port: Option<u16>,
    /// TLS enabled
    pub tls: bool,
    /// Authentication required
    pub auth_required: bool,
}

/// MCP request with routing information
#[derive(Debug, Clone)]
pub struct MCPRequest {
    /// Request ID
    pub request_id: String,
    /// Source peer
    pub source_peer: PeerId,
    /// Target peer
    pub target_peer: PeerId,
    /// MCP message
    pub message: MCPMessage,
    /// Request timestamp
    pub timestamp: SystemTime,
    /// Request timeout
    pub timeout: Duration,
    /// Authentication token
    pub auth_token: Option<String>,
}

/// P2P MCP message wrapper for network transmission
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct P2PMCPMessage {
    /// Message type
    pub message_type: P2PMCPMessageType,
    /// Request/Response ID for correlation
    pub message_id: String,
    /// Source peer ID
    pub source_peer: PeerId,
    /// Target peer ID (optional for broadcasts)
    pub target_peer: Option<PeerId>,
    /// Timestamp
    pub timestamp: u64,
    /// MCP message payload
    pub payload: MCPMessage,
    /// Message TTL for routing
    pub ttl: u8,
}

/// P2P MCP message types
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum P2PMCPMessageType {
    /// Direct request to a specific peer
    Request,
    /// Response to a request
    Response,
    /// Service advertisement
    ServiceAdvertisement,
    /// Service discovery query
    ServiceDiscovery,
}

/// MCP response with metadata
#[derive(Debug, Clone)]
pub struct MCPResponse {
    /// Request ID this response corresponds to
    pub request_id: String,
    /// Response message
    pub message: MCPMessage,
    /// Response timestamp
    pub timestamp: SystemTime,
    /// Processing time
    pub processing_time: Duration,
}

/// MCP call context
#[derive(Debug, Clone)]
pub struct MCPCallContext {
    /// Caller peer ID
    pub caller_id: PeerId,
    /// Call timestamp
    pub timestamp: SystemTime,
    /// Call timeout
    pub timeout: Duration,
    /// Authentication information
    pub auth_info: Option<MCPAuthInfo>,
    /// Call metadata
    pub metadata: HashMap<String, String>,
}

/// MCP authentication information
#[derive(Debug, Clone)]
pub struct MCPAuthInfo {
    /// Authentication token
    pub token: String,
    /// Token type
    pub token_type: String,
    /// Token expiration
    pub expires_at: Option<SystemTime>,
    /// Granted permissions
    pub permissions: Vec<String>,
}

/// MCP server configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MCPServerConfig {
    /// Server name
    pub server_name: String,
    /// Server version
    pub server_version: String,
    /// Enable tool discovery via DHT
    pub enable_dht_discovery: bool,
    /// Maximum concurrent requests
    pub max_concurrent_requests: usize,
    /// Request timeout
    pub request_timeout: Duration,
    /// Enable authentication
    pub enable_auth: bool,
    /// Enable rate limiting
    pub enable_rate_limiting: bool,
    /// Rate limit: requests per minute
    pub rate_limit_rpm: u32,
    /// Enable request logging
    pub enable_logging: bool,
    /// Maximum tool execution time
    pub max_tool_execution_time: Duration,
    /// Tool memory limit
    pub tool_memory_limit: u64,
}

impl Default for MCPServerConfig {
    fn default() -> Self {
        Self {
            server_name: "P2P-MCP-Server".to_string(),
            server_version: crate::VERSION.to_string(),
            enable_dht_discovery: true,
            max_concurrent_requests: 100,
            request_timeout: DEFAULT_CALL_TIMEOUT,
            enable_auth: true,
            enable_rate_limiting: true,
            rate_limit_rpm: 60,
            enable_logging: true,
            max_tool_execution_time: Duration::from_secs(30),
            tool_memory_limit: 100 * 1024 * 1024, // 100MB
        }
    }
}

/// Main MCP server implementation
pub struct MCPServer {
    /// Server configuration
    config: MCPServerConfig,
    /// Registered tools
    tools: Arc<RwLock<HashMap<String, Tool>>>,
    /// Registered prompts (reserved for future implementation)
    #[allow(dead_code)]
    prompts: Arc<RwLock<HashMap<String, MCPPrompt>>>,
    /// Registered resources (reserved for future implementation)
    #[allow(dead_code)]
    resources: Arc<RwLock<HashMap<String, MCPResource>>>,
    /// Active sessions
    sessions: Arc<RwLock<HashMap<String, MCPSession>>>,
    /// Request handlers
    request_handlers: Arc<RwLock<HashMap<String, oneshot::Sender<MCPResponse>>>>,
    /// DHT reference for service discovery
    dht: Option<Arc<RwLock<DHT>>>,
    /// Local service registry
    local_services: Arc<RwLock<HashMap<String, MCPService>>>,
    /// Remote service cache
    remote_services: Arc<RwLock<HashMap<String, MCPService>>>,
    /// Request statistics
    stats: Arc<RwLock<MCPServerStats>>,
    /// Message channel for incoming requests
    request_tx: mpsc::UnboundedSender<MCPRequest>,
    /// Message channel for outgoing responses (reserved for future implementation)
    #[allow(dead_code)]
    response_rx: Arc<RwLock<mpsc::UnboundedReceiver<MCPResponse>>>,
    /// Security manager
    security_manager: Option<Arc<MCPSecurityManager>>,
    /// Audit logger
    audit_logger: Arc<SecurityAuditLogger>,
}

/// MCP session information
#[derive(Debug, Clone)]
pub struct MCPSession {
    /// Session ID
    pub session_id: String,
    /// Peer ID
    pub peer_id: PeerId,
    /// Client capabilities
    pub client_capabilities: Option<MCPCapabilities>,
    /// Session start time
    pub started_at: SystemTime,
    /// Last activity time
    pub last_activity: SystemTime,
    /// Session state
    pub state: MCPSessionState,
    /// Subscribed resources
    pub subscribed_resources: Vec<String>,
}

/// MCP session state
#[derive(Debug, Clone, PartialEq)]
pub enum MCPSessionState {
    /// Session is initializing
    Initializing,
    /// Session is active
    Active,
    /// Session is inactive
    Inactive,
    /// Session is terminated
    Terminated,
}

/// MCP server statistics
#[derive(Debug, Clone)]
pub struct MCPServerStats {
    /// Total requests processed
    pub total_requests: u64,
    /// Total responses sent
    pub total_responses: u64,
    /// Total errors
    pub total_errors: u64,
    /// Average response time
    pub avg_response_time: Duration,
    /// Active sessions
    pub active_sessions: u32,
    /// Total tools registered
    pub total_tools: u32,
    /// Most called tools
    pub popular_tools: HashMap<String, u64>,
    /// Server start time
    pub server_started_at: SystemTime,
}

impl Default for MCPServerStats {
    fn default() -> Self {
        Self {
            total_requests: 0,
            total_responses: 0,
            total_errors: 0,
            avg_response_time: Duration::from_millis(0),
            active_sessions: 0,
            total_tools: 0,
            popular_tools: HashMap::new(),
            server_started_at: SystemTime::now(),
        }
    }
}

impl MCPServer {
    /// Create a new MCP server
    pub fn new(config: MCPServerConfig) -> Self {
        let (request_tx, _request_rx) = mpsc::unbounded_channel();
        let (_response_tx, response_rx) = mpsc::unbounded_channel();
        
        // Initialize security manager if authentication is enabled
        let security_manager = if config.enable_auth {
            // Generate a random secret key for token signing
            let secret_key = (0..32).map(|_| rand::random::<u8>()).collect();
            Some(Arc::new(MCPSecurityManager::new(secret_key, config.rate_limit_rpm)))
        } else {
            None
        };
        
        let server = Self {
            config,
            tools: Arc::new(RwLock::new(HashMap::new())),
            prompts: Arc::new(RwLock::new(HashMap::new())),
            resources: Arc::new(RwLock::new(HashMap::new())),
            sessions: Arc::new(RwLock::new(HashMap::new())),
            request_handlers: Arc::new(RwLock::new(HashMap::new())),
            dht: None,
            local_services: Arc::new(RwLock::new(HashMap::new())),
            remote_services: Arc::new(RwLock::new(HashMap::new())),
            stats: Arc::new(RwLock::new(MCPServerStats::default())),
            request_tx,
            response_rx: Arc::new(RwLock::new(response_rx)),
            security_manager,
            audit_logger: Arc::new(SecurityAuditLogger::new(10000)), // Keep 10k audit entries
        };
        
        server
    }
    
    /// Create MCP server with DHT integration
    pub fn with_dht(mut self, dht: Arc<RwLock<DHT>>) -> Self {
        self.dht = Some(dht);
        self
    }
    
    /// Start the MCP server
    pub async fn start(&self) -> Result<()> {
        info!("Starting MCP server: {}", self.config.server_name);
        
        // Start request processing task
        self.start_request_processor().await?;
        
        // Start service discovery if DHT is available
        if self.dht.is_some() {
            self.start_service_discovery().await?;
        }
        
        // Start health monitoring
        self.start_health_monitor().await?;
        
        info!("MCP server started successfully");
        Ok(())
    }
    
    /// Register a tool
    pub async fn register_tool(&self, tool: Tool) -> Result<()> {
        let tool_name = tool.definition.name.clone();
        
        // Validate tool
        self.validate_tool(&tool).await?;
        
        // Register locally
        {
            let mut tools = self.tools.write().await;
            tools.insert(tool_name.clone(), tool);
        }
        
        // Update statistics
        {
            let mut stats = self.stats.write().await;
            stats.total_tools += 1;
        }
        
        // Register in DHT if available
        if let Some(dht) = &self.dht {
            self.register_tool_in_dht(&tool_name, dht).await?;
        }
        
        info!("Registered tool: {}", tool_name);
        Ok(())
    }
    
    /// Validate tool before registration
    async fn validate_tool(&self, tool: &Tool) -> Result<()> {
        // Check for duplicate names
        let tools = self.tools.read().await;
        if tools.contains_key(&tool.definition.name) {
            return Err(P2PError::MCP(format!("Tool already exists: {}", tool.definition.name)).into());
        }
        
        // Validate tool name
        if tool.definition.name.is_empty() || tool.definition.name.len() > 100 {
            return Err(P2PError::MCP("Invalid tool name".to_string()).into());
        }
        
        // Validate schema
        if !tool.definition.input_schema.is_object() {
            return Err(P2PError::MCP("Tool input schema must be an object".to_string()).into());
        }
        
        Ok(())
    }
    
    /// Register tool in DHT for discovery
    async fn register_tool_in_dht(&self, tool_name: &str, dht: &Arc<RwLock<DHT>>) -> Result<()> {
        let key = Key::new(format!("mcp:tool:{}", tool_name).as_bytes());
        let service_info = json!({
            "tool_name": tool_name,
            "node_id": "local_node", // TODO: Get actual node ID
            "registered_at": SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_err(|e| P2PError::Network(format!("Time error: {}", e)))?.as_secs(),
            "capabilities": self.get_server_capabilities().await
        });
        
        let dht_guard = dht.read().await;
        dht_guard.put(key, serde_json::to_vec(&service_info)?).await?;
        
        Ok(())
    }
    
    /// Get server capabilities
    async fn get_server_capabilities(&self) -> MCPCapabilities {
        MCPCapabilities {
            experimental: None,
            sampling: None,
            tools: Some(MCPToolsCapability {
                list_changed: Some(true),
            }),
            prompts: Some(MCPPromptsCapability {
                list_changed: Some(true),
            }),
            resources: Some(MCPResourcesCapability {
                subscribe: Some(true),
                list_changed: Some(true),
            }),
            logging: Some(MCPLoggingCapability {
                levels: Some(vec![
                    MCPLogLevel::Debug,
                    MCPLogLevel::Info,
                    MCPLogLevel::Warning,
                    MCPLogLevel::Error,
                ]),
            }),
        }
    }
    
    /// Call a tool by name
    pub async fn call_tool(&self, tool_name: &str, arguments: Value, context: MCPCallContext) -> Result<Value> {
        let start_time = Instant::now();
        
        // Security checks
        
        // 1. Check rate limiting
        if !self.check_rate_limit(&context.caller_id).await? {
            return Err(P2PError::MCP("Rate limit exceeded".to_string()));
        }
        
        // 2. Check tool execution permission
        if !self.check_permission(&context.caller_id, &MCPPermission::ExecuteTools).await? {
            return Err(P2PError::MCP("Permission denied: execute tools".to_string()));
        }
        
        // 3. Check tool-specific security policy
        let tool_security_level = self.get_tool_security_policy(tool_name).await;
        let is_trusted = self.is_trusted_peer(&context.caller_id).await;
        
        match tool_security_level {
            SecurityLevel::Admin => {
                if !self.check_permission(&context.caller_id, &MCPPermission::Admin).await? {
                    return Err(P2PError::MCP("Permission denied: admin access required".to_string()));
                }
            }
            SecurityLevel::Strong => {
                if !is_trusted {
                    return Err(P2PError::MCP("Permission denied: trusted peer required".to_string()));
                }
            }
            SecurityLevel::Basic => {
                // Check if authentication is enabled and token is valid
                if self.config.enable_auth {
                    if let Some(auth_info) = &context.auth_info {
                        self.verify_auth_token(&auth_info.token).await?;
                    } else {
                        return Err(P2PError::MCP("Authentication required".to_string()));
                    }
                }
            }
            SecurityLevel::Public => {
                // No additional checks needed
            }
        }
        
        // Log the tool call attempt
        let mut details = HashMap::new();
        details.insert("action".to_string(), "tool_call".to_string());
        details.insert("tool_name".to_string(), tool_name.to_string());
        details.insert("security_level".to_string(), format!("{:?}", tool_security_level));
        
        self.audit_logger.log_event(
            "tool_execution".to_string(),
            context.caller_id.clone(),
            details,
            AuditSeverity::Info,
        ).await;
        
        // Check if tool exists
        let tool_exists = {
            let tools = self.tools.read().await;
            tools.contains_key(tool_name)
        };
        
        if !tool_exists {
            return Err(P2PError::MCP(format!("Tool not found: {}", tool_name)).into());
        }
        
        // Validate arguments and get requirements
        let requirements = {
            let tools = self.tools.read().await;
            let tool = tools.get(tool_name).unwrap(); // Safe because we checked exists above
            
            // Validate arguments
            if let Err(e) = tool.handler.validate(&arguments) {
                return Err(P2PError::MCP(format!("Tool validation failed: {}", e)).into());
            }
            
            // Get resource requirements
            tool.handler.get_requirements()
        };
        
        // Check resource requirements
        self.check_resource_requirements(&requirements).await?;
        
        // Execute tool in a spawned task to avoid borrow checker issues
        let tools_clone = self.tools.clone();
        let tool_name_owned = tool_name.to_string();
        let execution_timeout = context.timeout.min(requirements.max_execution_time.unwrap_or(context.timeout));
        
        let result = timeout(execution_timeout, async move {
            let tools = tools_clone.read().await;
            let tool = tools.get(&tool_name_owned).unwrap(); // Safe because we checked exists above
            tool.handler.execute(arguments).await
        }).await
        .map_err(|_| P2PError::MCP("Tool execution timeout".to_string()))?
        .map_err(|e| P2PError::MCP(format!("Tool execution failed: {}", e)))?;
        
        let execution_time = start_time.elapsed();
        
        // Update tool statistics
        self.update_tool_stats(tool_name, execution_time, true).await;
        
        // Update server statistics
        {
            let mut stats = self.stats.write().await;
            stats.total_requests += 1;
            stats.total_responses += 1;
            
            // Update average response time
            let new_total_time = stats.avg_response_time.mul_f64(stats.total_responses as f64 - 1.0) + execution_time;
            stats.avg_response_time = new_total_time.div_f64(stats.total_responses as f64);
            
            // Update popular tools
            *stats.popular_tools.entry(tool_name.to_string()).or_insert(0) += 1;
        }
        
        debug!("Tool '{}' executed in {:?}", tool_name, execution_time);
        Ok(result)
    }
    
    /// Check if resource requirements can be met
    async fn check_resource_requirements(&self, requirements: &ToolRequirements) -> Result<()> {
        // Check memory limit
        if let Some(max_memory) = requirements.max_memory {
            if max_memory > self.config.tool_memory_limit {
                return Err(P2PError::MCP("Tool memory requirement exceeds limit".to_string()).into());
            }
        }
        
        // Check execution time limit
        if let Some(max_execution_time) = requirements.max_execution_time {
            if max_execution_time > self.config.max_tool_execution_time {
                return Err(P2PError::MCP("Tool execution time requirement exceeds limit".to_string()).into());
            }
        }
        
        // TODO: Check other requirements (capabilities, network, filesystem)
        
        Ok(())
    }
    
    /// Update tool execution statistics
    async fn update_tool_stats(&self, tool_name: &str, execution_time: Duration, success: bool) {
        let mut tools = self.tools.write().await;
        if let Some(tool) = tools.get_mut(tool_name) {
            tool.metadata.call_count += 1;
            tool.metadata.last_called = Some(SystemTime::now());
            
            // Update average execution time
            let new_total_time = tool.metadata.avg_execution_time.mul_f64(tool.metadata.call_count as f64 - 1.0) + execution_time;
            tool.metadata.avg_execution_time = new_total_time.div_f64(tool.metadata.call_count as f64);
            
            // Update health status based on success
            if !success {
                tool.metadata.health_status = match tool.metadata.health_status {
                    ToolHealthStatus::Healthy => ToolHealthStatus::Degraded,
                    ToolHealthStatus::Degraded => ToolHealthStatus::Unhealthy,
                    other => other,
                };
            } else if tool.metadata.health_status != ToolHealthStatus::Disabled {
                tool.metadata.health_status = ToolHealthStatus::Healthy;
            }
        }
    }
    
    /// List available tools
    pub async fn list_tools(&self, _cursor: Option<String>) -> Result<(Vec<MCPTool>, Option<String>)> {
        let tools = self.tools.read().await;
        let tool_definitions: Vec<MCPTool> = tools.values()
            .map(|tool| tool.definition.clone())
            .collect();
        
        // For simplicity, return all tools without pagination
        // In a real implementation, you'd implement proper cursor-based pagination
        Ok((tool_definitions, None))
    }
    
    /// Start request processing task
    async fn start_request_processor(&self) -> Result<()> {
        let _request_tx = self.request_tx.clone();
        let _server_clone = Arc::new(self);
        
        tokio::spawn(async move {
            info!("MCP request processor started");
            
            // In a real implementation, this would listen for incoming MCP requests
            // and process them through a receiver channel. For now, we'll implement
            // the message handling infrastructure without the actual network loop.
            
            loop {
                // Sleep to prevent busy loop - in real implementation,
                // this would block on receiving messages
                tokio::time::sleep(Duration::from_millis(100)).await;
                
                // Check if we should shutdown
                // This is a placeholder - real implementation would have proper shutdown signaling
                break;
            }
            
            info!("MCP request processor stopped");
        });
        
        Ok(())
    }
    
    /// Start service discovery task
    async fn start_service_discovery(&self) -> Result<()> {
        if let Some(dht) = self.dht.clone() {
            let _stats = self.stats.clone();
            let remote_services = self.remote_services.clone();
            
            tokio::spawn(async move {
                info!("MCP service discovery started");
                
                loop {
                    // Periodically discover services
                    tokio::time::sleep(SERVICE_DISCOVERY_INTERVAL).await;
                    
                    // Query DHT for MCP services
                    let key = Key::new(b"mcp:services");
                    let dht_guard = dht.read().await;
                    
                    match dht_guard.get(&key).await {
                        Some(record) => {
                            match serde_json::from_slice::<Vec<MCPService>>(&record.value) {
                                Ok(services) => {
                                    debug!("Discovered {} MCP services", services.len());
                                    
                                    // Update remote services cache
                                    {
                                        let mut remote_cache = remote_services.write().await;
                                        for service in services {
                                            remote_cache.insert(service.service_id.clone(), service);
                                        }
                                    }
                                }
                                Err(e) => {
                                    debug!("Failed to deserialize services: {}", e);
                                }
                            }
                        }
                        None => {
                            debug!("No MCP services found in DHT");
                        }
                    }
                }
            });
        }
        
        Ok(())
    }
    
    /// Start health monitoring task
    async fn start_health_monitor(&self) -> Result<()> {
        // TODO: Implement health monitoring
        // This would check tool health and update status
        Ok(())
    }
    
    /// Get server statistics
    pub async fn get_stats(&self) -> MCPServerStats {
        self.stats.read().await.clone()
    }
    
    /// Discover remote services in the network
    pub async fn discover_remote_services(&self) -> Result<Vec<MCPService>> {
        if let Some(dht) = &self.dht {
            let key = Key::new(b"mcp:services");
            let dht_guard = dht.read().await;
            
            match dht_guard.get(&key).await {
                Some(record) => {
                    match serde_json::from_slice::<Vec<MCPService>>(&record.value) {
                        Ok(services) => {
                            // Update remote services cache
                            {
                                let mut remote_services = self.remote_services.write().await;
                                for service in &services {
                                    remote_services.insert(service.service_id.clone(), service.clone());
                                }
                            }
                            Ok(services)
                        }
                        Err(e) => {
                            debug!("Failed to deserialize services: {}", e);
                            Ok(Vec::new())
                        }
                    }
                }
                None => Ok(Vec::new()),
            }
        } else {
            Ok(Vec::new())
        }
    }
    
    /// Call a tool on a remote node
    pub async fn call_remote_tool(&self, peer_id: &PeerId, tool_name: &str, arguments: Value, context: MCPCallContext) -> Result<Value> {
        let request_id = uuid::Uuid::new_v4().to_string();
        
        // Create MCP call tool message
        let mcp_message = MCPMessage::CallTool {
            name: tool_name.to_string(),
            arguments,
        };
        
        // Create P2P message wrapper
        let p2p_message = P2PMCPMessage {
            message_type: P2PMCPMessageType::Request,
            message_id: request_id.clone(),
            source_peer: context.caller_id.clone(),
            target_peer: Some(peer_id.clone()),
            timestamp: SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map_err(|e| P2PError::Network(format!("Time error: {}", e)))?
                .as_secs(),
            payload: mcp_message,
            ttl: 5, // Max 5 hops
        };
        
        // Serialize the message
        let message_data = serde_json::to_vec(&p2p_message)
            .map_err(|e| P2PError::Serialization(e))?;
        
        if message_data.len() > MAX_MESSAGE_SIZE {
            return Err(P2PError::MCP("Message too large".to_string()));
        }
        
        // Create response channel
        let (response_tx, _response_rx) = oneshot::channel::<MCPResponse>();
        
        // Store response handler
        {
            let mut handlers = self.request_handlers.write().await;
            handlers.insert(request_id.clone(), response_tx);
        }
        
        // Send via P2P network - this will need to be connected to the network layer
        // For now, return an error indicating the need for network integration
        Err(P2PError::MCP("Remote tool calling requires P2P network integration".to_string()))
    }
    
    /// Handle incoming P2P MCP message
    pub async fn handle_p2p_message(&self, message_data: &[u8], source_peer: &PeerId) -> Result<Option<Vec<u8>>> {
        // Deserialize the P2P message
        let p2p_message: P2PMCPMessage = serde_json::from_slice(message_data)
            .map_err(|e| P2PError::Serialization(e))?;
        
        debug!("Received MCP message from {}: {:?}", source_peer, p2p_message.message_type);
        
        match p2p_message.message_type {
            P2PMCPMessageType::Request => {
                self.handle_remote_request(p2p_message).await
            }
            P2PMCPMessageType::Response => {
                self.handle_remote_response(p2p_message).await?;
                Ok(None) // Responses don't generate replies
            }
            P2PMCPMessageType::ServiceAdvertisement => {
                self.handle_service_advertisement(p2p_message).await?;
                Ok(None)
            }
            P2PMCPMessageType::ServiceDiscovery => {
                self.handle_service_discovery(p2p_message).await
            }
        }
    }
    
    /// Handle remote tool call request
    async fn handle_remote_request(&self, message: P2PMCPMessage) -> Result<Option<Vec<u8>>> {
        match message.payload {
            MCPMessage::CallTool { name, arguments } => {
                let context = MCPCallContext {
                    caller_id: message.source_peer.clone(),
                    timestamp: SystemTime::now(),
                    timeout: DEFAULT_CALL_TIMEOUT,
                    auth_info: None,
                    metadata: HashMap::new(),
                };
                
                // Call the local tool
                let result = self.call_tool(&name, arguments, context).await;
                
                // Create response message
                let response_payload = match result {
                    Ok(value) => MCPMessage::CallToolResult {
                        content: vec![MCPContent::Text { text: value.to_string() }],
                        is_error: false,
                    },
                    Err(e) => MCPMessage::Error {
                        code: -1,
                        message: e.to_string(),
                        data: None,
                    },
                };
                
                let response_message = P2PMCPMessage {
                    message_type: P2PMCPMessageType::Response,
                    message_id: message.message_id,
                    source_peer: "local".to_string(), // TODO: Get actual local peer ID
                    target_peer: Some(message.source_peer),
                    timestamp: SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map_err(|e| P2PError::Network(format!("Time error: {}", e)))?
                        .as_secs(),
                    payload: response_payload,
                    ttl: message.ttl.saturating_sub(1),
                };
                
                // Serialize response
                let response_data = serde_json::to_vec(&response_message)
                    .map_err(|e| P2PError::Serialization(e))?;
                
                Ok(Some(response_data))
            }
            MCPMessage::ListTools { cursor: _ } => {
                let (tools, _) = self.list_tools(None).await?;
                
                let response_payload = MCPMessage::ListToolsResult {
                    tools,
                    next_cursor: None,
                };
                
                let response_message = P2PMCPMessage {
                    message_type: P2PMCPMessageType::Response,
                    message_id: message.message_id,
                    source_peer: "local".to_string(), // TODO: Get actual local peer ID
                    target_peer: Some(message.source_peer),
                    timestamp: SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map_err(|e| P2PError::Network(format!("Time error: {}", e)))?
                        .as_secs(),
                    payload: response_payload,
                    ttl: message.ttl.saturating_sub(1),
                };
                
                let response_data = serde_json::to_vec(&response_message)
                    .map_err(|e| P2PError::Serialization(e))?;
                
                Ok(Some(response_data))
            }
            _ => {
                // Unsupported request type
                let error_response = P2PMCPMessage {
                    message_type: P2PMCPMessageType::Response,
                    message_id: message.message_id,
                    source_peer: "local".to_string(), // TODO: Get actual local peer ID
                    target_peer: Some(message.source_peer),
                    timestamp: SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map_err(|e| P2PError::Network(format!("Time error: {}", e)))?
                        .as_secs(),
                    payload: MCPMessage::Error {
                        code: -2,
                        message: "Unsupported request type".to_string(),
                        data: None,
                    },
                    ttl: message.ttl.saturating_sub(1),
                };
                
                let response_data = serde_json::to_vec(&error_response)
                    .map_err(|e| P2PError::Serialization(e))?;
                
                Ok(Some(response_data))
            }
        }
    }
    
    // Security-related methods
    
    /// Generate authentication token for peer
    pub async fn generate_auth_token(&self, peer_id: &PeerId, permissions: Vec<MCPPermission>, ttl: Duration) -> Result<String> {
        if let Some(security_manager) = &self.security_manager {
            let token = security_manager.generate_token(peer_id, permissions, ttl).await?;
            
            // Log authentication event
            let mut details = HashMap::new();
            details.insert("action".to_string(), "token_generated".to_string());
            details.insert("ttl_seconds".to_string(), ttl.as_secs().to_string());
            
            self.audit_logger.log_event(
                "authentication".to_string(),
                peer_id.clone(),
                details,
                AuditSeverity::Info,
            ).await;
            
            Ok(token)
        } else {
            Err(P2PError::MCP("Authentication not enabled".to_string()))
        }
    }
    
    /// Verify authentication token
    pub async fn verify_auth_token(&self, token: &str) -> Result<TokenPayload> {
        if let Some(security_manager) = &self.security_manager {
            match security_manager.verify_token(token).await {
                Ok(payload) => {
                    // Log successful verification
                    let mut details = HashMap::new();
                    details.insert("action".to_string(), "token_verified".to_string());
                    details.insert("subject".to_string(), payload.sub.clone());
                    
                    self.audit_logger.log_event(
                        "authentication".to_string(),
                        payload.iss.clone(),
                        details,
                        AuditSeverity::Info,
                    ).await;
                    
                    Ok(payload)
                }
                Err(e) => {
                    // Log failed verification
                    let mut details = HashMap::new();
                    details.insert("action".to_string(), "token_verification_failed".to_string());
                    details.insert("error".to_string(), e.to_string());
                    
                    self.audit_logger.log_event(
                        "authentication".to_string(),
                        "unknown".to_string(),
                        details,
                        AuditSeverity::Warning,
                    ).await;
                    
                    Err(e)
                }
            }
        } else {
            Err(P2PError::MCP("Authentication not enabled".to_string()))
        }
    }
    
    /// Check if peer has permission for operation
    pub async fn check_permission(&self, peer_id: &PeerId, permission: &MCPPermission) -> Result<bool> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.check_permission(peer_id, permission).await
        } else {
            // If security is disabled, allow all operations
            Ok(true)
        }
    }
    
    /// Check rate limit for peer
    pub async fn check_rate_limit(&self, peer_id: &PeerId) -> Result<bool> {
        if let Some(security_manager) = &self.security_manager {
            let allowed = security_manager.check_rate_limit(peer_id).await?;
            
            if !allowed {
                // Log rate limit violation
                let mut details = HashMap::new();
                details.insert("action".to_string(), "rate_limit_exceeded".to_string());
                
                self.audit_logger.log_event(
                    "rate_limiting".to_string(),
                    peer_id.clone(),
                    details,
                    AuditSeverity::Warning,
                ).await;
            }
            
            Ok(allowed)
        } else {
            // If rate limiting is disabled, allow all requests
            Ok(true)
        }
    }
    
    /// Grant permission to peer
    pub async fn grant_permission(&self, peer_id: &PeerId, permission: MCPPermission) -> Result<()> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.grant_permission(peer_id, permission.clone()).await?;
            
            // Log permission grant
            let mut details = HashMap::new();
            details.insert("action".to_string(), "permission_granted".to_string());
            details.insert("permission".to_string(), permission.as_str().to_string());
            
            self.audit_logger.log_event(
                "authorization".to_string(),
                peer_id.clone(),
                details,
                AuditSeverity::Info,
            ).await;
            
            Ok(())
        } else {
            Err(P2PError::MCP("Security not enabled".to_string()))
        }
    }
    
    /// Revoke permission from peer
    pub async fn revoke_permission(&self, peer_id: &PeerId, permission: &MCPPermission) -> Result<()> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.revoke_permission(peer_id, permission).await?;
            
            // Log permission revocation
            let mut details = HashMap::new();
            details.insert("action".to_string(), "permission_revoked".to_string());
            details.insert("permission".to_string(), permission.as_str().to_string());
            
            self.audit_logger.log_event(
                "authorization".to_string(),
                peer_id.clone(),
                details,
                AuditSeverity::Info,
            ).await;
            
            Ok(())
        } else {
            Err(P2PError::MCP("Security not enabled".to_string()))
        }
    }
    
    /// Add trusted peer
    pub async fn add_trusted_peer(&self, peer_id: PeerId) -> Result<()> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.add_trusted_peer(peer_id.clone()).await?;
            
            // Log trusted peer addition
            let mut details = HashMap::new();
            details.insert("action".to_string(), "trusted_peer_added".to_string());
            
            self.audit_logger.log_event(
                "trust_management".to_string(),
                peer_id,
                details,
                AuditSeverity::Info,
            ).await;
            
            Ok(())
        } else {
            Err(P2PError::MCP("Security not enabled".to_string()))
        }
    }
    
    /// Check if peer is trusted
    pub async fn is_trusted_peer(&self, peer_id: &PeerId) -> bool {
        if let Some(security_manager) = &self.security_manager {
            security_manager.is_trusted_peer(peer_id).await
        } else {
            false
        }
    }
    
    /// Set security policy for tool
    pub async fn set_tool_security_policy(&self, tool_name: String, level: SecurityLevel) -> Result<()> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.set_tool_policy(tool_name.clone(), level.clone()).await?;
            
            // Log policy change
            let mut details = HashMap::new();
            details.insert("action".to_string(), "tool_policy_set".to_string());
            details.insert("tool_name".to_string(), tool_name);
            details.insert("security_level".to_string(), format!("{:?}", level));
            
            self.audit_logger.log_event(
                "security_policy".to_string(),
                "system".to_string(),
                details,
                AuditSeverity::Info,
            ).await;
            
            Ok(())
        } else {
            Err(P2PError::MCP("Security not enabled".to_string()))
        }
    }
    
    /// Get security policy for tool
    pub async fn get_tool_security_policy(&self, tool_name: &str) -> SecurityLevel {
        if let Some(security_manager) = &self.security_manager {
            security_manager.get_tool_policy(tool_name).await
        } else {
            SecurityLevel::Public
        }
    }
    
    /// Get peer security statistics
    pub async fn get_peer_security_stats(&self, peer_id: &PeerId) -> Option<PeerACL> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.get_peer_stats(peer_id).await
        } else {
            None
        }
    }
    
    /// Get recent security audit entries
    pub async fn get_security_audit(&self, limit: Option<usize>) -> Vec<SecurityAuditEntry> {
        self.audit_logger.get_recent_entries(limit).await
    }
    
    /// Perform security housekeeping
    pub async fn security_cleanup(&self) -> Result<()> {
        if let Some(security_manager) = &self.security_manager {
            security_manager.cleanup().await?;
        }
        Ok(())
    }
    
    /// Handle remote response
    async fn handle_remote_response(&self, message: P2PMCPMessage) -> Result<()> {
        // Find the waiting request handler
        let response_tx = {
            let mut handlers = self.request_handlers.write().await;
            handlers.remove(&message.message_id)
        };
        
        if let Some(tx) = response_tx {
            let response = MCPResponse {
                request_id: message.message_id,
                message: message.payload,
                timestamp: SystemTime::now(),
                processing_time: Duration::from_millis(0), // TODO: Calculate actual processing time
            };
            
            // Send response to waiting caller
            let _ = tx.send(response);
        } else {
            debug!("Received response for unknown request: {}", message.message_id);
        }
        
        Ok(())
    }
    
    /// Handle service advertisement
    async fn handle_service_advertisement(&self, _message: P2PMCPMessage) -> Result<()> {
        // TODO: Parse service advertisement and update remote services cache
        Ok(())
    }
    
    /// Handle service discovery request
    async fn handle_service_discovery(&self, message: P2PMCPMessage) -> Result<Option<Vec<u8>>> {
        // Create service advertisement with our local services
        let local_services: Vec<MCPService> = {
            let services = self.local_services.read().await;
            services.values().cloned().collect()
        };
        
        if !local_services.is_empty() {
            let advertisement = P2PMCPMessage {
                message_type: P2PMCPMessageType::ServiceAdvertisement,
                message_id: uuid::Uuid::new_v4().to_string(),
                source_peer: "local".to_string(), // TODO: Get actual local peer ID
                target_peer: Some(message.source_peer),
                timestamp: SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .map_err(|e| P2PError::Network(format!("Time error: {}", e)))?
                    .as_secs(),
                payload: MCPMessage::ListToolsResult {
                    tools: local_services.into_iter()
                        .flat_map(|s| s.tools.into_iter().map(|t| MCPTool {
                            name: t,
                            description: "Remote tool".to_string(),
                            input_schema: json!({"type": "object"}),
                        }))
                        .collect(),
                    next_cursor: None,
                },
                ttl: message.ttl.saturating_sub(1),
            };
            
            let response_data = serde_json::to_vec(&advertisement)
                .map_err(|e| P2PError::Serialization(e))?;
            
            Ok(Some(response_data))
        } else {
            Ok(None)
        }
    }
    
    /// Shutdown the server
    pub async fn shutdown(&self) -> Result<()> {
        info!("Shutting down MCP server");
        
        // Close all sessions
        {
            let mut sessions = self.sessions.write().await;
            for session in sessions.values_mut() {
                session.state = MCPSessionState::Terminated;
            }
            sessions.clear();
        }
        
        // TODO: Cleanup tasks and channels
        
        info!("MCP server shutdown complete");
        Ok(())
    }
}

impl Tool {
    /// Create a new tool
    pub fn new(name: &str, description: &str, input_schema: Value) -> ToolBuilder {
        ToolBuilder {
            name: name.to_string(),
            description: description.to_string(),
            input_schema,
            handler: None,
            tags: Vec::new(),
        }
    }
}

/// Builder for creating tools
pub struct ToolBuilder {
    name: String,
    description: String,
    input_schema: Value,
    handler: Option<Box<dyn ToolHandler + Send + Sync>>,
    tags: Vec<String>,
}

impl ToolBuilder {
    /// Set tool handler
    pub fn handler<H: ToolHandler + Send + Sync + 'static>(mut self, handler: H) -> Self {
        self.handler = Some(Box::new(handler));
        self
    }
    
    /// Add tags
    pub fn tags(mut self, tags: Vec<String>) -> Self {
        self.tags = tags;
        self
    }
    
    /// Build the tool
    pub fn build(self) -> Result<Tool> {
        let handler = self.handler
            .ok_or_else(|| P2PError::MCP("Tool handler is required".to_string()))?;
        
        let definition = MCPTool {
            name: self.name,
            description: self.description,
            input_schema: self.input_schema,
        };
        
        let metadata = ToolMetadata {
            created_at: SystemTime::now(),
            last_called: None,
            call_count: 0,
            avg_execution_time: Duration::from_millis(0),
            health_status: ToolHealthStatus::Healthy,
            tags: self.tags,
        };
        
        Ok(Tool {
            definition,
            handler,
            metadata,
        })
    }
}

/// Simple function-based tool handler
pub struct FunctionToolHandler<F> {
    function: F,
}

impl<F, Fut> ToolHandler for FunctionToolHandler<F>
where
    F: Fn(Value) -> Fut + Send + Sync,
    Fut: std::future::Future<Output = Result<Value>> + Send + 'static,
{
    fn execute(&self, arguments: Value) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Value>> + Send + '_>> {
        Box::pin((self.function)(arguments))
    }
}

impl<F> FunctionToolHandler<F> {
    /// Create a new function-based tool handler
    pub fn new(function: F) -> Self {
        Self { function }
    }
}

/// MCP service descriptor for discovery and routing
impl MCPService {
    /// Create a new MCP service descriptor
    pub fn new(service_id: String, node_id: PeerId) -> Self {
        Self {
            service_id,
            node_id,
            tools: Vec::new(),
            capabilities: MCPCapabilities {
                experimental: None,
                sampling: None,
                tools: Some(MCPToolsCapability {
                    list_changed: Some(true),
                }),
                prompts: None,
                resources: None,
                logging: None,
            },
            metadata: MCPServiceMetadata {
                name: "MCP Service".to_string(),
                version: "1.0.0".to_string(),
                description: None,
                tags: Vec::new(),
                health_status: ServiceHealthStatus::Healthy,
                load_metrics: ServiceLoadMetrics {
                    active_requests: 0,
                    requests_per_second: 0.0,
                    avg_response_time_ms: 0.0,
                    error_rate: 0.0,
                    cpu_usage: 0.0,
                    memory_usage: 0,
                },
            },
            registered_at: SystemTime::now(),
            endpoint: MCPEndpoint {
                protocol: "p2p".to_string(),
                address: "".to_string(),
                port: None,
                tls: false,
                auth_required: false,
            },
        }
    }
}

impl Default for MCPCapabilities {
    fn default() -> Self {
        Self {
            experimental: None,
            sampling: None,
            tools: Some(MCPToolsCapability {
                list_changed: Some(true),
            }),
            prompts: Some(MCPPromptsCapability {
                list_changed: Some(true),
            }),
            resources: Some(MCPResourcesCapability {
                subscribe: Some(true),
                list_changed: Some(true),
            }),
            logging: Some(MCPLoggingCapability {
                levels: Some(vec![
                    MCPLogLevel::Debug,
                    MCPLogLevel::Info,
                    MCPLogLevel::Warning,
                    MCPLogLevel::Error,
                ]),
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dht::{DHT, DHTConfig, Key};
    use std::pin::Pin;
    use std::future::Future;
    use tokio::time::timeout;

    /// Test implementation of ToolHandler for unit tests
    struct TestTool {
        name: String,
        should_error: bool,
        execution_time: Duration,
    }

    impl TestTool {
        fn new(name: &str) -> Self {
            Self {
                name: name.to_string(),
                should_error: false,
                execution_time: Duration::from_millis(10),
            }
        }

        fn with_error(mut self) -> Self {
            self.should_error = true;
            self
        }

        fn with_execution_time(mut self, duration: Duration) -> Self {
            self.execution_time = duration;
            self
        }
    }

    impl ToolHandler for TestTool {
        fn execute(&self, arguments: Value) -> Pin<Box<dyn Future<Output = Result<Value>> + Send + '_>> {
            let should_error = self.should_error;
            let execution_time = self.execution_time;
            let name = self.name.clone();

            Box::pin(async move {
                tokio::time::sleep(execution_time).await;

                if should_error {
                    return Err(P2PError::MCP(format!("Test error from tool {}", name)).into());
                }

                // Echo back the arguments with a response marker
                Ok(json!({
                    "tool": name,
                    "arguments": arguments,
                    "result": "success"
                }))
            })
        }

        fn validate(&self, arguments: &Value) -> Result<()> {
            if !arguments.is_object() {
                return Err(P2PError::MCP("Arguments must be an object".to_string()).into());
            }
            Ok(())
        }

        fn get_requirements(&self) -> ToolRequirements {
            ToolRequirements {
                max_memory: Some(1024 * 1024), // 1MB
                max_execution_time: Some(Duration::from_secs(5)),
                required_capabilities: vec!["test".to_string()],
                requires_network: false,
                requires_filesystem: false,
            }
        }
    }

    /// Helper function to create a test MCP server
    async fn create_test_mcp_server() -> MCPServer {
        let config = MCPServerConfig {
            server_name: "test_server".to_string(),
            server_version: "1.0.0".to_string(),
            enable_auth: false,
            enable_rate_limiting: false,
            max_concurrent_requests: 10,
            request_timeout: Duration::from_secs(30),
            enable_dht_discovery: true,
            rate_limit_rpm: 60,
            enable_logging: true,
            max_tool_execution_time: Duration::from_secs(30),
            tool_memory_limit: 100 * 1024 * 1024,
        };

        MCPServer::new(config)
    }

    /// Helper function to create a test tool
    fn create_test_tool(name: &str) -> Tool {
        Tool {
            definition: MCPTool {
                name: name.to_string(),
                description: format!("Test tool: {}", name),
                input_schema: json!({
                    "type": "object",
                    "properties": {
                        "input": { "type": "string" }
                    }
                }),
            },
            handler: Box::new(TestTool::new(name)),
            metadata: ToolMetadata {
                created_at: SystemTime::now(),
                last_called: None,
                call_count: 0,
                avg_execution_time: Duration::from_millis(0),
                health_status: ToolHealthStatus::Healthy,
                tags: vec!["test".to_string()],
            },
        }
    }

    /// Helper function to create a test DHT
    async fn create_test_dht() -> DHT {
        let local_id = Key::new(b"test_node_id");
        let config = DHTConfig::default();
        DHT::new(local_id, config)
    }

    /// Helper function to create an MCP call context
    fn create_test_context(caller_id: PeerId) -> MCPCallContext {
        MCPCallContext {
            caller_id,
            timestamp: SystemTime::now(),
            timeout: Duration::from_secs(30),
            auth_info: None,
            metadata: HashMap::new(),
        }
    }

    #[tokio::test]
    async fn test_mcp_server_creation() {
        let server = create_test_mcp_server().await;
        assert_eq!(server.config.server_name, "test_server");
        assert_eq!(server.config.server_version, "1.0.0");
        assert!(!server.config.enable_auth);
        assert!(!server.config.enable_rate_limiting);
    }

    #[tokio::test]
    async fn test_tool_registration() -> Result<()> {
        let server = create_test_mcp_server().await;
        let tool = create_test_tool("test_calculator");

        // Register the tool
        server.register_tool(tool).await?;

        // Verify tool is registered
        let tools = server.tools.read().await;
        assert!(tools.contains_key("test_calculator"));
        assert_eq!(tools.get("test_calculator").unwrap().definition.name, "test_calculator");

        // Verify stats updated
        let stats = server.stats.read().await;
        assert_eq!(stats.total_tools, 1);

        Ok(())
    }

    #[tokio::test]
    async fn test_tool_registration_duplicate() -> Result<()> {
        let server = create_test_mcp_server().await;
        let tool1 = create_test_tool("duplicate_tool");
        let tool2 = create_test_tool("duplicate_tool");

        // Register first tool
        server.register_tool(tool1).await?;

        // Try to register duplicate - should fail
        let result = server.register_tool(tool2).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Tool already exists"));

        Ok(())
    }

    #[tokio::test]
    async fn test_tool_validation() {
        let server = create_test_mcp_server().await;

        // Test invalid tool name (empty)
        let mut invalid_tool = create_test_tool("");
        let result = server.validate_tool(&invalid_tool).await;
        assert!(result.is_err());

        // Test invalid tool name (too long)
        invalid_tool.definition.name = "a".repeat(200);
        let result = server.validate_tool(&invalid_tool).await;
        assert!(result.is_err());

        // Test invalid schema (not an object)
        let mut invalid_schema_tool = create_test_tool("valid_name");
        invalid_schema_tool.definition.input_schema = json!("not an object");
        let result = server.validate_tool(&invalid_schema_tool).await;
        assert!(result.is_err());

        // Test valid tool
        let valid_tool = create_test_tool("valid_tool");
        let result = server.validate_tool(&valid_tool).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_tool_call_success() -> Result<()> {
        let server = create_test_mcp_server().await;
        let tool = create_test_tool("success_tool");
        server.register_tool(tool).await?;

        let caller_id = "test_peer_123".to_string();
        let context = create_test_context(caller_id);
        let arguments = json!({"input": "test data"});

        let result = server.call_tool("success_tool", arguments.clone(), context).await?;

        // Verify response structure
        assert_eq!(result["tool"], "success_tool");
        assert_eq!(result["arguments"], arguments);
        assert_eq!(result["result"], "success");

        // Verify tool metadata updated
        let tools = server.tools.read().await;
        let tool_metadata = &tools.get("success_tool").unwrap().metadata;
        assert_eq!(tool_metadata.call_count, 1);
        assert!(tool_metadata.last_called.is_some());

        Ok(())
    }

    #[tokio::test]
    async fn test_tool_call_nonexistent() -> Result<()> {
        let server = create_test_mcp_server().await;
        let caller_id = "test_peer_456".to_string();
        let context = create_test_context(caller_id);
        let arguments = json!({"input": "test"});

        let result = server.call_tool("nonexistent_tool", arguments, context).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Tool not found"));

        Ok(())
    }

    #[tokio::test]
    async fn test_tool_call_handler_error() -> Result<()> {
        let server = create_test_mcp_server().await;
        let tool = Tool {
            definition: MCPTool {
                name: "error_tool".to_string(),
                description: "Tool that always errors".to_string(),
                input_schema: json!({"type": "object"}),
            },
            handler: Box::new(TestTool::new("error_tool").with_error()),
            metadata: ToolMetadata {
                created_at: SystemTime::now(),
                last_called: None,
                call_count: 0,
                avg_execution_time: Duration::from_millis(0),
                health_status: ToolHealthStatus::Healthy,
                tags: vec![],
            },
        };

        server.register_tool(tool).await?;

        let caller_id = "test_peer_error".to_string();
        let context = create_test_context(caller_id);
        let arguments = json!({"input": "test"});

        let result = server.call_tool("error_tool", arguments, context).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Test error from tool error_tool"));

        Ok(())
    }

    #[tokio::test]
    async fn test_tool_call_timeout() -> Result<()> {
        let server = create_test_mcp_server().await;
        let slow_tool = Tool {
            definition: MCPTool {
                name: "slow_tool".to_string(),
                description: "Tool that takes too long".to_string(),
                input_schema: json!({"type": "object"}),
            },
            handler: Box::new(TestTool::new("slow_tool").with_execution_time(Duration::from_secs(2))),
            metadata: ToolMetadata {
                created_at: SystemTime::now(),
                last_called: None,
                call_count: 0,
                avg_execution_time: Duration::from_millis(0),
                health_status: ToolHealthStatus::Healthy,
                tags: vec![],
            },
        };

        server.register_tool(slow_tool).await?;

        let caller_id = "test_peer_error".to_string();
        let context = create_test_context(caller_id);
        let arguments = json!({"input": "test"});

        // Test with very short timeout
        let result = timeout(
            Duration::from_millis(100),
            server.call_tool("slow_tool", arguments, context)
        ).await;

        assert!(result.is_err()); // Should timeout

        Ok(())
    }

    #[tokio::test]
    async fn test_tool_requirements() {
        let tool = TestTool::new("req_tool");
        let requirements = tool.get_requirements();

        assert_eq!(requirements.max_memory, Some(1024 * 1024));
        assert_eq!(requirements.max_execution_time, Some(Duration::from_secs(5)));
        assert_eq!(requirements.required_capabilities, vec!["test"]);
        assert!(!requirements.requires_network);
        assert!(!requirements.requires_filesystem);
    }

    #[tokio::test]
    async fn test_tool_validation_handler() {
        let tool = TestTool::new("validation_tool");

        // Valid arguments (object)
        let valid_args = json!({"key": "value"});
        assert!(tool.validate(&valid_args).is_ok());

        // Invalid arguments (not an object)
        let invalid_args = json!("not an object");
        assert!(tool.validate(&invalid_args).is_err());

        let invalid_args = json!(123);
        assert!(tool.validate(&invalid_args).is_err());
    }

    #[tokio::test]
    async fn test_tool_health_status() {
        let mut metadata = ToolMetadata {
            created_at: SystemTime::now(),
            last_called: None,
            call_count: 0,
            avg_execution_time: Duration::from_millis(0),
            health_status: ToolHealthStatus::Healthy,
            tags: vec![],
        };

        // Test different health statuses
        assert_eq!(metadata.health_status, ToolHealthStatus::Healthy);

        metadata.health_status = ToolHealthStatus::Degraded;
        assert_eq!(metadata.health_status, ToolHealthStatus::Degraded);

        metadata.health_status = ToolHealthStatus::Unhealthy;
        assert_eq!(metadata.health_status, ToolHealthStatus::Unhealthy);

        metadata.health_status = ToolHealthStatus::Disabled;
        assert_eq!(metadata.health_status, ToolHealthStatus::Disabled);
    }

    #[tokio::test]
    async fn test_mcp_capabilities() {
        let server = create_test_mcp_server().await;
        let capabilities = server.get_server_capabilities().await;

        assert!(capabilities.tools.is_some());
        assert!(capabilities.prompts.is_some());
        assert!(capabilities.resources.is_some());
        assert!(capabilities.logging.is_some());

        let tools_cap = capabilities.tools.unwrap();
        assert_eq!(tools_cap.list_changed, Some(true));

        let logging_cap = capabilities.logging.unwrap();
        let levels = logging_cap.levels.unwrap();
        assert!(levels.contains(&MCPLogLevel::Debug));
        assert!(levels.contains(&MCPLogLevel::Info));
        assert!(levels.contains(&MCPLogLevel::Warning));
        assert!(levels.contains(&MCPLogLevel::Error));
    }

    #[tokio::test]
    async fn test_mcp_message_serialization() {
        // Test Initialize message
        let init_msg = MCPMessage::Initialize {
            protocol_version: MCP_VERSION.to_string(),
            capabilities: MCPCapabilities {
                experimental: None,
                sampling: None,
                tools: Some(MCPToolsCapability { list_changed: Some(true) }),
                prompts: None,
                resources: None,
                logging: None,
            },
            client_info: MCPClientInfo {
                name: "test_client".to_string(),
                version: "1.0.0".to_string(),
            },
        };

        let serialized = serde_json::to_string(&init_msg).unwrap();
        let deserialized: MCPMessage = serde_json::from_str(&serialized).unwrap();

        match deserialized {
            MCPMessage::Initialize { protocol_version, client_info, .. } => {
                assert_eq!(protocol_version, MCP_VERSION);
                assert_eq!(client_info.name, "test_client");
                assert_eq!(client_info.version, "1.0.0");
            }
            _ => panic!("Wrong message type after deserialization"),
        }
    }

    #[tokio::test]
    async fn test_mcp_content_types() {
        // Test text content
        let text_content = MCPContent::Text {
            text: "Hello, world!".to_string(),
        };

        let serialized = serde_json::to_string(&text_content).unwrap();
        let deserialized: MCPContent = serde_json::from_str(&serialized).unwrap();

        match deserialized {
            MCPContent::Text { text } => assert_eq!(text, "Hello, world!"),
            _ => panic!("Wrong content type"),
        }

        // Test image content
        let image_content = MCPContent::Image {
            data: "base64data".to_string(),
            mime_type: "image/png".to_string(),
        };

        let serialized = serde_json::to_string(&image_content).unwrap();
        let deserialized: MCPContent = serde_json::from_str(&serialized).unwrap();

        match deserialized {
            MCPContent::Image { data, mime_type } => {
                assert_eq!(data, "base64data");
                assert_eq!(mime_type, "image/png");
            }
            _ => panic!("Wrong content type"),
        }
    }

    #[tokio::test]
    async fn test_service_health_status() {
        let mut metrics = ServiceLoadMetrics {
            active_requests: 0,
            requests_per_second: 0.0,
            avg_response_time_ms: 0.0,
            error_rate: 0.0,
            cpu_usage: 0.0,
            memory_usage: 0,
        };

        // Test healthy service
        let metadata = MCPServiceMetadata {
            name: "test_service".to_string(),
            version: "1.0.0".to_string(),
            description: Some("Test service".to_string()),
            tags: vec!["test".to_string()],
            health_status: ServiceHealthStatus::Healthy,
            load_metrics: metrics.clone(),
        };

        assert_eq!(metadata.health_status, ServiceHealthStatus::Healthy);

        // Test different health statuses
        metrics.error_rate = 0.5; // 50% error rate
        let degraded_metadata = MCPServiceMetadata {
            health_status: ServiceHealthStatus::Degraded,
            load_metrics: metrics.clone(),
            ..metadata.clone()
        };

        assert_eq!(degraded_metadata.health_status, ServiceHealthStatus::Degraded);

        let unhealthy_metadata = MCPServiceMetadata {
            health_status: ServiceHealthStatus::Unhealthy,
            ..metadata.clone()
        };

        assert_eq!(unhealthy_metadata.health_status, ServiceHealthStatus::Unhealthy);
    }

    #[tokio::test]
    async fn test_p2p_mcp_message() {
        let source_peer = "source_peer_123".to_string();
        let target_peer = "target_peer_456".to_string();

        let p2p_message = P2PMCPMessage {
            message_type: P2PMCPMessageType::Request,
            message_id: uuid::Uuid::new_v4().to_string(),
            source_peer: source_peer.clone(),
            target_peer: Some(target_peer.clone()),
            timestamp: SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
            payload: MCPMessage::ListTools { cursor: None },
            ttl: 10,
        };

        // Test serialization
        let serialized = serde_json::to_string(&p2p_message).unwrap();
        let deserialized: P2PMCPMessage = serde_json::from_str(&serialized).unwrap();

        assert_eq!(deserialized.message_type, P2PMCPMessageType::Request);
        assert_eq!(deserialized.source_peer, source_peer);
        assert_eq!(deserialized.target_peer, Some(target_peer));
        assert_eq!(deserialized.ttl, 10);

        match deserialized.payload {
            MCPMessage::ListTools { cursor } => assert_eq!(cursor, None),
            _ => panic!("Wrong message payload type"),
        }
    }

    #[tokio::test]
    async fn test_tool_requirements_default() {
        let default_requirements = ToolRequirements::default();

        assert_eq!(default_requirements.max_memory, Some(100 * 1024 * 1024));
        assert_eq!(default_requirements.max_execution_time, Some(Duration::from_secs(30)));
        assert!(default_requirements.required_capabilities.is_empty());
        assert!(!default_requirements.requires_network);
        assert!(!default_requirements.requires_filesystem);
    }

    #[tokio::test]
    async fn test_mcp_server_stats() {
        let server = create_test_mcp_server().await;

        // Initial stats should be zero
        let stats = server.stats.read().await;
        assert_eq!(stats.total_tools, 0);
        assert_eq!(stats.total_requests, 0);
        assert_eq!(stats.total_responses, 0);
        assert_eq!(stats.total_errors, 0);

        drop(stats);

        // Register a tool and verify stats update
        let tool = create_test_tool("stats_test_tool");
        server.register_tool(tool).await.unwrap();

        let stats = server.stats.read().await;
        assert_eq!(stats.total_tools, 1);
    }

    #[tokio::test]
    async fn test_log_levels() {
        // Test all log levels serialize/deserialize correctly
        let levels = vec![
            MCPLogLevel::Debug,
            MCPLogLevel::Info,
            MCPLogLevel::Notice,
            MCPLogLevel::Warning,
            MCPLogLevel::Error,
            MCPLogLevel::Critical,
            MCPLogLevel::Alert,
            MCPLogLevel::Emergency,
        ];

        for level in levels {
            let serialized = serde_json::to_string(&level).unwrap();
            let deserialized: MCPLogLevel = serde_json::from_str(&serialized).unwrap();
            assert_eq!(level as u8, deserialized as u8);
        }
    }

    #[tokio::test]
    async fn test_mcp_endpoint() {
        let endpoint = MCPEndpoint {
            protocol: "p2p".to_string(),
            address: "127.0.0.1".to_string(),
            port: Some(9000),
            tls: true,
            auth_required: true,
        };

        let serialized = serde_json::to_string(&endpoint).unwrap();
        let deserialized: MCPEndpoint = serde_json::from_str(&serialized).unwrap();

        assert_eq!(deserialized.protocol, "p2p");
        assert_eq!(deserialized.address, "127.0.0.1");
        assert_eq!(deserialized.port, Some(9000));
        assert!(deserialized.tls);
        assert!(deserialized.auth_required);
    }

    #[tokio::test]
    async fn test_mcp_service_metadata() {
        let load_metrics = ServiceLoadMetrics {
            active_requests: 5,
            requests_per_second: 10.5,
            avg_response_time_ms: 250.0,
            error_rate: 0.01,
            cpu_usage: 45.5,
            memory_usage: 1024 * 1024 * 100, // 100MB
        };

        let metadata = MCPServiceMetadata {
            name: "test_service".to_string(),
            version: "2.1.0".to_string(),
            description: Some("A test service for unit testing".to_string()),
            tags: vec!["test".to_string(), "unit".to_string(), "mcp".to_string()],
            health_status: ServiceHealthStatus::Healthy,
            load_metrics,
        };

        // Test serialization
        let serialized = serde_json::to_string(&metadata).unwrap();
        let deserialized: MCPServiceMetadata = serde_json::from_str(&serialized).unwrap();

        assert_eq!(deserialized.name, "test_service");
        assert_eq!(deserialized.version, "2.1.0");
        assert_eq!(deserialized.description, Some("A test service for unit testing".to_string()));
        assert_eq!(deserialized.tags, vec!["test", "unit", "mcp"]);
        assert_eq!(deserialized.health_status, ServiceHealthStatus::Healthy);
        assert_eq!(deserialized.load_metrics.active_requests, 5);
        assert_eq!(deserialized.load_metrics.requests_per_second, 10.5);
    }

    #[tokio::test]
    async fn test_function_tool_handler() {
        // Test function tool handler creation and execution
        let handler = FunctionToolHandler::new(|args: Value| async move {
            let name = args.get("name").and_then(|v| v.as_str()).unwrap_or("world");
            Ok(json!({"greeting": format!("Hello, {}!", name)}))
        });

        let args = json!({"name": "Alice"});
        let result = handler.execute(args).await.unwrap();
        assert_eq!(result["greeting"], "Hello, Alice!");

        // Test with missing argument
        let empty_args = json!({});
        let result = handler.execute(empty_args).await.unwrap();
        assert_eq!(result["greeting"], "Hello, world!");
    }

    #[tokio::test]
    async fn test_mcp_service_creation() {
        let service_id = "test_service_123".to_string();
        let node_id = "test_node_789".to_string();

        let service = MCPService::new(service_id.clone(), node_id.clone());

        assert_eq!(service.service_id, service_id);
        assert_eq!(service.node_id, node_id);
        assert!(service.tools.is_empty());
        assert_eq!(service.metadata.name, "MCP Service");
        assert_eq!(service.metadata.version, "1.0.0");
        assert_eq!(service.metadata.health_status, ServiceHealthStatus::Healthy);
        assert_eq!(service.endpoint.protocol, "p2p");
        assert!(!service.endpoint.tls);
        assert!(!service.endpoint.auth_required);
    }

    #[tokio::test]
    async fn test_mcp_capabilities_default() {
        let capabilities = MCPCapabilities::default();

        assert!(capabilities.tools.is_some());
        assert!(capabilities.prompts.is_some());
        assert!(capabilities.resources.is_some());
        assert!(capabilities.logging.is_some());

        let tools_cap = capabilities.tools.unwrap();
        assert_eq!(tools_cap.list_changed, Some(true));

        let resources_cap = capabilities.resources.unwrap();
        assert_eq!(resources_cap.subscribe, Some(true));
        assert_eq!(resources_cap.list_changed, Some(true));

        let logging_cap = capabilities.logging.unwrap();
        let levels = logging_cap.levels.unwrap();
        assert!(levels.contains(&MCPLogLevel::Debug));
        assert!(levels.contains(&MCPLogLevel::Info));
        assert!(levels.contains(&MCPLogLevel::Warning));
        assert!(levels.contains(&MCPLogLevel::Error));
    }

    #[tokio::test]
    async fn test_mcp_request_creation() {
        let source_peer = "source_peer_123".to_string();
        let target_peer = "target_peer_456".to_string();

        let request = MCPRequest {
            request_id: uuid::Uuid::new_v4().to_string(),
            source_peer: source_peer.clone(),
            target_peer: target_peer.clone(),
            message: MCPMessage::ListTools { cursor: None },
            timestamp: SystemTime::now(),
            timeout: Duration::from_secs(30),
            auth_token: Some("test_token".to_string()),
        };

        assert_eq!(request.source_peer, source_peer);
        assert_eq!(request.target_peer, target_peer);
        assert_eq!(request.timeout, Duration::from_secs(30));
        assert_eq!(request.auth_token, Some("test_token".to_string()));

        match request.message {
            MCPMessage::ListTools { cursor } => assert_eq!(cursor, None),
            _ => panic!("Wrong message type"),
        }
    }

    #[tokio::test]
    async fn test_p2p_message_types() {
        // Test all P2P message types
        assert_eq!(P2PMCPMessageType::Request, P2PMCPMessageType::Request);
        assert_eq!(P2PMCPMessageType::Response, P2PMCPMessageType::Response);
        assert_eq!(P2PMCPMessageType::ServiceAdvertisement, P2PMCPMessageType::ServiceAdvertisement);
        assert_eq!(P2PMCPMessageType::ServiceDiscovery, P2PMCPMessageType::ServiceDiscovery);

        // Test serialization of each type
        for msg_type in [
            P2PMCPMessageType::Request,
            P2PMCPMessageType::Response,
            P2PMCPMessageType::ServiceAdvertisement,
            P2PMCPMessageType::ServiceDiscovery,
        ] {
            let serialized = serde_json::to_string(&msg_type).unwrap();
            let deserialized: P2PMCPMessageType = serde_json::from_str(&serialized).unwrap();
            assert_eq!(msg_type, deserialized);
        }
    }
}