use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context as _, Result, anyhow};
use rmcp::model::{CallToolRequestParams, CallToolResult};
use rmcp::service::{RoleClient, RunningService, serve_client};
use rmcp::transport::StreamableHttpClientTransport;
use rmcp::transport::child_process::TokioChildProcess;
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
use tokio::process::Command;
use tokio::sync::Mutex;
use tokio::time;
use tracing::{debug, error, info, warn};

use crate::client_handler::LoggingClientHandler;
use crate::config::{McpHttpProtocol, McpServerSettings, McpTransport};
use crate::http_transport::headers_from_pairs;
use crate::sse_client::SseClientTransport;

enum ClientState {
    /// Transport created but handshake not yet performed.
    Connecting(Option<PendingTransport>),
    /// Handshake complete, ready for tool calls.
    Ready(Arc<RunningService<RoleClient, LoggingClientHandler>>),
    /// Connection was explicitly closed.
    Closed,
}

enum PendingTransport {
    Stdio(TokioChildProcess),
    Http(StreamableHttpClientTransport<fabro_http::HttpClient>),
    Sse(SseClientTransport),
}

/// MCP client wrapping the rmcp SDK. Handles stdio and HTTP transports.
pub struct McpClient {
    server_name: String,
    state:       Mutex<ClientState>,
}

impl McpClient {
    /// Create a new MCP client from config. Does not connect yet — call
    /// `initialize()`.
    pub fn new(config: &McpServerSettings) -> Result<Self> {
        let transport = match &config.transport {
            McpTransport::Stdio { command, env } => {
                let (program, args) = command.split_first().ok_or_else(|| {
                    anyhow!("MCP server '{}': command must not be empty", config.name)
                })?;
                let mut cmd = Command::new(program);
                cmd.args(args)
                    .stdin(Stdio::piped())
                    .stdout(Stdio::piped())
                    .stderr(Stdio::piped())
                    .kill_on_drop(true);

                if config.clear_env {
                    cmd.env_clear();
                }
                if !env.is_empty() {
                    cmd.envs(env);
                }
                if let Some(current_dir) = config.current_dir.as_ref() {
                    cmd.current_dir(current_dir);
                }

                #[cfg(unix)]
                cmd.process_group(0);

                let transport = TokioChildProcess::new(cmd)
                    .with_context(|| format!("failed to spawn MCP server '{}'", config.name))?;

                PendingTransport::Stdio(transport)
            }
            McpTransport::Http {
                protocol,
                url,
                headers,
            } => {
                let headers = headers_from_pairs(headers)?;
                let mut builder = fabro_http::HttpClientBuilder::new();
                if !headers.is_empty() {
                    builder = builder.default_headers(headers);
                }

                let http_client = builder.build()?;
                match protocol {
                    McpHttpProtocol::StreamableHttp => {
                        let http_config =
                            StreamableHttpClientTransportConfig::with_uri(url.clone());
                        let transport =
                            StreamableHttpClientTransport::with_client(http_client, http_config);
                        PendingTransport::Http(transport)
                    }
                    McpHttpProtocol::Sse => {
                        PendingTransport::Sse(SseClientTransport::new(url, http_client)?)
                    }
                }
            }
            McpTransport::Sandbox { .. } => {
                return Err(anyhow!(
                    "MCP server '{}': Sandbox transport must be resolved to Http before connecting",
                    config.name
                ));
            }
        };

        let transport_type = match &transport {
            PendingTransport::Stdio(_) => "stdio",
            PendingTransport::Http(_) => "http",
            PendingTransport::Sse(_) => "sse",
        };
        debug!(server = %config.name, transport = transport_type, "Creating MCP client");

        Ok(Self {
            server_name: config.name.clone(),
            state:       Mutex::new(ClientState::Connecting(Some(transport))),
        })
    }

    /// Perform the initialization handshake with the MCP server.
    pub async fn initialize(&self, timeout: Duration) -> Result<()> {
        let handler = LoggingClientHandler;

        let service = {
            let mut guard = self.state.lock().await;
            let transport = match &mut *guard {
                ClientState::Connecting(t) => t
                    .take()
                    .ok_or_else(|| anyhow!("client already initializing"))?,
                ClientState::Ready(_) => return Err(anyhow!("client already initialized")),
                ClientState::Closed => return Err(anyhow!("MCP client is shut down")),
            };

            // Drop the lock before the blocking handshake
            drop(guard);

            debug!(server = %self.server_name, "Starting MCP server handshake");

            let handshake = async {
                match transport {
                    PendingTransport::Stdio(t) => serve_client(handler.clone(), t).await,
                    PendingTransport::Http(t) => serve_client(handler.clone(), t).await,
                    PendingTransport::Sse(t) => serve_client(handler.clone(), t).await,
                }
            };

            let service = time::timeout(timeout, handshake)
                .await
                .map_err(|_| {
                    error!(server = %self.server_name, timeout_secs = timeout.as_secs(), "MCP server handshake timed out");
                    anyhow!(
                        "timed out initializing MCP server '{}' after {:?}",
                        self.server_name,
                        timeout
                    )
                })?
                .map_err(|e| {
                    error!(server = %self.server_name, error = %e, "MCP server handshake failed");
                    anyhow!(
                        "failed to initialize MCP server '{}': {}",
                        self.server_name,
                        e
                    )
                })?;

            let peer_info = service.peer().peer_info();
            if let Some(info) = peer_info {
                info!(
                    server = %self.server_name,
                    server_name = %info.server_info.name,
                    server_version = %info.server_info.version,
                    "MCP server initialized"
                );
            }

            Arc::new(service)
        };

        let mut guard = self.state.lock().await;
        *guard = ClientState::Ready(service);
        Ok(())
    }

    /// List all tools exposed by this server.
    /// Returns `(name, description, input_schema)` tuples.
    pub async fn list_tools(&self) -> Result<Vec<(String, String, serde_json::Value)>> {
        let service = self.service().await?;
        let result = service.list_all_tools().await.map_err(|e| {
            anyhow!(
                "failed to list tools from MCP server '{}': {}",
                self.server_name,
                e
            )
        })?;

        let tools: Vec<_> = result
            .into_iter()
            .map(|tool| {
                let name = tool.name.to_string();
                let description = tool.description.as_deref().unwrap_or("").to_string();
                let input_schema = serde_json::to_value(&*tool.input_schema).unwrap_or_default();
                (name, description, input_schema)
            })
            .collect();

        debug!(server = %self.server_name, tool_count = tools.len(), "Listed MCP server tools");

        Ok(tools)
    }

    /// Call a tool on this server.
    pub async fn call_tool(
        &self,
        name: &str,
        arguments: serde_json::Value,
        timeout: Duration,
    ) -> Result<CallToolResult> {
        let service = self.service().await?;

        let args = match arguments {
            serde_json::Value::Object(map) => Some(map),
            serde_json::Value::Null => None,
            other => {
                return Err(anyhow!(
                    "MCP tool arguments must be a JSON object, got {other}"
                ));
            }
        };

        let mut params = CallToolRequestParams::new(name.to_string());
        if let Some(arguments) = args {
            params = params.with_arguments(arguments);
        }

        debug!(server = %self.server_name, tool = %name, "Calling MCP tool");

        let result = time::timeout(timeout, service.call_tool(params))
            .await
            .map_err(|_| {
                warn!(server = %self.server_name, tool = %name, timeout_secs = timeout.as_secs(), "MCP tool call timed out");
                anyhow!(
                    "timed out calling tool '{}' on MCP server '{}' after {:?}",
                    name,
                    self.server_name,
                    timeout
                )
            })?
            .map_err(|e| {
                anyhow!(
                    "failed to call tool '{}' on MCP server '{}': {}",
                    name,
                    self.server_name,
                    e
                )
            })?;

        Ok(result)
    }

    pub async fn shutdown(self) -> Result<()> {
        let service = {
            let mut guard = self.state.lock().await;
            match std::mem::replace(&mut *guard, ClientState::Closed) {
                ClientState::Connecting(_) | ClientState::Closed => None,
                ClientState::Ready(service) => Some(service),
            }
        };

        if let Some(service) = service {
            match Arc::try_unwrap(service) {
                Ok(mut service) => {
                    service
                        .close_with_timeout(Duration::from_secs(2))
                        .await
                        .context("failed to shut down MCP client")?;
                }
                Err(service) => {
                    service.cancellation_token().cancel();
                }
            }
        }

        Ok(())
    }

    async fn service(&self) -> Result<Arc<RunningService<RoleClient, LoggingClientHandler>>> {
        let guard = self.state.lock().await;
        match &*guard {
            ClientState::Ready(service) => Ok(Arc::clone(service)),
            ClientState::Connecting(_) => Err(anyhow!("MCP client not initialized")),
            ClientState::Closed => Err(anyhow!("MCP client is shut down")),
        }
    }
}
