use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;

use async_trait::async_trait;
use motosan_agent_tool::ToolDef;
use rmcp::model::{CallToolRequestParams, CallToolResult, RawContent};
use rmcp::service::{Peer, RoleClient, RunningService};
use rmcp::transport::child_process::TokioChildProcess;
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
use rmcp::transport::StreamableHttpClientTransport;
use rmcp::ServiceExt;

use crate::error::Result;
use crate::AgentError;

/// Trait implemented by all MCP server backends.
#[async_trait]
pub trait McpServer: Send + Sync {
    /// The logical name of this MCP server (used for tool namespacing).
    fn name(&self) -> &str;

    /// Connect and initialize the MCP server.
    async fn connect(&self) -> Result<()>;

    /// List all tools available on this server.
    async fn list_tools(&self) -> Result<Vec<ToolDef>>;

    /// Execute a tool by name with the given arguments.
    async fn call_tool(&self, name: &str, args: serde_json::Value) -> Result<String>;

    /// Disconnect cleanly.
    async fn disconnect(&self) -> Result<()>;
}

type StdioService = RunningService<RoleClient, ()>;
type HttpService = RunningService<RoleClient, ()>;

/// MCP server over stdio transport (spawns a child process).
pub struct McpServerStdio {
    pub name: String,
    pub command: String,
    pub args: Vec<String>,
    pub env: HashMap<String, String>,
    pub cwd: Option<PathBuf>,
    session: futures::lock::Mutex<Option<StdioService>>,
}

impl McpServerStdio {
    pub fn new(
        command: impl Into<String>,
        args: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        let command = command.into();
        let name = command.clone();
        Self {
            name,
            command,
            args: args.into_iter().map(Into::into).collect(),
            env: HashMap::new(),
            cwd: None,
            session: futures::lock::Mutex::new(None),
        }
    }

    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = name.into();
        self
    }
}

#[async_trait]
impl McpServer for McpServerStdio {
    fn name(&self) -> &str {
        &self.name
    }

    async fn connect(&self) -> Result<()> {
        let mut cmd = tokio::process::Command::new(&self.command);
        cmd.args(&self.args);
        for (k, v) in &self.env {
            cmd.env(k, v);
        }
        if let Some(cwd) = &self.cwd {
            cmd.current_dir(cwd);
        }

        let transport = TokioChildProcess::new(cmd).map_err(|e| AgentError::Mcp(e.to_string()))?;

        let service: StdioService = ().serve(transport).await.map_err(map_init_error)?;

        let mut session = self.session.lock().await;
        *session = Some(service);
        Ok(())
    }

    async fn list_tools(&self) -> Result<Vec<ToolDef>> {
        let session = self.session.lock().await;
        let peer = peer_from_session(&session)?;
        let tools = peer.list_all_tools().await.map_err(map_service_error)?;
        Ok(tools.into_iter().map(rmcp_tool_to_tool_def).collect())
    }

    async fn call_tool(&self, name: &str, args: serde_json::Value) -> Result<String> {
        let session = self.session.lock().await;
        let peer = peer_from_session(&session)?;
        let params = build_call_params(name, args);
        let result = peer.call_tool(params).await.map_err(map_service_error)?;
        Ok(call_tool_result_to_string(&result))
    }

    async fn disconnect(&self) -> Result<()> {
        let mut session = self.session.lock().await;
        if let Some(s) = session.take() {
            let _ = s.cancel().await;
        }
        Ok(())
    }
}

/// MCP server over HTTP (Streamable HTTP or SSE).
pub struct McpServerHttp {
    pub name: String,
    pub url: String,
    pub auth_token: Option<String>,
    pub headers: HashMap<String, String>,
    session: futures::lock::Mutex<Option<HttpService>>,
}

impl McpServerHttp {
    pub fn new(url: impl Into<String>) -> Self {
        let url = url.into();
        let name = url.clone();
        Self {
            name,
            url,
            auth_token: None,
            headers: HashMap::new(),
            session: futures::lock::Mutex::new(None),
        }
    }

    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = name.into();
        self
    }
}

#[async_trait]
impl McpServer for McpServerHttp {
    fn name(&self) -> &str {
        &self.name
    }

    async fn connect(&self) -> Result<()> {
        let mut config = StreamableHttpClientTransportConfig::with_uri(self.url.as_str());
        if let Some(token) = &self.auth_token {
            config.auth_header = Some(format!("Bearer {token}"));
        }
        // Custom headers are applied through the config's custom_headers field.
        // Header names/values are validated by the http crate during parsing.
        for (k, v) in &self.headers {
            let header_name = k
                .parse()
                .map_err(|e: http::header::InvalidHeaderName| AgentError::Mcp(e.to_string()))?;
            let header_value = v
                .parse()
                .map_err(|e: http::header::InvalidHeaderValue| AgentError::Mcp(e.to_string()))?;
            config.custom_headers.insert(header_name, header_value);
        }

        let transport = StreamableHttpClientTransport::<reqwest::Client>::from_config(config);
        let service: HttpService = ().serve(transport).await.map_err(map_init_error)?;

        let mut session = self.session.lock().await;
        *session = Some(service);
        Ok(())
    }

    async fn list_tools(&self) -> Result<Vec<ToolDef>> {
        let session = self.session.lock().await;
        let peer = peer_from_session(&session)?;
        let tools = peer.list_all_tools().await.map_err(map_service_error)?;
        Ok(tools.into_iter().map(rmcp_tool_to_tool_def).collect())
    }

    async fn call_tool(&self, name: &str, args: serde_json::Value) -> Result<String> {
        let session = self.session.lock().await;
        let peer = peer_from_session(&session)?;
        let params = build_call_params(name, args);
        let result = peer.call_tool(params).await.map_err(map_service_error)?;
        Ok(call_tool_result_to_string(&result))
    }

    async fn disconnect(&self) -> Result<()> {
        let mut session = self.session.lock().await;
        if let Some(s) = session.take() {
            let _ = s.cancel().await;
        }
        Ok(())
    }
}

// ── helpers ──────────────────────────────────────────────────────────────

fn peer_from_session(
    session: &Option<impl std::ops::Deref<Target = Peer<RoleClient>>>,
) -> Result<Peer<RoleClient>> {
    session
        .as_ref()
        .map(|s| std::ops::Deref::deref(s).clone())
        .ok_or_else(|| AgentError::Mcp("not connected".into()))
}

fn rmcp_tool_to_tool_def(tool: rmcp::model::Tool) -> ToolDef {
    ToolDef {
        name: tool.name.into_owned(),
        description: tool.description.map(|d| d.into_owned()).unwrap_or_default(),
        input_schema: serde_json::Value::Object(tool.input_schema.as_ref().clone()),
    }
}

fn build_call_params(name: &str, args: serde_json::Value) -> CallToolRequestParams {
    let arguments = match args {
        serde_json::Value::Object(map) => Some(map),
        _ => None,
    };
    let mut params = CallToolRequestParams::new(name.to_string());
    params.arguments = arguments;
    params
}

fn call_tool_result_to_string(result: &CallToolResult) -> String {
    result
        .content
        .iter()
        .filter_map(|c| match &c.raw {
            RawContent::Text(t) => Some(t.text.clone()),
            _ => None,
        })
        .collect::<Vec<_>>()
        .join("\n")
}

fn map_service_error(e: rmcp::service::ServiceError) -> AgentError {
    AgentError::Mcp(e.to_string())
}

fn map_init_error(e: rmcp::service::ClientInitializeError) -> AgentError {
    AgentError::Mcp(e.to_string())
}

// Compile-time check: McpServer must be object-safe and Arc-compatible.
#[allow(dead_code)]
fn assert_arc_compatible(_server: Arc<dyn McpServer>) {}

#[cfg(test)]
mod tests {
    use super::*;
    use rmcp::model::Content;

    #[test]
    fn stdio_new_sets_fields() {
        let server = McpServerStdio::new("npx", ["-y", "@modelcontextprotocol/server-linear"]);
        assert_eq!(server.command, "npx");
        assert_eq!(
            server.args,
            vec!["-y", "@modelcontextprotocol/server-linear"]
        );
        assert!(server.env.is_empty());
        assert!(server.cwd.is_none());
    }

    #[test]
    fn http_new_sets_fields() {
        let server = McpServerHttp::new("http://localhost:3000/mcp");
        assert_eq!(server.url, "http://localhost:3000/mcp");
        assert!(server.auth_token.is_none());
        assert!(server.headers.is_empty());
    }

    #[test]
    fn rmcp_tool_conversion() {
        let tool = rmcp::model::Tool::new(
            "search",
            "Search the web",
            Arc::new(serde_json::Map::from_iter([(
                "type".to_string(),
                serde_json::Value::String("object".to_string()),
            )])),
        );
        let def = rmcp_tool_to_tool_def(tool);
        assert_eq!(def.name, "search");
        assert_eq!(def.description, "Search the web");
        assert_eq!(def.input_schema["type"], "object");
    }

    #[test]
    fn build_call_params_with_object() {
        let args = serde_json::json!({"query": "rust"});
        let params = build_call_params("search", args);
        assert_eq!(params.name.as_ref(), "search");
        assert!(params.arguments.is_some());
        assert_eq!(
            params.arguments.unwrap()["query"],
            serde_json::Value::String("rust".into())
        );
    }

    #[test]
    fn build_call_params_with_non_object() {
        let params = build_call_params("ping", serde_json::Value::Null);
        assert_eq!(params.name.as_ref(), "ping");
        assert!(params.arguments.is_none());
    }

    #[test]
    fn call_tool_result_to_string_text() {
        let result = CallToolResult::success(vec![Content::text("hello world")]);
        assert_eq!(call_tool_result_to_string(&result), "hello world");
    }

    #[test]
    fn call_tool_result_to_string_multiple() {
        let result = CallToolResult::success(vec![Content::text("line1"), Content::text("line2")]);
        assert_eq!(call_tool_result_to_string(&result), "line1\nline2");
    }

    #[test]
    fn call_tool_result_to_string_empty() {
        let result = CallToolResult::success(vec![]);
        assert_eq!(call_tool_result_to_string(&result), "");
    }

    #[tokio::test]
    async fn stdio_not_connected_errors() {
        let server = McpServerStdio::new("echo", ["hello"]);
        let err = server.list_tools().await.unwrap_err();
        assert!(err.to_string().contains("not connected"));
    }

    #[tokio::test]
    async fn http_not_connected_errors() {
        let server = McpServerHttp::new("http://localhost:9999");
        let err = server
            .call_tool("test", serde_json::json!({}))
            .await
            .unwrap_err();
        assert!(err.to_string().contains("not connected"));
    }

    #[tokio::test]
    async fn disconnect_when_not_connected_is_ok() {
        let server = McpServerStdio::new("echo", ["hello"]);
        assert!(server.disconnect().await.is_ok());
    }
}
