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

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

    /// 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.
    pub async fn connect_all(&self, servers: Vec<(String, McpServerConfig)>) -> anyhow::Result<()> {
        for (key, config) in servers {
            if !config.is_active {
                tracing::info!(key = %key, "Skipping inactive MCP server");
                continue;
            }

            // Each server connection is independent
            match self.connect_server(&key, &config).await {
                Ok(info) => {
                    tracing::info!(
                        key = %key,
                        name = %info.name,
                        tools_count = info.tools.len(),
                        "Connected to MCP server"
                    );
                }
                Err(e) => {
                    // Log error but continue with other servers
                    tracing::error!(
                        key = %key,
                        error = %e,
                        "Failed to connect to MCP server"
                    );
                }
            }
        }
        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?;

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

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

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 {
        if !config.has_active_servers() {
            return Self::disabled();
        }

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