instance.rs

//! Server instance representation
//!
//! Each (space_id, server_id) pair gets its own isolated ServerInstance.
//! No sharing between spaces - this is a security boundary.

use std::collections::HashMap;
use std::time::Instant;

use mcpmux_core::DomainEvent;
use parking_lot::RwLock;
use rmcp::model::{ClientCapabilities, ClientInfo, Implementation};
use rmcp::service::{NotificationContext, RunningService};
use rmcp::RoleClient;
use serde::{Deserialize, Serialize};
use tracing::{debug, info, warn};
use uuid::Uuid;

// Re-export TransportType from mcpmux-core as the single source of truth
pub use mcpmux_core::TransportType;

/// Type alias for the MCP client service
pub type McpClient = RunningService<RoleClient, McpClientHandler>;

/// Client handler for MCP connections
#[derive(Clone, Debug)]
pub struct McpClientHandler {
    info: ClientInfo,
    server_id: String,
    space_id: Uuid,
    event_tx: Option<tokio::sync::broadcast::Sender<DomainEvent>>,
}

impl McpClientHandler {
    pub fn new(
        server_id: &str,
        space_id: Uuid,
        event_tx: Option<tokio::sync::broadcast::Sender<DomainEvent>>,
    ) -> Self {
        Self {
            info: ClientInfo {
                protocol_version: Default::default(),
                capabilities: ClientCapabilities::default(),
                client_info: Implementation {
                    name: format!("mcpmux-{}", server_id),
                    version: env!("CARGO_PKG_VERSION").to_string(),
                    title: Some("McpMux Gateway".to_string()),
                    icons: None,
                    website_url: None,
                    ..Default::default()
                },
                meta: None,
            },
            server_id: server_id.to_string(),
            space_id,
            event_tx,
        }
    }
}

impl rmcp::ClientHandler for McpClientHandler {
    fn get_info(&self) -> ClientInfo {
        self.info.clone()
    }

    // Handle notifications from backend MCP servers
    fn on_tool_list_changed(
        &self,
        _context: NotificationContext<RoleClient>,
    ) -> impl std::future::Future<Output = ()> + Send + '_ {
        let server_id = self.server_id.clone();
        let space_id = self.space_id;
        let event_tx = self.event_tx.clone();
        async move {
            info!(
                server_id = %server_id,
                space_id = %space_id,
                "[McpClientHandler] 🔔 Backend server sent tools/list_changed notification"
            );

            if let Some(tx) = &event_tx {
                let event = DomainEvent::ToolsChanged {
                    server_id: server_id.clone(),
                    space_id,
                };
                if let Err(e) = tx.send(event) {
                    warn!(
                        server_id = %server_id,
                        space_id = %space_id,
                        error = %e,
                        "[McpClientHandler] ⚠️ Failed to emit ToolsChanged event (no subscribers)"
                    );
                } else {
                    debug!(
                        server_id = %server_id,
                        space_id = %space_id,
                        "[McpClientHandler] ✅ Emitted ToolsChanged event to domain event bus"
                    );
                }
            } else {
                warn!(
                    server_id = %server_id,
                    space_id = %space_id,
                    "[McpClientHandler] ⚠️ No event_tx available - cannot forward tools/list_changed"
                );
            }
        }
    }

    fn on_prompt_list_changed(
        &self,
        _context: NotificationContext<RoleClient>,
    ) -> impl std::future::Future<Output = ()> + Send + '_ {
        let server_id = self.server_id.clone();
        let space_id = self.space_id;
        let event_tx = self.event_tx.clone();
        async move {
            info!(
                server_id = %server_id,
                space_id = %space_id,
                "[McpClientHandler] 🔔 Backend server sent prompts/list_changed notification"
            );

            if let Some(tx) = &event_tx {
                let event = DomainEvent::PromptsChanged {
                    server_id: server_id.clone(),
                    space_id,
                };
                if let Err(e) = tx.send(event) {
                    warn!(
                        server_id = %server_id,
                        space_id = %space_id,
                        error = %e,
                        "[McpClientHandler] ⚠️ Failed to emit PromptsChanged event (no subscribers)"
                    );
                } else {
                    debug!(
                        server_id = %server_id,
                        space_id = %space_id,
                        "[McpClientHandler] ✅ Emitted PromptsChanged event to domain event bus"
                    );
                }
            } else {
                warn!(
                    server_id = %server_id,
                    space_id = %space_id,
                    "[McpClientHandler] ⚠️ No event_tx available - cannot forward prompts/list_changed"
                );
            }
        }
    }

    fn on_resource_list_changed(
        &self,
        _context: NotificationContext<RoleClient>,
    ) -> impl std::future::Future<Output = ()> + Send + '_ {
        let server_id = self.server_id.clone();
        let space_id = self.space_id;
        let event_tx = self.event_tx.clone();
        async move {
            info!(
                server_id = %server_id,
                space_id = %space_id,
                "[McpClientHandler] 🔔 Backend server sent resources/list_changed notification"
            );

            if let Some(tx) = &event_tx {
                let event = DomainEvent::ResourcesChanged {
                    server_id: server_id.clone(),
                    space_id,
                };
                if let Err(e) = tx.send(event) {
                    warn!(
                        server_id = %server_id,
                        space_id = %space_id,
                        error = %e,
                        "[McpClientHandler] ⚠️ Failed to emit ResourcesChanged event (no subscribers)"
                    );
                } else {
                    debug!(
                        server_id = %server_id,
                        space_id = %space_id,
                        "[McpClientHandler] ✅ Emitted ResourcesChanged event to domain event bus"
                    );
                }
            } else {
                warn!(
                    server_id = %server_id,
                    space_id = %space_id,
                    "[McpClientHandler] ⚠️ No event_tx available - cannot forward resources/list_changed"
                );
            }
        }
    }
}

/// Instance key - identifies a server instance for debugging/logging.
/// Note: Actual instance lookup uses (space_id, server_id) tuple in PoolService.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct InstanceKey {
    /// Space ID that owns this instance
    pub space_id: Uuid,
    /// Human-readable description (e.g., "http:https://mcp.example.com")
    pub description: String,
}

impl InstanceKey {
    /// Create instance key for STDIO transport.
    pub fn stdio(
        space_id: Uuid,
        command: &str,
        _args: &[String],
        _env: &HashMap<String, String>,
    ) -> Self {
        Self {
            space_id,
            description: format!("stdio:{}", command),
        }
    }

    /// Create instance key for HTTP transport.
    pub fn http(space_id: Uuid, url: &str, _headers: &HashMap<String, String>) -> Self {
        Self {
            space_id,
            description: format!("http:{}", url),
        }
    }
}

/// Connection state for a server instance.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum InstanceState {
    /// Not connected
    Disconnected,
    /// Currently connecting
    Connecting,
    /// Connected and ready
    Connected,
    /// Connection failed
    Failed,
    /// OAuth flow in progress
    OAuthPending,
}

/// Features discovered from an MCP server.
#[derive(Debug, Clone, Default)]
pub struct DiscoveredFeatures {
    /// Available tools
    pub tools: Vec<serde_json::Value>,
    /// Available prompts  
    pub prompts: Vec<serde_json::Value>,
    /// Available resources
    pub resources: Vec<serde_json::Value>,
}

/// Statistics for a server instance.
#[derive(Debug, Clone)]
pub struct InstanceStats {
    /// Current state
    pub state: InstanceState,
    /// When instance was connected (if connected)
    pub connected_at: Option<Instant>,
    /// Last connection attempt
    pub last_attempt: Option<Instant>,
    /// Consecutive failures
    pub consecutive_failures: u32,
    /// Total requests served
    pub requests_served: u64,
    /// Last error message
    pub last_error: Option<String>,
}

impl Default for InstanceStats {
    fn default() -> Self {
        Self {
            state: InstanceState::Disconnected,
            connected_at: None,
            last_attempt: None,
            consecutive_failures: 0,
            requests_served: 0,
            last_error: None,
        }
    }
}

/// An isolated MCP server instance.
/// Each (space_id, server_id) pair gets exactly one instance - no sharing.
pub struct ServerInstance {
    /// Instance key for debugging/logging
    pub key: InstanceKey,
    /// Server ID from the registry
    pub server_id: String,
    /// Transport type  
    pub transport_type: TransportType,
    /// Connection statistics
    pub stats: RwLock<InstanceStats>,
    /// Discovered features (populated after connection)
    pub features: RwLock<Option<DiscoveredFeatures>>,
    /// The actual MCP client connection
    client: RwLock<Option<McpClientConnection>>,
}

/// The actual MCP client connection.
pub enum McpClientConnection {
    /// STDIO transport - child process with MCP client
    Stdio { client: McpClient },
    /// HTTP transport - streamable HTTP
    Http { client: McpClient },
}

impl McpClientConnection {
    /// Get the MCP client for issuing requests.
    pub fn client(&self) -> Option<&McpClient> {
        match self {
            Self::Stdio { client } => Some(client),
            Self::Http { client } => Some(client),
        }
    }
}

impl ServerInstance {
    /// Create a new disconnected instance.
    pub fn new(key: InstanceKey, server_id: String, transport_type: TransportType) -> Self {
        Self {
            key,
            server_id,
            transport_type,
            stats: RwLock::new(InstanceStats::default()),
            features: RwLock::new(None),
            client: RwLock::new(None),
        }
    }

    /// Get the current state.
    pub fn state(&self) -> InstanceState {
        self.stats.read().state
    }

    /// Check if connected and healthy.
    pub fn is_healthy(&self) -> bool {
        self.stats.read().state == InstanceState::Connected && self.client.read().is_some()
    }

    /// Update state to connecting.
    pub fn mark_connecting(&self) {
        let mut stats = self.stats.write();
        stats.state = InstanceState::Connecting;
        stats.last_attempt = Some(Instant::now());
    }

    /// Update state to connected with discovered features.
    pub fn mark_connected(&self, features: DiscoveredFeatures, connection: McpClientConnection) {
        let mut stats = self.stats.write();
        stats.state = InstanceState::Connected;
        stats.connected_at = Some(Instant::now());
        stats.consecutive_failures = 0;
        stats.last_error = None;

        *self.features.write() = Some(features);
        *self.client.write() = Some(connection);
    }

    /// Update state to failed.
    pub fn mark_failed(&self, error: String) {
        let mut stats = self.stats.write();
        stats.state = InstanceState::Failed;
        stats.consecutive_failures += 1;
        stats.last_error = Some(error);
    }

    /// Update state to OAuth pending.
    pub fn mark_oauth_pending(&self) {
        let mut stats = self.stats.write();
        stats.state = InstanceState::OAuthPending;
    }

    /// Record a successful request.
    pub fn record_success(&self) {
        self.stats.write().requests_served += 1;
    }

    /// Record a failed request.
    pub fn record_failure(&self, error: &str) {
        let mut stats = self.stats.write();
        stats.consecutive_failures += 1;
        stats.last_error = Some(error.to_string());
    }

    /// Get discovered features.
    pub fn get_features(&self) -> Option<DiscoveredFeatures> {
        self.features.read().clone()
    }

    /// Execute an operation with the MCP client.
    ///
    /// This is the primary API for accessing the MCP client. The closure receives
    /// a reference to the `McpClient` and can perform operations like calling tools
    /// or reading resources.
    ///
    /// Returns `None` if the client is not connected.
    ///
    /// # Example
    /// ```ignore
    /// let result = instance.with_client(|client| {
    ///     client.peer().clone()  // Get client handle for async operations
    /// });
    /// ```
    pub fn with_client<F, R>(&self, f: F) -> Option<R>
    where
        F: FnOnce(&McpClient) -> R,
    {
        let guard = self.client.read();
        guard.as_ref().and_then(|conn| conn.client().map(f))
    }

    /// Get the server URL from the instance key (for HTTP/SSE transports).
    /// Returns None for STDIO transports.
    pub fn get_url(&self) -> Option<String> {
        // The description format is "transport:url" (e.g., "sse:https://mcp.atlassian.com/v1/sse")
        let desc = &self.key.description;
        if desc.starts_with("http:") {
            Some(desc.strip_prefix("http:").unwrap_or(desc).to_string())
        } else if desc.starts_with("sse:") {
            Some(desc.strip_prefix("sse:").unwrap_or(desc).to_string())
        } else {
            None
        }
    }
}