use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use anyhow::{Result, Context};
use rmcp::model::{CallToolRequestParam, RawContent};
use rmcp::ServiceExt;
use rmcp::transport::child_process::TokioChildProcess;
use serde_json::Value;

use super::types::{ExternalSystem, McpServerConfig};

/// Manages MCP client connections to external system servers.
///
/// Maintains a registry of `McpServerConfig` entries keyed by system name.
/// Each operation spawns a fresh child process, performs the MCP handshake,
/// executes the request, and shuts down. This avoids relying on long-lived
/// connections to external MCP servers that may not support them well.
pub struct McpConnectionManager {
    configs: Arc<Mutex<HashMap<String, McpServerConfig>>>,
}

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

impl McpConnectionManager {
    pub fn new() -> Self {
        Self {
            configs: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Register an external MCP server configuration.
    pub async fn register(&self, config: McpServerConfig) {
        let mut configs = self.configs.lock().await;
        configs.insert(config.system.to_string(), config);
    }

    /// Check if a system has a registered MCP server.
    pub async fn is_registered(&self, system: &ExternalSystem) -> bool {
        let configs = self.configs.lock().await;
        configs.contains_key(&system.to_string())
    }

    /// Spawn an external MCP server, call a tool, and return the result.
    ///
    /// Creates a fresh connection per call (spawn -> init -> call -> shutdown).
    pub async fn call_tool(
        &self,
        system: &ExternalSystem,
        tool_name: &str,
        arguments: Value,
    ) -> Result<Value> {
        let config = {
            let configs = self.configs.lock().await;
            configs
                .get(&system.to_string())
                .cloned()
                .context(format!("No MCP server configured for {}", system))?
        };

        // Build the command
        let mut cmd = tokio::process::Command::new(&config.command);
        for arg in &config.args {
            cmd.arg(arg);
        }
        for (key, val) in &config.env_vars {
            cmd.env(key, val);
        }

        // Spawn and connect via MCP client protocol
        let transport = TokioChildProcess::new(&mut cmd)
            .context(format!("Failed to spawn MCP server for {}", system))?;
        let client = ().serve(transport)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to connect to {} MCP server: {}", system, e))?;

        // Call the tool
        let result = client
            .call_tool(CallToolRequestParam {
                name: tool_name.to_string().into(),
                arguments: Some(arguments.as_object().cloned().unwrap_or_default()),
            })
            .await
            .map_err(|e| {
                anyhow::anyhow!("Tool call '{}' failed on {}: {}", tool_name, system, e)
            })?;

        // Extract text content from the result.
        // Content = Annotated<RawContent>, which Derefs to RawContent.
        // RawContent::Text(RawTextContent { text }) holds the string.
        let content_text = result
            .content
            .iter()
            .filter_map(|c| {
                if let RawContent::Text(t) = &c.raw {
                    Some(t.text.clone())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
            .join("\n");

        // Try to parse as JSON, otherwise wrap as string
        let value = serde_json::from_str(&content_text)
            .unwrap_or(Value::String(content_text));

        // Shutdown gracefully
        let _ = client.cancel().await;

        Ok(value)
    }

    /// List available tools on an external MCP server.
    pub async fn list_tools(&self, system: &ExternalSystem) -> Result<Vec<String>> {
        let config = {
            let configs = self.configs.lock().await;
            configs
                .get(&system.to_string())
                .cloned()
                .context(format!("No MCP server configured for {}", system))?
        };

        let mut cmd = tokio::process::Command::new(&config.command);
        for arg in &config.args {
            cmd.arg(arg);
        }
        for (key, val) in &config.env_vars {
            cmd.env(key, val);
        }

        let transport = TokioChildProcess::new(&mut cmd)
            .context(format!("Failed to spawn MCP server for {}", system))?;
        let client = ().serve(transport)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to connect to {} MCP server: {}", system, e))?;

        let tools = client
            .list_all_tools()
            .await
            .map_err(|e| anyhow::anyhow!("Failed to list tools on {}: {}", system, e))?;

        let tool_names: Vec<String> = tools.iter().map(|t| t.name.to_string()).collect();

        let _ = client.cancel().await;

        Ok(tool_names)
    }
}
