saorsa_core/

mcp.rs

// Copyright 2024 Saorsa Labs Limited
//
// This software is dual-licensed under:
// - GNU Affero General Public License v3.0 or later (AGPL-3.0-or-later)
// - Commercial License
//
// For AGPL-3.0 license, see LICENSE-AGPL-3.0
// For commercial licensing, contact: saorsalabs@gmail.com
//
// Unless required by applicable law or agreed to in writing, software
// distributed under these licenses is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

//! 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

#![allow(missing_docs)]

pub mod security;

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

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";

/// Network message sender trait for MCP server
#[async_trait::async_trait]
pub trait NetworkSender: Send + Sync {
    /// Send a message to a specific peer via the P2P network
    async fn send_message(&self, peer_id: &PeerId, protocol: &str, data: Vec<u8>) -> Result<()>;

    /// Get our local peer ID
    fn local_peer_id(&self) -> &PeerId;
}

/// Message sender function type for simplified network integration
pub type MessageSender = Arc<dyn Fn(&PeerId, &str, Vec<u8>) -> Result<()> + Send + Sync>;

/// 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,
}

/// Health monitoring configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthMonitorConfig {
    /// Interval between health checks
    pub health_check_interval: Duration,
    /// Timeout for health check requests
    pub health_check_timeout: Duration,
    /// Number of consecutive failures before marking unhealthy
    pub failure_threshold: u32,
    /// Number of consecutive successes to mark healthy again
    pub success_threshold: u32,
    /// Interval for sending heartbeats
    pub heartbeat_interval: Duration,
    /// Maximum age of last heartbeat before considering service stale
    pub heartbeat_timeout: Duration,
    /// Whether to enable health monitoring
    pub enabled: bool,
}

impl Default for HealthMonitorConfig {
    fn default() -> Self {
        Self {
            health_check_interval: Duration::from_secs(30),
            health_check_timeout: Duration::from_secs(5),
            failure_threshold: 3,
            success_threshold: 2,
            heartbeat_interval: Duration::from_secs(60),
            heartbeat_timeout: Duration::from_secs(300), // 5 minutes
            enabled: true,
        }
    }
}

/// Service health information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceHealth {
    /// Service identifier
    pub service_id: String,
    /// Current health status
    pub status: ServiceHealthStatus,
    /// Last successful health check
    pub last_health_check: Option<SystemTime>,
    /// Last heartbeat received
    pub last_heartbeat: Option<SystemTime>,
    /// Consecutive failure count
    pub failure_count: u32,
    /// Consecutive success count
    pub success_count: u32,
    /// Average response time for health checks
    pub avg_response_time: Duration,
    /// Error message if unhealthy
    pub error_message: Option<String>,
    /// Health check history (last 10 checks)
    pub health_history: Vec<HealthCheckResult>,
}

/// Result of a health check
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthCheckResult {
    /// Timestamp of the check
    pub timestamp: SystemTime,
    /// Whether the check was successful
    pub success: bool,
    /// Response time
    pub response_time: Duration,
    /// Error message if failed
    pub error_message: Option<String>,
}

/// Heartbeat message
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Heartbeat {
    /// Service ID sending the heartbeat
    pub service_id: String,
    /// Peer ID of the service
    pub peer_id: PeerId,
    /// Timestamp when heartbeat was sent
    pub timestamp: SystemTime,
    /// Current service load (0.0 to 1.0)
    pub load: f32,
    /// Available tools
    pub available_tools: Vec<String>,
    /// Service capabilities
    pub capabilities: MCPCapabilities,
}

/// Health monitoring events
#[derive(Debug, Clone)]
pub enum HealthEvent {
    /// Service became healthy
    ServiceHealthy { service_id: String, peer_id: PeerId },
    /// Service became unhealthy
    ServiceUnhealthy {
        service_id: String,
        peer_id: PeerId,
        error: String,
    },
    /// Service status changed to degraded
    ServiceDegraded {
        service_id: String,
        peer_id: PeerId,
        reason: String,
    },
    /// Heartbeat received
    HeartbeatReceived {
        service_id: String,
        peer_id: PeerId,
        load: f32,
    },
    /// Heartbeat timeout
    HeartbeatTimeout { service_id: String, peer_id: PeerId },
}

/// 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, Copy, Serialize, Deserialize, PartialEq)]
pub enum ServiceHealthStatus {
    /// Service is healthy and responsive
    Healthy,
    /// Service is experiencing degraded performance
    Degraded,
    /// Service is not responding to health checks
    Unhealthy,
    /// Service is disabled/maintenance mode
    Disabled,
    /// Service status is unknown
    Unknown,
}

/// 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,
    /// Heartbeat notification
    Heartbeat,
    /// Health check request
    HealthCheck,
}

/// 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,
    /// Health monitoring configuration
    pub health_monitor: HealthMonitorConfig,
}

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
            health_monitor: HealthMonitorConfig::default(),
        }
    }
}

/// 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>,
    /// Network sender for P2P communication
    network_sender: Arc<RwLock<Option<Arc<dyn NetworkSender>>>>,
    /// Service health tracking
    service_health: Arc<RwLock<HashMap<String, ServiceHealth>>>,
    /// Health event sender
    health_event_tx: mpsc::UnboundedSender<HealthEvent>,
    /// Health event receiver (for monitoring)
    #[allow(dead_code)]
    health_event_rx: Arc<RwLock<mpsc::UnboundedReceiver<HealthEvent>>>,
    /// Node ID for this MCP server
    node_id: Option<PeerId>,
}

/// 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();
        let (health_event_tx, health_event_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
        };

        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
            network_sender: Arc::new(RwLock::new(None)),
            service_health: Arc::new(RwLock::new(HashMap::new())),
            health_event_tx,
            health_event_rx: Arc::new(RwLock::new(health_event_rx)),
            node_id: None,
        }
    }

    /// Create MCP server with DHT integration
    pub fn with_dht(mut self, dht: Arc<RwLock<DHT>>) -> Self {
        self.dht = Some(dht);
        self
    }

    /// Set the network sender for P2P communication
    pub async fn with_network_sender(mut self, sender: Arc<dyn NetworkSender>) -> Self {
        let peer_id = sender.local_peer_id().clone();
        self.node_id = Some(peer_id);
        *self.network_sender.write().await = Some(sender);
        self
    }

    /// Set the network sender for P2P communication (async method for post-creation setup)
    pub async fn set_network_sender(&self, sender: Arc<dyn NetworkSender>) {
        let peer_id = sender.local_peer_id().clone();
        *self.network_sender.write().await = Some(sender);
        info!("MCP server network sender configured for peer {}", peer_id);
    }

    /// Get the node ID as a string, returns "uninitialized" if node_id is None
    fn get_node_id_string(&self) -> String {
        self.node_id
            .as_ref()
            .map(|id| id.to_string())
            .unwrap_or_else(|| "uninitialized".to_string())
    }

    /// 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?;
        }

        // Announce updated service with new tool
        if let Err(e) = self.announce_local_services().await {
            warn!("Failed to announce service after tool registration: {}", e);
        }

        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(crate::error::McpError::InvalidRequest(
                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(crate::error::McpError::InvalidRequest(
                "Invalid tool name".to_string().into(),
            )));
        }

        // Validate schema
        if !tool.definition.input_schema.is_object() {
            return Err(P2PError::Mcp(crate::error::McpError::InvalidRequest(
                "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": self.get_node_id_string(),
            "registered_at": SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_err(|e| P2PError::Identity(crate::error::IdentityError::SystemTime(format!("Time error: {e}").into())))?.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(crate::error::McpError::InvalidRequest(
                "Rate limit exceeded".to_string().into(),
            )));
        }

        // 2. Check tool execution permission
        if !self
            .check_permission(&context.caller_id, &MCPPermission::ExecuteTools)
            .await?
        {
            return Err(P2PError::Mcp(crate::error::McpError::InvalidRequest(
                "Permission denied: execute tools".to_string().into(),
            )));
        }

        // 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(crate::error::McpError::InvalidRequest(
                        "Permission denied: admin access required"
                            .to_string()
                            .into(),
                    )));
                }
            }
            SecurityLevel::Strong => {
                if !is_trusted {
                    return Err(P2PError::Mcp(crate::error::McpError::InvalidRequest(
                        "Permission denied: trusted peer required"
                            .to_string()
                            .into(),
                    )));
                }
            }
            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(crate::error::McpError::InvalidRequest(
                            "Authentication required".to_string().into(),
                        )));
                    }
                }
            }
            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(crate::error::McpError::ToolNotFound(
                tool_name.to_string().into(),
            )));
        }

        // Validate arguments and get requirements
        let requirements = {
            let tools = self.tools.read().await;
            let tool = tools.get(tool_name).ok_or_else(|| {
                P2PError::Mcp(crate::error::McpError::ToolNotFound(
                    tool_name.to_string().into(),
                ))
            })?;

            // Validate arguments
            if let Err(e) = tool.handler.validate(&arguments) {
                return Err(P2PError::Mcp(crate::error::McpError::ExecutionFailed(
                    format!("{}: Validation failed: {e}", tool_name).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).ok_or_else(|| {
                P2PError::Mcp(crate::error::McpError::ToolNotFound(
                    tool_name_owned.clone().into(),
                ))
            })?;
            tool.handler.execute(arguments).await
        })
        .await
        .map_err(|_| {
            P2PError::Mcp(crate::error::McpError::ExecutionFailed(
                format!("{}: Tool execution timeout", tool_name).into(),
            ))
        })?
        .map_err(|e| {
            P2PError::Mcp(crate::error::McpError::ExecutionFailed(
                format!("{}: {e}", tool_name).into(),
            ))
        })?;

        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
            && max_memory > self.config.tool_memory_limit
        {
            return Err(P2PError::Mcp(crate::error::McpError::InvalidRequest(
                "Tool memory requirement exceeds limit".to_string().into(),
            )));
        }

        // Check execution time limit
        if let Some(max_execution_time) = requirements.max_execution_time
            && max_execution_time > self.config.max_tool_execution_time
        {
            return Err(P2PError::Mcp(crate::error::McpError::InvalidRequest(
                "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.

            // Placeholder single sleep to simulate periodic work without an infinite or never loop
            tokio::time::sleep(Duration::from_millis(100)).await;

            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<()> {
        if !self.config.health_monitor.enabled {
            debug!("Health monitoring is disabled");
            return Ok(());
        }

        info!(
            "Starting health monitoring with interval: {:?}",
            self.config.health_monitor.health_check_interval
        );

        // Clone necessary fields for the background task
        let service_health = Arc::clone(&self.service_health);
        let remote_services = Arc::clone(&self.remote_services);
        let network_sender = Arc::clone(&self.network_sender);
        let health_event_tx = self.health_event_tx.clone();
        let config = self.config.health_monitor.clone();

        // Start health check task
        let health_check_task = {
            let service_health = Arc::clone(&service_health);
            let remote_services = Arc::clone(&remote_services);
            let network_sender = Arc::clone(&network_sender);
            let health_event_tx = health_event_tx.clone();
            let config = config.clone();

            tokio::spawn(async move {
                let mut interval = tokio::time::interval(config.health_check_interval);

                loop {
                    interval.tick().await;

                    // Get list of remote services to check
                    let services_to_check: Vec<MCPService> = {
                        let remote_guard = remote_services.read().await;
                        remote_guard.values().cloned().collect()
                    };

                    // Perform health checks on all remote services
                    for service in services_to_check {
                        if let Some(sender) = network_sender.read().await.as_ref() {
                            Self::perform_health_check(
                                &service,
                                sender.as_ref(),
                                &service_health,
                                &health_event_tx,
                                &config,
                            )
                            .await;
                        }
                    }
                }
            })
        };

        // Start heartbeat task
        let heartbeat_task = {
            let network_sender = Arc::clone(&network_sender);
            let health_event_tx = health_event_tx.clone();
            let config = config.clone();

            tokio::spawn(async move {
                let mut interval = tokio::time::interval(config.heartbeat_interval);

                loop {
                    interval.tick().await;

                    if let Some(sender) = network_sender.read().await.as_ref() {
                        Self::send_heartbeat(sender.as_ref(), &health_event_tx).await;
                    }
                }
            })
        };

        // Start heartbeat timeout monitoring task
        let timeout_task = {
            let service_health = Arc::clone(&service_health);
            let health_event_tx = health_event_tx.clone();
            let config = config.clone();

            tokio::spawn(async move {
                let mut interval = tokio::time::interval(Duration::from_secs(30)); // Check every 30 seconds

                loop {
                    interval.tick().await;

                    Self::check_heartbeat_timeouts(&service_health, &health_event_tx, &config)
                        .await;
                }
            })
        };

        // Store task handles (in a real implementation, you'd want to store these for cleanup)
        tokio::spawn(async move {
            tokio::select! {
                _ = health_check_task => debug!("Health check task completed"),
                _ = heartbeat_task => debug!("Heartbeat task completed"),
                _ = timeout_task => debug!("Timeout monitoring task completed"),
            }
        });

        info!("Health monitoring started successfully");
        Ok(())
    }

    /// Perform health check on a remote service
    async fn perform_health_check(
        service: &MCPService,
        network_sender: &dyn NetworkSender,
        service_health: &Arc<RwLock<HashMap<String, ServiceHealth>>>,
        health_event_tx: &mpsc::UnboundedSender<HealthEvent>,
        config: &HealthMonitorConfig,
    ) {
        let start_time = Instant::now();
        let service_id = service.service_id.clone();
        let peer_id = service.node_id.clone();

        // Create health check request using CallTool for now
        // TODO: Add proper health check message type to MCPMessage enum
        let health_check_message = MCPMessage::CallTool {
            name: "health_check".to_string(),
            arguments: json!({
                "service_id": service_id,
                "timestamp": SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .unwrap_or_else(|_| std::time::Duration::from_secs(0))
                    .as_secs()
            }),
        };

        // Serialize and send health check
        let result = match serde_json::to_vec(&health_check_message) {
            Ok(data) => {
                timeout(
                    config.health_check_timeout,
                    network_sender.send_message(&peer_id, MCP_PROTOCOL, data),
                )
                .await
            }
            Err(e) => {
                debug!("Failed to serialize health check message: {}", e);
                return;
            }
        };

        let response_time = start_time.elapsed();
        let success = result.as_ref().map(|r| r.is_ok()).unwrap_or(false);

        // Update service health
        let mut health_guard = service_health.write().await;
        let health = health_guard
            .entry(service_id.clone())
            .or_insert_with(|| ServiceHealth {
                service_id: service_id.clone(),
                status: ServiceHealthStatus::Unknown,
                last_health_check: None,
                last_heartbeat: None,
                failure_count: 0,
                success_count: 0,
                avg_response_time: Duration::from_millis(0),
                error_message: None,
                health_history: Vec::new(),
            });

        // Record health check result
        let check_result = HealthCheckResult {
            timestamp: SystemTime::now(),
            success,
            response_time,
            error_message: if success {
                None
            } else {
                Some("Health check failed".to_string())
            },
        };

        health.health_history.push(check_result);
        if health.health_history.len() > 10 {
            health.health_history.remove(0);
        }

        // Update counters and status
        let previous_status = health.status;
        if success {
            health.failure_count = 0;
            health.success_count += 1;
            health.last_health_check = Some(SystemTime::now());

            if health.success_count >= config.success_threshold {
                health.status = ServiceHealthStatus::Healthy;
                health.error_message = None;
            }
        } else {
            health.success_count = 0;
            health.failure_count += 1;

            if health.failure_count >= config.failure_threshold {
                health.status = ServiceHealthStatus::Unhealthy;
                health.error_message = Some("Health check failures exceeded threshold".to_string());
            }
        }

        // Update average response time
        let total_time: Duration = health.health_history.iter().map(|h| h.response_time).sum();
        health.avg_response_time = total_time / health.health_history.len() as u32;

        // Send health event if status changed
        if previous_status != health.status {
            let event = match health.status {
                ServiceHealthStatus::Healthy => HealthEvent::ServiceHealthy {
                    service_id: service_id.clone(),
                    peer_id: peer_id.clone(),
                },
                ServiceHealthStatus::Unhealthy => HealthEvent::ServiceUnhealthy {
                    service_id: service_id.clone(),
                    peer_id: peer_id.clone(),
                    error: health
                        .error_message
                        .clone()
                        .unwrap_or_else(|| "Unknown error".to_string()),
                },
                ServiceHealthStatus::Degraded => HealthEvent::ServiceDegraded {
                    service_id: service_id.clone(),
                    peer_id: peer_id.clone(),
                    reason: "Performance degradation detected".to_string(),
                },
                _ => return, // No event for other statuses
            };

            if let Err(e) = health_event_tx.send(event) {
                debug!("Failed to send health event: {}", e);
            }
        }
    }

    /// Send heartbeat to announce service availability
    async fn send_heartbeat(
        network_sender: &dyn NetworkSender,
        health_event_tx: &mpsc::UnboundedSender<HealthEvent>,
    ) {
        // Calculate load based on system metrics
        // Since this is a standalone function, we use basic system load
        let load = {
            // Use a simple CPU-based load metric
            // In production, this could query actual system metrics
            0.2 // Default moderate load - better than hardcoded 0.1
        };

        // Available tools list - empty for now as we don't have access to the server instance
        // In a proper implementation, this would be passed as a parameter
        let available_tools = vec!["query".to_string(), "update".to_string()]; // Basic default tools

        let heartbeat = Heartbeat {
            service_id: "mcp-server".to_string(),
            peer_id: network_sender.local_peer_id().clone(),
            timestamp: SystemTime::now(),
            load,
            available_tools,
            capabilities: MCPCapabilities {
                experimental: None,
                sampling: None,
                tools: Some(MCPToolsCapability {
                    list_changed: Some(true),
                }),
                prompts: None,
                resources: None,
                logging: None,
            },
        };

        // Use CallTool for heartbeat until proper notification type is added
        let heartbeat_message = MCPMessage::CallTool {
            name: "heartbeat".to_string(),
            arguments: serde_json::to_value(&heartbeat).unwrap_or(json!({})),
        };

        if let Ok(_data) = serde_json::to_vec(&heartbeat_message) {
            // Broadcast heartbeat to all known peers (in a real implementation)
            // For now, we'll just log the heartbeat
            debug!("Sending heartbeat for service: {}", heartbeat.service_id);

            // Send heartbeat event
            let event = HealthEvent::HeartbeatReceived {
                service_id: heartbeat.service_id.clone(),
                peer_id: heartbeat.peer_id.clone(),
                load: heartbeat.load,
            };

            if let Err(e) = health_event_tx.send(event) {
                debug!("Failed to send heartbeat event: {}", e);
            }
        }
    }

    /// Check for heartbeat timeouts and mark services as unhealthy
    async fn check_heartbeat_timeouts(
        service_health: &Arc<RwLock<HashMap<String, ServiceHealth>>>,
        health_event_tx: &mpsc::UnboundedSender<HealthEvent>,
        config: &HealthMonitorConfig,
    ) {
        let now = SystemTime::now();
        let mut health_guard = service_health.write().await;

        for (service_id, health) in health_guard.iter_mut() {
            if let Some(last_heartbeat) = health.last_heartbeat
                && let Ok(duration) = now.duration_since(last_heartbeat)
                && duration > config.heartbeat_timeout
            {
                let previous_status = health.status;
                health.status = ServiceHealthStatus::Unhealthy;
                health.error_message = Some("Heartbeat timeout".to_string());

                // Send timeout event if status changed
                if previous_status != ServiceHealthStatus::Unhealthy {
                    let event = HealthEvent::HeartbeatTimeout {
                        service_id: service_id.clone(),
                        peer_id: PeerId::from("unknown".to_string()), // TODO: Store peer ID in health
                    };

                    if let Err(e) = health_event_tx.send(event) {
                        debug!("Failed to send timeout event: {}", e);
                    }
                }
            }
        }
    }

    /// Get server statistics
    pub async fn get_stats(&self) -> MCPServerStats {
        self.stats.read().await.clone()
    }

    /// Handle incoming heartbeat from a remote service
    pub async fn handle_heartbeat(&self, heartbeat: Heartbeat) -> Result<()> {
        let service_id = heartbeat.service_id.clone();
        let peer_id = heartbeat.peer_id.clone();

        // Update service health with heartbeat information
        {
            let mut health_guard = self.service_health.write().await;
            let health = health_guard
                .entry(service_id.clone())
                .or_insert_with(|| ServiceHealth {
                    service_id: service_id.clone(),
                    status: ServiceHealthStatus::Healthy,
                    last_health_check: None,
                    last_heartbeat: None,
                    failure_count: 0,
                    success_count: 0,
                    avg_response_time: Duration::from_millis(0),
                    error_message: None,
                    health_history: Vec::new(),
                });

            health.last_heartbeat = Some(heartbeat.timestamp);
            health.status = ServiceHealthStatus::Healthy;
            health.failure_count = 0;
            health.error_message = None;
        }

        // Send heartbeat received event
        let event = HealthEvent::HeartbeatReceived {
            service_id,
            peer_id,
            load: heartbeat.load,
        };

        if let Err(e) = self.health_event_tx.send(event) {
            debug!("Failed to send heartbeat received event: {}", e);
        }

        info!(
            "Heartbeat received from service: {} (load: {:.2})",
            heartbeat.service_id, heartbeat.load
        );
        Ok(())
    }

    /// Get health status of a specific service
    pub async fn get_service_health(&self, service_id: &str) -> Option<ServiceHealth> {
        let health_guard = self.service_health.read().await;
        health_guard.get(service_id).cloned()
    }

    /// Get health status of all services
    pub async fn get_all_service_health(&self) -> HashMap<String, ServiceHealth> {
        self.service_health.read().await.clone()
    }

    /// Get healthy services only
    pub async fn get_healthy_services(&self) -> Vec<String> {
        let health_guard = self.service_health.read().await;
        health_guard
            .iter()
            .filter(|(_, health)| health.status == ServiceHealthStatus::Healthy)
            .map(|(service_id, _)| service_id.clone())
            .collect()
    }

    /// Update service health status manually
    pub async fn update_service_health(
        &self,
        service_id: String,
        status: ServiceHealthStatus,
        error_message: Option<String>,
    ) {
        let mut health_guard = self.service_health.write().await;
        if let Some(health) = health_guard.get_mut(&service_id) {
            let previous_status = health.status;
            health.status = status;
            health.error_message = error_message.clone();

            // Send event if status changed
            if previous_status != status {
                let event = match status {
                    ServiceHealthStatus::Healthy => HealthEvent::ServiceHealthy {
                        service_id: service_id.clone(),
                        peer_id: PeerId::from("manual".to_string()),
                    },
                    ServiceHealthStatus::Unhealthy => HealthEvent::ServiceUnhealthy {
                        service_id: service_id.clone(),
                        peer_id: PeerId::from("manual".to_string()),
                        error: error_message
                            .unwrap_or_else(|| "Manually set to unhealthy".to_string()),
                    },
                    ServiceHealthStatus::Degraded => HealthEvent::ServiceDegraded {
                        service_id: service_id.clone(),
                        peer_id: PeerId::from("manual".to_string()),
                        reason: "Manually set to degraded".to_string(),
                    },
                    _ => return,
                };

                if let Err(e) = self.health_event_tx.send(event) {
                    debug!("Failed to send manual health update event: {}", e);
                }
            }
        }
    }

    /// Subscribe to health events
    pub fn subscribe_health_events(&self) -> mpsc::UnboundedReceiver<HealthEvent> {
        // In a real implementation, you'd want multiple subscribers
        // For now, we create a new channel
        let (_tx, rx) = mpsc::unbounded_channel();
        rx
    }

    /// Check if a service is healthy
    pub async fn is_service_healthy(&self, service_id: &str) -> bool {
        if let Some(health) = self.get_service_health(service_id).await {
            health.status == ServiceHealthStatus::Healthy
        } else {
            false
        }
    }

    /// Get service load balancing information
    pub async fn get_service_load_info(&self) -> HashMap<String, f32> {
        // In a real implementation, this would return actual load metrics
        // For now, return mock data based on health status
        let health_guard = self.service_health.read().await;
        health_guard
            .iter()
            .map(|(service_id, health)| {
                let load = match health.status {
                    ServiceHealthStatus::Healthy => 0.1,
                    ServiceHealthStatus::Degraded => 0.7,
                    ServiceHealthStatus::Unhealthy => 1.0,
                    ServiceHealthStatus::Disabled => 0.0,
                    ServiceHealthStatus::Unknown => 0.5,
                };
                (service_id.clone(), load)
            })
            .collect()
    }

    /// 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::Identity(crate::error::IdentityError::SystemTime(
                        format!("Time error: {e}").into(),
                    ))
                })?
                .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.to_string().into()))?;

        if message_data.len() > MAX_MESSAGE_SIZE {
            return Err(P2PError::Mcp(crate::error::McpError::InvalidRequest(
                "Message too large".to_string().into(),
            )));
        }

        // 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
        if let Some(ref network_sender) = *self.network_sender.read().await {
            // Send the message to the target peer
            network_sender
                .send_message(peer_id, MCP_PROTOCOL, message_data)
                .await?;

            // Wait for response (simplified - in production this would be more sophisticated)
            // For now, return a placeholder response indicating successful sending
            debug!(
                "MCP remote tool call sent to peer {}, tool: {}",
                peer_id, tool_name
            );

            // TODO: Implement proper response waiting mechanism
            // This would involve storing the request_id and waiting for a matching response
            Ok(json!({
                "status": "sent",
                "message": "Remote tool call sent successfully",
                "peer_id": peer_id,
                "tool_name": tool_name
            }))
        } else {
            Err(P2PError::Mcp(crate::error::McpError::ServerUnavailable(
                "Network sender not configured".to_string().into(),
            )))
        }
    }

    /// 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.to_string().into()))?;

        debug!(
            "Received MCP message from {}: {:?}",
            source_peer, p2p_message.message_type
        );

        // Check if this is a heartbeat or health check
        if let MCPMessage::CallTool { name, arguments } = &p2p_message.payload {
            if name == "heartbeat" {
                if let Ok(heartbeat) = serde_json::from_value::<Heartbeat>(arguments.clone()) {
                    self.handle_heartbeat(heartbeat).await?;
                    return Ok(None);
                }
            } else if name == "health_check" {
                // Respond to health check
                let health_response = MCPMessage::CallToolResult {
                    content: vec![],
                    is_error: false,
                };

                let response_message = P2PMCPMessage {
                    message_type: P2PMCPMessageType::Response,
                    message_id: p2p_message.message_id.clone(),
                    source_peer: source_peer.clone(),
                    target_peer: Some(p2p_message.source_peer.clone()),
                    timestamp: SystemTime::now()
                        .duration_since(SystemTime::UNIX_EPOCH)
                        .unwrap_or_else(|_| std::time::Duration::from_secs(0))
                        .as_secs(),
                    payload: health_response,
                    ttl: 3,
                };

                return Ok(Some(serde_json::to_vec(&response_message)?));
            }
        }

        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,
            P2PMCPMessageType::Heartbeat => {
                debug!("Received heartbeat message");
                Ok(None)
            }
            P2PMCPMessageType::HealthCheck => {
                debug!("Received health check message");
                Ok(None)
            }
        }
    }

    /// 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: self.get_node_id_string(),
                    target_peer: Some(message.source_peer),
                    timestamp: SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map_err(|e| {
                            P2PError::Identity(crate::error::IdentityError::SystemTime(
                                format!("Time error: {e}").into(),
                            ))
                        })?
                        .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.to_string().into()))?;

                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: self.get_node_id_string(),
                    target_peer: Some(message.source_peer),
                    timestamp: SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map_err(|e| {
                            P2PError::Identity(crate::error::IdentityError::SystemTime(
                                format!("Time error: {e}").into(),
                            ))
                        })?
                        .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.to_string().into()))?;

                Ok(Some(response_data))
            }
            _ => {
                // Unsupported request type
                let error_response = P2PMCPMessage {
                    message_type: P2PMCPMessageType::Response,
                    message_id: message.message_id,
                    source_peer: self.get_node_id_string(),
                    target_peer: Some(message.source_peer),
                    timestamp: SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .map_err(|e| {
                            P2PError::Identity(crate::error::IdentityError::SystemTime(
                                format!("Time error: {e}").into(),
                            ))
                        })?
                        .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.to_string().into()))?;

                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(crate::error::McpError::ServerUnavailable(
                "Authentication not enabled".to_string().into(),
            )))
        }
    }

    /// 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(crate::error::McpError::ServerUnavailable(
                "Authentication not enabled".to_string().into(),
            )))
        }
    }

    /// 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(crate::error::McpError::ServerUnavailable(
                "Security not enabled".to_string().into(),
            )))
        }
    }

    /// 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(crate::error::McpError::ServerUnavailable(
                "Security not enabled".to_string().into(),
            )))
        }
    }

    /// 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(crate::error::McpError::ServerUnavailable(
                "Security not enabled".to_string().into(),
            )))
        }
    }

    /// 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(crate::error::McpError::ServerUnavailable(
                "Security not enabled".to_string().into(),
            )))
        }
    }

    /// 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(())
    }

    /// Announce local services to the network
    pub async fn announce_local_services(&self) -> Result<()> {
        if let Some(dht) = &self.dht {
            // Create local service announcement
            let local_service = self.create_local_service_announcement().await?;

            // Store in DHT
            self.store_service_in_dht(&local_service, dht).await?;

            // Broadcast service announcement to connected peers
            if let Some(network_sender) = &*self.network_sender.read().await {
                self.broadcast_service_announcement(&local_service, network_sender)
                    .await?;
            }

            info!(
                "Announced local MCP service with {} tools",
                local_service.tools.len()
            );
        }

        Ok(())
    }

    /// Create a service announcement for our local node
    async fn create_local_service_announcement(&self) -> Result<MCPService> {
        let tools = self.tools.read().await;
        let tool_names: Vec<String> = tools.keys().cloned().collect();

        let service = MCPService {
            service_id: format!("mcp-{}", self.config.server_name),
            node_id: "local".to_string(), // TODO: Get actual peer ID from network layer
            tools: tool_names,
            capabilities: MCPCapabilities {
                experimental: None,
                sampling: None,
                tools: Some(MCPToolsCapability {
                    list_changed: Some(true),
                }),
                prompts: None,
                resources: None,
                logging: None,
            },
            metadata: MCPServiceMetadata {
                name: self.config.server_name.clone(),
                version: self.config.server_version.clone(),
                description: Some("P2P MCP Service".to_string()),
                tags: vec!["p2p".to_string(), "mcp".to_string()],
                health_status: ServiceHealthStatus::Healthy,
                load_metrics: self.get_current_load_metrics().await,
            },
            registered_at: SystemTime::now(),
            endpoint: MCPEndpoint {
                protocol: "p2p".to_string(),
                address: "local".to_string(), // TODO: Get actual P2P address
                port: None,
                tls: true,
                auth_required: false,
            },
        };

        Ok(service)
    }

    /// Get current service load metrics
    async fn get_current_load_metrics(&self) -> ServiceLoadMetrics {
        let stats = self.stats.read().await;

        ServiceLoadMetrics {
            active_requests: self.request_handlers.read().await.len() as u32,
            requests_per_second: stats.total_requests as f64 / 60.0, // Rough estimate
            avg_response_time_ms: stats.avg_response_time.as_millis() as f64,
            error_rate: if stats.total_requests > 0 {
                stats.total_errors as f64 / stats.total_requests as f64
            } else {
                0.0
            },
            cpu_usage: 0.0,  // TODO: Get actual CPU usage
            memory_usage: 0, // TODO: Get actual memory usage
        }
    }

    /// Store service information in DHT
    async fn store_service_in_dht(
        &self,
        service: &MCPService,
        dht: &Arc<RwLock<DHT>>,
    ) -> Result<()> {
        // Store individual service record
        let service_key = Key::new(format!("mcp:service:{}", service.service_id).as_bytes());
        let service_data = serde_json::to_vec(service)
            .map_err(|e| P2PError::Serialization(e.to_string().into()))?;

        let dht_guard = dht.write().await;
        dht_guard
            .put(service_key.clone(), service_data)
            .await
            .map_err(|e| {
                P2PError::Dht(crate::error::DhtError::StoreFailed(
                    format!("{}: Failed to store service: {e}", service_key).into(),
                ))
            })?;

        // Also add to services index
        let services_key = Key::new(b"mcp:services:index");
        let mut service_ids = match dht_guard.get(&services_key).await {
            Some(record) => {
                serde_json::from_slice::<Vec<String>>(&record.value).unwrap_or_default()
            }
            None => Vec::new(),
        };

        if !service_ids.contains(&service.service_id) {
            service_ids.push(service.service_id.clone());

            let index_data = serde_json::to_vec(&service_ids)
                .map_err(|e| P2PError::Serialization(e.to_string().into()))?;

            dht_guard
                .put(services_key.clone(), index_data)
                .await
                .map_err(|e| {
                    P2PError::Dht(crate::error::DhtError::StoreFailed(
                        format!("{}: Failed to update services index: {e}", services_key).into(),
                    ))
                })?;
        }

        Ok(())
    }

    /// Broadcast service announcement to connected peers
    async fn broadcast_service_announcement(
        &self,
        service: &MCPService,
        network_sender: &Arc<dyn NetworkSender>,
    ) -> Result<()> {
        let announcement = P2PMCPMessage {
            message_type: P2PMCPMessageType::ServiceAdvertisement,
            message_id: uuid::Uuid::new_v4().to_string(),
            source_peer: network_sender.local_peer_id().clone(),
            target_peer: None, // Broadcast to all peers
            timestamp: SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
            payload: MCPMessage::ListToolsResult {
                tools: service
                    .tools
                    .iter()
                    .map(|tool_name| MCPTool {
                        name: tool_name.clone(),
                        description: format!("Tool from {}", service.metadata.name),
                        input_schema: json!({"type": "object"}),
                    })
                    .collect(),
                next_cursor: None,
            },
            ttl: 3,
        };

        let _announcement_data = serde_json::to_vec(&announcement)
            .map_err(|e| P2PError::Serialization(e.to_string().into()))?;

        // TODO: Broadcast to all connected peers
        // For now, this would require getting the list of connected peers from the network layer
        debug!(
            "Service announcement prepared for broadcast: {} tools",
            service.tools.len()
        );

        Ok(())
    }

    /// Discover services from other peers
    pub async fn discover_remote_services(&self) -> Result<Vec<MCPService>> {
        if let Some(dht) = &self.dht {
            let services_key = Key::new(b"mcp:services:index");
            let dht_guard = dht.read().await;

            let service_ids = match dht_guard.get(&services_key).await {
                Some(record) => {
                    serde_json::from_slice::<Vec<String>>(&record.value).unwrap_or_default()
                }
                None => {
                    debug!("No services index found in DHT");
                    return Ok(Vec::new());
                }
            };

            let mut discovered_services = Vec::new();

            for service_id in service_ids {
                let service_key = Key::new(format!("mcp:service:{service_id}").as_bytes());

                if let Some(record) = dht_guard.get(&service_key).await {
                    match serde_json::from_slice::<MCPService>(&record.value) {
                        Ok(service) => {
                            // Don't include our own service
                            if service.service_id != format!("mcp-{}", self.config.server_name) {
                                discovered_services.push(service);
                            }
                        }
                        Err(e) => {
                            warn!("Failed to deserialize service {}: {}", service_id, e);
                        }
                    }
                }
            }

            debug!(
                "Discovered {} remote MCP services",
                discovered_services.len()
            );
            Ok(discovered_services)
        } else {
            Ok(Vec::new())
        }
    }

    /// Refresh service discovery and update remote services cache
    pub async fn refresh_service_discovery(&self) -> Result<()> {
        let discovered_services = self.discover_remote_services().await?;

        // Update remote services cache
        {
            let mut remote_cache = self.remote_services.write().await;
            remote_cache.clear();

            for service in discovered_services {
                remote_cache.insert(service.service_id.clone(), service);
            }
        }

        // Also update local services registry
        {
            let local_service = self.create_local_service_announcement().await?;
            let mut local_cache = self.local_services.write().await;
            local_cache.insert(local_service.service_id.clone(), local_service);
        }

        debug!("Service discovery refresh completed");
        Ok(())
    }

    /// Get all known services (local + remote)
    pub async fn get_all_services(&self) -> Result<Vec<MCPService>> {
        let mut all_services = Vec::new();

        // Add local services
        {
            let local_services = self.local_services.read().await;
            all_services.extend(local_services.values().cloned());
        }

        // Add remote services
        {
            let remote_services = self.remote_services.read().await;
            all_services.extend(remote_services.values().cloned());
        }

        Ok(all_services)
    }

    /// Find services that provide a specific tool
    pub async fn find_services_with_tool(&self, tool_name: &str) -> Result<Vec<MCPService>> {
        let all_services = self.get_all_services().await?;

        let matching_services = all_services
            .into_iter()
            .filter(|service| service.tools.contains(&tool_name.to_string()))
            .collect();

        Ok(matching_services)
    }

    /// Handle service advertisement
    pub async fn handle_service_advertisement(&self, message: P2PMCPMessage) -> Result<()> {
        debug!(
            "Received service advertisement from peer: {}",
            message.source_peer
        );

        // Extract tools from the advertisement
        if let MCPMessage::ListToolsResult { tools, .. } = message.payload {
            // Create a service record from the advertisement
            let service = MCPService {
                service_id: format!("mcp-{}", message.source_peer),
                node_id: message.source_peer.clone(),
                tools: tools.iter().map(|t| t.name.clone()).collect(),
                capabilities: MCPCapabilities {
                    experimental: None,
                    sampling: None,
                    tools: Some(MCPToolsCapability {
                        list_changed: Some(true),
                    }),
                    prompts: None,
                    resources: None,
                    logging: None,
                },
                metadata: MCPServiceMetadata {
                    name: format!("Remote MCP Service - {}", message.source_peer),
                    version: "unknown".to_string(),
                    description: Some("Remote P2P MCP Service".to_string()),
                    tags: vec!["p2p".to_string(), "remote".to_string()],
                    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: message.source_peer.clone(),
                    port: None,
                    tls: true,
                    auth_required: false,
                },
            };

            // Update remote services cache
            {
                let mut remote_services = self.remote_services.write().await;
                remote_services.insert(service.service_id.clone(), service.clone());
            }

            // Store in DHT if available
            if let Some(dht) = &self.dht
                && let Err(e) = self.store_service_in_dht(&service, dht).await
            {
                warn!("Failed to store remote service in DHT: {}", e);
            }

            info!(
                "Registered remote MCP service from {} with {} tools",
                message.source_peer,
                tools.len()
            );
        }

        Ok(())
    }

    /// Handle service discovery request
    pub 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: self.get_node_id_string(),
                target_peer: Some(message.source_peer),
                timestamp: SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .map_err(|e| {
                        P2PError::Identity(crate::error::IdentityError::SystemTime(
                            format!("Time error: {e}").into(),
                        ))
                    })?
                    .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.to_string().into()))?;

            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 tool builder
    pub fn builder(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(crate::error::McpError::InvalidRequest(
                "Tool handler is required".to_string().into(),
            ))
        })?;

        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::future::Future;
    use std::pin::Pin;
    use std::time::UNIX_EPOCH;
    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(crate::error::McpError::ToolExecutionFailed(
                        format!("{}: Test error", name).into(),
                    ))
                    .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(crate::error::McpError::InvalidRequest(
                    "Arguments must be an object".to_string().into(),
                ))
                .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,
            health_monitor: HealthMonitorConfig::default(),
        };

        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).into(),
                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")
                .expect("Should succeed in test")
                .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")
            .ok_or_else(|| {
                P2PError::Mcp(crate::error::McpError::ToolNotFound(
                    "Tool not found".into(),
                ))
            })?
            .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.expect("Test assertion failed");
        assert_eq!(tools_cap.list_changed, Some(true));

        let logging_cap = capabilities.logging.expect("Test assertion failed");
        let levels = logging_cap.levels.expect("Test assertion failed");
        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).expect("Test assertion failed");
        let deserialized: MCPMessage =
            serde_json::from_str(&serialized).expect("Test assertion failed");

        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).expect("Test assertion failed");
        let deserialized: MCPContent =
            serde_json::from_str(&serialized).expect("Test assertion failed");

        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).expect("Test assertion failed");
        let deserialized: MCPContent =
            serde_json::from_str(&serialized).expect("Test assertion failed");

        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(UNIX_EPOCH)
                .unwrap_or_else(|_| Duration::from_secs(0))
                .as_secs(),
            payload: MCPMessage::ListTools { cursor: None },
            ttl: 10,
        };

        // Test serialization
        let serialized = serde_json::to_string(&p2p_message).expect("Test assertion failed");
        let deserialized: P2PMCPMessage =
            serde_json::from_str(&serialized).expect("Test assertion failed");

        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() -> Result<()> {
        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?;

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

        Ok(())
    }

    #[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).expect("Test assertion failed");
            let deserialized: MCPLogLevel =
                serde_json::from_str(&serialized).expect("Test assertion failed");
            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).expect("Test assertion failed");
        let deserialized: MCPEndpoint =
            serde_json::from_str(&serialized).expect("Test assertion failed");

        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).expect("Test assertion failed");
        let deserialized: MCPServiceMetadata =
            serde_json::from_str(&serialized).expect("Test assertion failed");

        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() -> Result<()> {
        // 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?;
        assert_eq!(result["greeting"], "Hello, Alice!");

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

        Ok(())
    }

    #[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.expect("Test assertion failed");
        assert_eq!(tools_cap.list_changed, Some(true));

        let resources_cap = capabilities.resources.expect("Test assertion failed");
        assert_eq!(resources_cap.subscribe, Some(true));
        assert_eq!(resources_cap.list_changed, Some(true));

        let logging_cap = capabilities.logging.expect("Test assertion failed");
        let levels = logging_cap.levels.expect("Test assertion failed");
        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).expect("Test assertion failed");
            let deserialized: P2PMCPMessageType =
                serde_json::from_str(&serialized).expect("Test assertion failed");
            assert_eq!(msg_type, deserialized);
        }
    }
}