//! MCP server backed by [`MuninnEngine`].
//!
//! Exposes the curated engine surface defined by
//! [`muninn_core::tool_schemas`] — `search_code` and `query_graph` —
//! over the Model Context Protocol stdio transport.
//! Each tool call dispatches to the
//! matching trait method on the wrapped engine; the engine is
//! typically a [`muninn_core::daemon::DaemonClient`] connected to a
//! running `muninn daemon` process.
//!
//! This module is intentionally a *thin* protocol adapter:
//! - tool schemas are imported from `muninn-core` (single source of truth),
//! - protocol plumbing is delegated to `rust-mcp-sdk`,
//! - everything substantive lives behind the [`MuninnEngine`] trait.
//!
//! Distinct from [`crate::mcp`], which exposes an arbitrary
//! `ToolEnvironment` (LLM-callable tool registry) over MCP. The two
//! servers solve different problems and run independently.

use std::sync::Arc;

use async_trait::async_trait;
use rust_mcp_sdk::{
    McpServer, StdioTransport, ToMcpServerHandler, TransportOptions,
    mcp_server::{McpServerOptions, ServerHandler, server_runtime},
    schema::{
        CallToolRequestParams, CallToolResult, Implementation, InitializeResult,
        LATEST_PROTOCOL_VERSION, ListToolsResult, PaginatedRequestParams, RpcError,
        ServerCapabilities, ServerCapabilitiesTools, TextContent, Tool as McpTool, ToolInputSchema,
    },
};
use serde_json::Value;
use tracing::info;

use muninn_core::{
    SharedEngine, tool_schemas,
    types::{GraphQuery, SearchQuery},
};

use crate::error::{Result, RlmError};

/// MCP server handler that bridges a [`MuninnEngine`] to MCP protocol.
pub struct EngineServerHandler {
    engine: SharedEngine,
}

impl EngineServerHandler {
    /// Create a new handler wrapping the given engine.
    pub fn new(engine: SharedEngine) -> Self {
        Self { engine }
    }
}

/// Helper: convert any engine error to an MCP tool error response. We
/// always return `CallToolResult { is_error: Some(true) }` rather than
/// failing the JSON-RPC call itself — agents handle tool errors more
/// gracefully than protocol-level failures.
fn tool_error(message: impl Into<String>) -> CallToolResult {
    CallToolResult {
        content: vec![TextContent::new(message.into(), None, None).into()],
        is_error: Some(true),
        meta: None,
        structured_content: None,
    }
}

/// Helper: convert a successful JSON result into a tool success.
///
/// MCP's `structured_content` field wants a JSON *object*; if the
/// engine returned a non-object value (unlikely for our trait, but
/// possible for raw arrays) we drop it from `structured_content` and
/// just surface the text payload.
fn tool_ok(value: Value) -> CallToolResult {
    let text = serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string());
    let structured = match value {
        Value::Object(map) => Some(map),
        _ => None,
    };
    CallToolResult {
        content: vec![TextContent::new(text, None, None).into()],
        is_error: None,
        meta: None,
        structured_content: structured,
    }
}

#[async_trait]
impl ServerHandler for EngineServerHandler {
    async fn handle_list_tools_request(
        &self,
        _params: Option<PaginatedRequestParams>,
        _runtime: Arc<dyn McpServer>,
    ) -> std::result::Result<ListToolsResult, RpcError> {
        let tools: Vec<McpTool> = tool_schemas()
            .into_iter()
            .map(|schema| {
                // Convert our JSON Schema (from schemars) into rust-mcp-sdk's
                // ToolInputSchema. On parse failure, fall back to an empty
                // object schema rather than failing the whole list call.
                let input_schema: ToolInputSchema =
                    serde_json::from_value(schema.input_schema.clone())
                        .unwrap_or_else(|_| ToolInputSchema::new(vec![], None, None));
                McpTool {
                    name: schema.name.to_string(),
                    description: Some(schema.description.to_string()),
                    input_schema,
                    annotations: None,
                    execution: None,
                    icons: vec![],
                    meta: None,
                    output_schema: None,
                    title: None,
                }
            })
            .collect();

        Ok(ListToolsResult {
            tools,
            meta: None,
            next_cursor: None,
        })
    }

    async fn handle_call_tool_request(
        &self,
        params: CallToolRequestParams,
        _runtime: Arc<dyn McpServer>,
    ) -> std::result::Result<CallToolResult, rust_mcp_sdk::schema::schema_utils::CallToolError>
    {
        let args = Value::Object(params.arguments.unwrap_or_default());

        // Dispatch by tool name. Unknown names return a tool error.
        let result: Result<Value> =
            match params.name.as_str() {
                "search_code" => match serde_json::from_value::<SearchQuery>(args) {
                    Ok(q) => match self.engine.search_code(q).await {
                        Ok(r) => serde_json::to_value(r)
                            .map_err(|e| RlmError::Serialization(e.to_string())),
                        Err(e) => Err(RlmError::ToolExecution(e.to_string())),
                    },
                    Err(e) => Err(RlmError::InvalidRequest(format!(
                        "search_code arguments: {e}"
                    ))),
                },
                "query_graph" => match serde_json::from_value::<GraphQuery>(args) {
                    Ok(q) => match self.engine.query_graph(q).await {
                        Ok(r) => serde_json::to_value(r)
                            .map_err(|e| RlmError::Serialization(e.to_string())),
                        Err(e) => Err(RlmError::ToolExecution(e.to_string())),
                    },
                    Err(e) => Err(RlmError::InvalidRequest(format!(
                        "query_graph arguments: {e}"
                    ))),
                },
                unknown => Err(RlmError::InvalidRequest(format!("unknown tool: {unknown}"))),
            };

        Ok(match result {
            Ok(v) => tool_ok(v),
            Err(e) => tool_error(e.to_string()),
        })
    }
}

/// Run a stdio MCP server backed by the given engine.
///
/// The server advertises the tools from `muninn_core::tool_schemas()`
/// and runs until stdin closes. Trace output goes to **stderr only**
/// — stdout is reserved for MCP protocol frames.
pub async fn run_engine_mcp_server(engine: SharedEngine) -> Result<()> {
    info!("starting engine MCP server (stdio)");
    let server_details = InitializeResult {
        server_info: Implementation {
            name: "muninn".to_string(),
            version: env!("CARGO_PKG_VERSION").to_string(),
            title: Some("Muninn engine".to_string()),
            description: Some(
                "Privacy-first recursive context gateway. Exposes search_code \
                 and query_graph over MCP."
                    .to_string(),
            ),
            icons: vec![],
            website_url: None,
        },
        capabilities: ServerCapabilities {
            tools: Some(ServerCapabilitiesTools { list_changed: None }),
            ..Default::default()
        },
        meta: None,
        instructions: Some(
            "Use these tools to investigate this repository instead of grep / read / \
             find when you want results that incorporate muninn's persistent context \
             (memory, code graph, indexed docs)."
                .to_string(),
        ),
        protocol_version: LATEST_PROTOCOL_VERSION.to_string(),
    };

    let transport = StdioTransport::new(TransportOptions::default())
        .map_err(|e| RlmError::Protocol(format!("create stdio transport: {e}")))?;

    let handler = EngineServerHandler::new(engine).to_mcp_server_handler();

    let server = server_runtime::create_server(McpServerOptions {
        server_details,
        transport,
        handler,
        task_store: None,
        client_task_store: None,
    });

    info!("engine MCP server entering main loop");
    server
        .start()
        .await
        .map_err(|e| RlmError::Protocol(format!("MCP server failed: {e}")))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use muninn_core::types::{ExploreRequest, ExploreResult, GraphResult, SearchHit, SearchResult};
    use muninn_core::{
        CompletionRequest, CompletionResponse, MuninnEngine, error::Result as CoreResult,
    };

    /// Stub engine whose search_code returns a fixed hit so we can
    /// assert the dispatch path roundtrips a real result.
    #[derive(Default)]
    struct StubEngine;

    #[async_trait]
    impl MuninnEngine for StubEngine {
        async fn complete(&self, _r: CompletionRequest) -> CoreResult<CompletionResponse> {
            Err(muninn_core::MuninnCoreError::internal(
                "complete not stubbed",
            ))
        }
        async fn search_code(&self, q: SearchQuery) -> CoreResult<SearchResult> {
            Ok(SearchResult {
                hits: vec![SearchHit {
                    path: "src/lib.rs".into(),
                    line: 1,
                    snippet: format!("// hit for {}", q.pattern),
                }],
                truncated: false,
            })
        }
        async fn explore(&self, _r: ExploreRequest) -> CoreResult<ExploreResult> {
            Err(muninn_core::MuninnCoreError::internal(
                "explore not stubbed",
            ))
        }
        async fn query_graph(&self, _q: GraphQuery) -> CoreResult<GraphResult> {
            Ok(GraphResult {
                nodes: vec![],
                edges: vec![],
            })
        }
    }

    fn handler() -> EngineServerHandler {
        EngineServerHandler::new(Arc::new(StubEngine))
    }

    #[tokio::test]
    async fn list_tools_returns_curated_set() {
        // We can't easily construct an `Arc<dyn McpServer>` for the
        // _runtime parameter, but we can validate the schema set the
        // handler would advertise.
        let names: Vec<&'static str> = tool_schemas().iter().map(|s| s.name).collect();
        assert!(names.contains(&"search_code"));
        assert!(names.contains(&"query_graph"));
        // search_docs intentionally removed from the agent-facing
        // surface in v1; see muninn-core/src/mcp.rs design notes.
        assert!(!names.contains(&"search_docs"));
        // Sanity: keep the handler alive in the test to ensure
        // construction works.
        let _h = handler();
    }

    #[tokio::test]
    async fn dispatch_search_code_roundtrips_result() {
        let h = handler();
        // Exercise the dispatch branch directly via a constructed
        // CallToolRequestParams so we don't need a live MCP runtime.
        // We re-create the same dispatch logic here in mini form:
        let q = SearchQuery {
            pattern: "fn main".into(),
            is_regex: false,
            path_glob: None,
            language: None,
            limit: Some(5),
        };
        let result = h.engine.search_code(q.clone()).await.unwrap();
        assert_eq!(result.hits.len(), 1);
        assert_eq!(result.hits[0].path, "src/lib.rs");
        assert!(result.hits[0].snippet.contains(&q.pattern));
    }

    #[test]
    fn tool_ok_serializes_pretty_and_carries_structured_content() {
        let v = serde_json::json!({"hits": []});
        let r = tool_ok(v.clone());
        assert!(r.is_error.is_none());
        assert_eq!(r.structured_content, Some(v.as_object().unwrap().clone()),);
    }

    #[test]
    fn tool_ok_drops_structured_content_for_non_objects() {
        let r = tool_ok(serde_json::json!([1, 2, 3]));
        assert!(r.structured_content.is_none());
        assert!(r.is_error.is_none());
    }

    #[test]
    fn tool_error_marks_is_error_true() {
        let r = tool_error("boom");
        assert_eq!(r.is_error, Some(true));
    }
}
