//! MCP Client Manager - manages connections to MCP servers
//!
//! This module provides a decoupled MCP integration that:
//! - Doesn't block gateway startup
//! - Is completely optional (disabled by default)
//! - Connects to MCP servers asynchronously
//! - Dynamically registers MCP tools via the Tool trait adapter

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

use rmcp::{
    model::{CallToolRequestParams, CallToolResult, ServerInfo, Tool},
    RoleClient, ServiceExt,
    service::RunningService,
    transport::TokioChildProcess,
    transport::streamable_http_client::{StreamableHttpClientTransport, StreamableHttpClientTransportConfig},
};
use http::{HeaderName, HeaderValue};
use tokio::process::Command;

use crate::mcp::config::{McpServerConfig, McpTransportConfig};
use std::env;

/// Resolve ${ENV_VAR} placeholders in a value string
fn resolve_env_placeholders_in_value(value: &str) -> String {
    let re = regex::Regex::new(r"\$\{([A-Z_][A-Z0-9_]*)\}").expect("invalid regex");
    re.replace_all(value, |caps: &regex::Captures| {
        let var_name = &caps[1];
        env::var(var_name).unwrap_or_else(|_| caps[0].to_string())
    })
    .to_string()
}

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

/// Information about a connected MCP server
#[derive(Debug, Clone)]
pub struct McpServerInfo {
    /// Server name (effective name from config)
    pub name: String,
    /// Server key (the key in mcpServers map)
    pub key: String,
    /// Server information from MCP protocol
    pub info: Option<ServerInfo>,
    /// Available tools
    pub tools: Vec<Tool>,
}

/// Manager for MCP client connections
///
/// This manager handles:
/// - Connecting to MCP servers (stdio and HTTP transports)
/// - Discovering available tools
/// - Calling tools on connected servers
/// - Connection lifecycle management
pub struct McpClientManager {
    /// Connected clients keyed by server key
    clients: RwLock<HashMap<String, Arc<McpClient>>>,
    /// Server information cache keyed by server key
    server_info: RwLock<HashMap<String, McpServerInfo>>,
    /// Count of active stdio (child process) connections
    stdio_client_count: std::sync::atomic::AtomicUsize,
    /// Connection errors per server key (last error message)
    connection_errors: RwLock<HashMap<String, String>>,
}

impl McpClientManager {
    /// Create a new manager (no connections yet)
    pub fn new() -> Self {
        Self {
            clients: RwLock::new(HashMap::new()),
            server_info: RwLock::new(HashMap::new()),
            stdio_client_count: std::sync::atomic::AtomicUsize::new(0),
            connection_errors: RwLock::new(HashMap::new()),
        }
    }

    /// Connect to all configured servers (async, non-blocking)
    ///
    /// This method is designed to be called asynchronously without
    /// blocking the main gateway startup flow.
    /// Takes a list of (key, config) pairs from the mcpServers map.
    ///
    /// Each server connection is retried up to 3 times with exponential backoff
    /// (1s, 2s, 4s) to handle transient failures during gateway restart.
    pub async fn connect_all(&self, servers: Vec<(String, McpServerConfig)>) -> anyhow::Result<()> {
        let total = servers.iter().filter(|(_, c)| c.is_active).count();
        if total == 0 {
            tracing::info!("No active MCP servers to connect");
            return Ok(());
        }

        tracing::info!(server_count = total, "Connecting to MCP servers");

        let mut connected = 0usize;
        let mut failed = 0usize;
        const MAX_RETRIES: u32 = 3;

        for (key, config) in servers {
            if !config.is_active {
                tracing::info!(key = %key, "Skipping inactive MCP server");
                continue;
            }

            let mut last_error = None;
            for attempt in 1..=MAX_RETRIES {
                match self.connect_server(&key, &config).await {
                    Ok(info) => {
                        tracing::info!(
                            key = %key,
                            name = %info.name,
                            tools_count = info.tools.len(),
                            attempt,
                            "Connected to MCP server"
                        );
                        connected += 1;
                        last_error = None;
                        break;
                    }
                    Err(e) => {
                        last_error = Some(e);
                        if attempt < MAX_RETRIES {
                            let delay_secs = 1u64 << (attempt - 1); // 1s, 2s, 4s
                            tracing::warn!(
                                key = %key,
                                attempt,
                                max_retries = MAX_RETRIES,
                                retry_delay_secs = delay_secs,
                                error = %last_error.as_ref().unwrap(),
                                "MCP connection failed, retrying"
                            );
                            tokio::time::sleep(std::time::Duration::from_secs(delay_secs)).await;
                        }
                    }
                }
            }

            if let Some(e) = last_error {
                tracing::error!(
                    key = %key,
                    error = %e,
                    attempts = MAX_RETRIES,
                    "Failed to connect to MCP server after all retries"
                );
                // Record error for status reporting
                self.connection_errors.write().await.insert(key.clone(), e.to_string());
                failed += 1;
            } else {
                // Clear any previous error on successful connection
                self.connection_errors.write().await.remove(&key);
            }
        }

        tracing::info!(
            total,
            connected,
            failed,
            "MCP connection summary: {}/{} connected, {} failed",
            connected,
            total,
            failed
        );

        // Only return error if all connections failed
        if connected == 0 && failed > 0 {
            return Err(anyhow::anyhow!(
                "All {} MCP server connection(s) failed",
                failed
            ));
        }

        Ok(())
    }

    /// Connect to a single MCP server
    pub async fn connect_server(&self, key: &str, config: &McpServerConfig) -> anyhow::Result<McpServerInfo> {
        let effective_name = config.effective_name(key);
        tracing::info!(key = %key, name = %effective_name, transport_type = %config.transport_type, "Connecting to MCP server");

        let transport = config.transport().map_err(|e| anyhow::anyhow!("{}", e))?;
        let client = match transport {
            McpTransportConfig::Stdio { command, args, env } => {
                self.connect_stdio(&command, &args, &env).await?
            }
            McpTransportConfig::Http { url, headers } => {
                self.connect_http(&url, &headers).await?
            }
        };

        // Get server info (returns Option<&ServerInfo>)
        let info = client.peer_info().cloned();

        // List available tools
        let tools = client.list_all_tools().await?;

        let server_info = McpServerInfo {
            key: key.to_string(),
            name: effective_name,
            info,
            tools,
        };

        // Store the client and info
        {
            let mut clients = self.clients.write().await;
            clients.insert(key.to_string(), Arc::new(client));
        }
        {
            let mut info_map = self.server_info.write().await;
            info_map.insert(key.to_string(), server_info.clone());
        }

        Ok(server_info)
    }

    /// Connect via stdio transport (spawn child process)
    async fn connect_stdio(
        &self,
        command: &str,
        args: &[String],
        env: &HashMap<String, String>,
    ) -> anyhow::Result<McpClient> {
        let mut cmd = Command::new(command);
        cmd.args(args);

        // Set environment variables
        for (key, value) in env {
            cmd.env(key, value);
        }

        let transport = TokioChildProcess::new(cmd)?;

        // Use default client handler (empty tuple)
        let client = ().serve(transport).await?;

        // Track that we have a stdio (child process) connection
        self.stdio_client_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);

        Ok(client)
    }

    /// Connect via HTTP transport (Streamable HTTP)
    async fn connect_http(
        &self,
        url: &str,
        headers: &HashMap<String, String>,
    ) -> anyhow::Result<McpClient> {
        // Resolve env placeholders in headers
        let resolved_headers: HashMap<String, String> = headers
            .iter()
            .map(|(key, value)| {
                // Resolve ${ENV_VAR} placeholders
                let resolved = if value.contains("${") {
                    resolve_env_placeholders_in_value(value)
                } else {
                    value.clone()
                };
                (key.clone(), resolved)
            })
            .collect();

        // Build custom headers
        let custom_headers: HashMap<HeaderName, HeaderValue> = resolved_headers
            .iter()
            .filter_map(|(key, value)| {
                // Try to parse header name and value
                HeaderName::try_from(key.clone())
                    .ok()
                    .and_then(|name| {
                        HeaderValue::try_from(value.clone())
                            .ok()
                            .map(|val| (name, val))
                    })
            })
            .collect();

        // Create transport config with custom headers
        let config = StreamableHttpClientTransportConfig::with_uri(url)
            .custom_headers(custom_headers);

        // Create transport using reqwest client (default)
        let transport = StreamableHttpClientTransport::with_client(
            reqwest::Client::default(),
            config,
        );

        // Connect
        let client = ().serve(transport).await?;

        Ok(client)
    }

    /// Get a client by server key
    pub async fn get_client(&self, key: &str) -> Option<Arc<McpClient>> {
        let clients = self.clients.read().await;
        clients.get(key).cloned()
    }

    /// Get server info by key
    pub async fn get_server_info(&self, key: &str) -> Option<McpServerInfo> {
        let info_map = self.server_info.read().await;
        info_map.get(key).cloned()
    }

    /// Get all connected server keys
    pub async fn connected_servers(&self) -> Vec<String> {
        let clients = self.clients.read().await;
        clients.keys().cloned().collect()
    }

    /// Get all tools from all connected servers
    /// Returns (server_key, tool) pairs for tool registration
    pub async fn all_tools(&self) -> Vec<(String, Tool)> {
        let info_map = self.server_info.read().await;
        info_map
            .values()
            .flat_map(|info| {
                info.tools.iter().map(|tool| (info.key.clone(), tool.clone()))
            })
            .collect()
    }

    /// Call a tool on a specific server by key
    pub async fn call_tool(
        &self,
        server_key: impl Into<String>,
        tool_name: impl Into<String>,
        args: serde_json::Value,
    ) -> anyhow::Result<CallToolResult> {
        let server_key = server_key.into();
        let tool_name = tool_name.into();

        let client = self
            .get_client(&server_key)
            .await
            .ok_or_else(|| anyhow::anyhow!("MCP server '{}' not connected", server_key))?;

        // Convert Value to JsonObject if it's an object
        let arguments = if args.is_object() {
            args.as_object().unwrap().clone()
        } else {
            // If not an object, use empty object
            serde_json::Map::new()
        };

        // Create params with owned String (converted to Cow<'static, str>)
        let params = CallToolRequestParams::new(tool_name).with_arguments(arguments);

        let result = client.call_tool(params).await?;

        Ok(result)
    }

    /// Disconnect from a server by key
    pub async fn disconnect(&self, key: impl Into<String>) -> anyhow::Result<()> {
        let key = key.into();
        let mut clients = self.clients.write().await;
        if clients.remove(&key).is_some() {
            tracing::info!(key = %key, "Disconnected MCP server");
        }
        self.server_info.write().await.remove(&key);
        Ok(())
    }

    /// Disconnect from all servers
    pub async fn disconnect_all(&self) -> anyhow::Result<()> {
        let mut clients = self.clients.write().await;
        for (key, _client) in clients.drain() {
            tracing::info!(key = %key, "Disconnected MCP server");
        }
        self.server_info.write().await.clear();
        Ok(())
    }

    /// Shut down all MCP connections with proper cleanup
    ///
    /// Drops all client connections and waits for child processes to terminate
    /// if stdio transports were in use. This prevents race conditions during
    /// gateway restart where old MCP processes may still be running when
    /// new ones start.
    pub async fn shutdown_all(&self) -> anyhow::Result<()> {
        let stdio_count = self.stdio_client_count.load(std::sync::atomic::Ordering::SeqCst);

        // Drop all clients (triggers cancellation + graceful shutdown in rmcp)
        self.disconnect_all().await?;

        // If stdio connections were active, wait for child processes to be killed.
        // rmcp's RunningService::drop() triggers async cancellation with:
        // - 2 second graceful drain period
        // - 3 second process kill timeout
        // Total: ~5 seconds. We add 1 second buffer.
        if stdio_count > 0 {
            tracing::info!(
                stdio_count,
                "Waiting for MCP child processes to terminate (up to 6s)..."
            );
            tokio::time::sleep(std::time::Duration::from_secs(6)).await;
            tracing::info!("MCP child process cleanup wait complete");
            self.stdio_client_count.store(0, std::sync::atomic::Ordering::SeqCst);
        }

        Ok(())
    }

    /// Check if any servers are connected
    pub async fn has_connections(&self) -> bool {
        !self.clients.read().await.is_empty()
    }

    /// Get current MCP connection status for all configured servers
    pub async fn get_status(
        &self,
        configured_servers: &HashMap<String, crate::mcp::McpServerConfig>,
    ) -> McpStatusResponse {
        let info_map = self.server_info.read().await;
        let errors = self.connection_errors.read().await;
        let clients = self.clients.read().await;

        let mut total_servers = 0usize;
        let mut connected_servers = 0usize;
        let mut failed_servers = 0usize;
        let mut total_tools = 0usize;
        let mut servers = Vec::new();

        for (key, config) in configured_servers {
            total_servers += 1;
            let name = config.effective_name(key);
            let transport_type = config.transport_type.clone();
            let is_active = config.is_active;
            let connected = clients.contains_key(key);
            let tool_count = info_map.get(key).map(|info| info.tools.len()).unwrap_or(0);
            let error = errors.get(key).cloned();

            if connected {
                connected_servers += 1;
            } else if is_active && error.is_some() {
                failed_servers += 1;
            }

            total_tools += tool_count;

            servers.push(McpServerStatus {
                key: key.clone(),
                name,
                transport_type,
                is_active,
                connected,
                tool_count,
                error,
            });
        }

        McpStatusResponse {
            enabled: !configured_servers.is_empty(),
            total_servers,
            connected_servers,
            failed_servers,
            total_tools,
            servers,
        }
    }
}

/// Status of a single MCP server connection
#[derive(Debug, Clone, serde::Serialize)]
pub struct McpServerStatus {
    pub key: String,
    pub name: String,
    pub transport_type: String,
    pub is_active: bool,
    pub connected: bool,
    pub tool_count: usize,
    pub error: Option<String>,
}

/// Overall MCP status response
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct McpStatusResponse {
    pub enabled: bool,
    pub total_servers: usize,
    pub connected_servers: usize,
    pub failed_servers: usize,
    pub total_tools: usize,
    pub servers: Vec<McpServerStatus>,
}

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

/// MCP Initializer - handles asynchronous MCP initialization
///
/// This struct provides a decoupled way to initialize MCP:
/// - Doesn't block gateway startup
/// - Can be initialized in a background task
/// - Tools are registered after connection is established
pub struct McpInitializer {
    /// The MCP client manager (None if MCP is disabled)
    manager: Option<Arc<McpClientManager>>,
    /// Connection task handle (for background initialization)
    connection_task: Option<tokio::task::JoinHandle<anyhow::Result<()>>>,
}

impl McpInitializer {
    /// Create a disabled initializer (MCP not configured)
    pub fn disabled() -> Self {
        Self {
            manager: None,
            connection_task: None,
        }
    }

    /// Create an initializer with MCP configuration
    ///
    /// This spawns a background task to connect to MCP servers,
    /// allowing the gateway to start immediately.
    pub fn with_config(config: crate::mcp::McpConfig) -> Self {
        let server_count = config.mcp_servers.len();
        let active_count = config.active_servers().len();

        if !config.has_active_servers() {
            if server_count > 0 {
                tracing::info!(
                    server_count,
                    active_count,
                    "MCP disabled: {} server(s) configured but none active",
                    server_count
                );
            } else {
                tracing::info!("MCP disabled: no mcpServers configured in config");
            }
            return Self::disabled();
        }

        tracing::info!(
            server_count,
            active_count,
            "MCP enabled: {} active server(s) out of {} configured",
            active_count,
            server_count
        );

        let manager = Arc::new(McpClientManager::new());
        let servers = config.active_servers();

        // Spawn background connection task
        let manager_clone = manager.clone();
        let connection_task = tokio::spawn(async move {
            tracing::info!("Starting MCP connection task...");
            manager_clone.connect_all(servers).await
        });

        Self {
            manager: Some(manager),
            connection_task: Some(connection_task),
        }
    }

    /// Get the manager (if MCP is enabled)
    pub fn manager(&self) -> Option<Arc<McpClientManager>> {
        self.manager.clone()
    }

    /// Check if MCP is enabled
    pub fn is_enabled(&self) -> bool {
        self.manager.is_some()
    }

    /// Wait for connections to complete (optional)
    ///
    /// This can be called if you want to ensure MCP servers are connected
    /// before proceeding, but it's not required.
    pub async fn wait_for_connections(&mut self) -> anyhow::Result<()> {
        if let Some(task) = self.connection_task.take() {
            // Handle JoinError and inner Result
            task.await??;
        }
        Ok(())
    }

    /// Register MCP tools to the tool registry
    ///
    /// This should be called after the gateway is ready to accept tools.
    /// Waits for connections to complete before registering tools.
    pub async fn register_tools(&mut self, registry: &mut crate::tools::ToolRegistry) -> anyhow::Result<()> {
        if let Some(manager) = self.manager.clone() {
            // Wait for connections to complete first
            self.wait_for_connections().await?;

            tracing::info!("Registering MCP tools after connections completed");
            crate::mcp::register_mcp_tools(manager, registry).await?;
        }
        Ok(())
    }
}