//! An MCP client, for the tools a llama session has that `llama-server` does
//! not provide itself.
//!
//! **Why this is here and not a flag on `llama-server`.** That server can host
//! MCP servers (`--mcp-servers-json`), but only ones it can *spawn*: its
//! configuration is Cursor's, and an entry without a `command` is skipped with
//! "MCP server 'exa' has no command". Exa's is a remote HTTP endpoint with
//! nothing to spawn, so reaching it that way means a local process bridging
//! stdio to HTTP -- a Node install on the machine serving the model, and a
//! package to keep current, for what is three JSON-RPC calls.
//!
//! llama.cpp's own web UI does not do that either. It ships Exa in a
//! "recommended servers" list and connects to `https://mcp.exa.ai/mcp`
//! *itself*, from the browser. This is the same arrangement with this server
//! in the browser's place, and it is the right one for a second reason: it
//! puts the search on the machine running the backend rather than on whichever
//! machine happens to be serving the model, which may have no route out at
//! all.
//!
//! **Only the three calls a tool needs.** `initialize`, `tools/list`,
//! `tools/call`. Nothing here implements resources, prompts, sampling or the
//! server-to-client stream, because nothing here uses them; a session's tools
//! are a list fetched once and a call made on demand. That is why this is a
//! file rather than a dependency on a protocol crate -- there is no spec
//! surface to get subtly wrong, only a request and its reply.
//!
//! Transport is "streamable HTTP": every message is a POST, and the reply is
//! either JSON or a one-event SSE stream carrying the same JSON. Both are
//! accepted because which one arrives is the server's choice, not ours.

use anyhow::{Context, Result, bail};
use serde_json::{Value, json};

/// Identifies this client to an MCP server.
///
/// Not politeness: Exa's endpoint is behind Cloudflare, which answers **403**
/// to a request with no `User-Agent` at all (measured 2026-09-19 -- the same
/// request with one succeeds). A client that omitted it would look exactly
/// like a server that was refusing us.
const USER_AGENT: &str = concat!("ai-server/", env!("CARGO_PKG_VERSION"));

/// The protocol version this speaks. Sent at `initialize`; a server that
/// prefers another says so in its answer and this goes along with whatever it
/// then sends, since none of the three calls here has changed between
/// versions.
const PROTOCOL_VERSION: &str = "2025-06-18";

/// How long any one call may take.
///
/// Generous because a web search is a search: Exa fetches and cleans pages
/// before answering. Bounded at all because this blocks a turn, and a tool
/// that never returns is a session that never speaks again.
const CALL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);

/// A connected MCP server, and the tools it offered.
pub struct McpServer {
    /// The name this server is configured under. It prefixes every tool, so
    /// two servers offering `search` are two different tools.
    name: String,
    url: String,
    /// What the server called this conversation, when it named one. Sent back
    /// on every later request; a server that keeps no session sends no header
    /// and this stays `None`.
    session: Option<String>,
    /// The tool names this server answers to, without the prefix, keyed by the
    /// prefixed name the model is given.
    tools: Vec<McpTool>,
}

/// One tool an MCP server offers, in both the names it has.
pub struct McpTool {
    /// `{server}_{tool}` -- what the model calls it, and what comes back in a
    /// tool call. Prefixed the way `llama-server` prefixes the MCP tools it
    /// hosts itself, so a reader sees one naming convention whichever side a
    /// tool came from.
    pub qualified: String,
    /// What the server calls it.
    bare: String,
    /// The OpenAI-shaped function definition sent to the model.
    pub definition: Value,
}

impl McpServer {
    /// Connects, handshakes, and asks what it can do.
    ///
    /// All three steps or none: a server that answered `initialize` and then
    /// failed to list its tools is not a server with no tools, and returning
    /// an empty list for it would put a session on screen that silently
    /// cannot search.
    pub fn connect(name: &str, url: &str) -> Result<Self> {
        let mut server = Self {
            name: name.to_string(),
            url: url.to_string(),
            session: None,
            tools: Vec::new(),
        };
        server
            .request(
                1,
                "initialize",
                json!({
                    "protocolVersion": PROTOCOL_VERSION,
                    "capabilities": {},
                    "clientInfo": {"name": "ai-server", "title": "AI Sessions", "version": env!("CARGO_PKG_VERSION")},
                }),
            )
            .with_context(|| format!("handshaking with the {name} MCP server at {url}"))?;
        // A notification: no id, and the server answers with no body. Sent
        // because the specification requires it before any other call, and
        // Exa's server does enforce it.
        server.notify("notifications/initialized")?;
        let listed = server
            .request(2, "tools/list", json!({}))
            .with_context(|| format!("asking the {name} MCP server what it offers"))?;
        server.tools = listed
            .get("tools")
            .and_then(Value::as_array)
            .map(|tools| {
                tools
                    .iter()
                    .filter_map(|tool| server.describe(tool))
                    .collect()
            })
            .unwrap_or_default();
        Ok(server)
    }

    /// Turns one entry of `tools/list` into the function definition a model is
    /// given, or `None` for one this cannot name or call.
    fn describe(&self, tool: &Value) -> Option<McpTool> {
        let bare = tool.get("name").and_then(Value::as_str)?.to_string();
        let qualified = format!("{}_{bare}", self.name);
        let mut function = serde_json::Map::new();
        function.insert("name".into(), json!(qualified));
        if let Some(description) = tool.get("description").and_then(Value::as_str) {
            function.insert("description".into(), json!(description));
        }
        // `inputSchema` in MCP, `parameters` in the OpenAI shape: the same
        // JSON Schema under two names. A tool that declares none takes no
        // arguments, which is an empty object rather than an absent key --
        // some templates render the key unconditionally.
        function.insert(
            "parameters".into(),
            tool.get("inputSchema")
                .cloned()
                .unwrap_or_else(|| json!({"type": "object", "properties": {}})),
        );
        Some(McpTool {
            qualified,
            bare,
            definition: json!({"type": "function", "function": function}),
        })
    }

    pub fn tools(&self) -> &[McpTool] {
        &self.tools
    }

    /// Runs one of this server's tools, named as the model named it.
    ///
    /// The result is the text a model is shown. A tool the server reports as
    /// failing is **not** an error here: `isError` means the tool ran and went
    /// wrong -- a search that found nothing, a page that would not fetch --
    /// and the model is the one that has to know, so it comes back as its own
    /// message. An error is reserved for not having reached the server at all.
    pub fn call(&mut self, qualified: &str, arguments: &Value) -> Result<String> {
        // The bare name is taken before the call, because the call needs the
        // whole of `self` and the tool list is part of it.
        let bare = self
            .tools
            .iter()
            .find(|tool| tool.qualified == qualified)
            .map(|tool| tool.bare.clone())
            .with_context(|| format!("{} does not offer {qualified}", self.name))?;
        let result = self.request(
            3,
            "tools/call",
            json!({"name": bare, "arguments": arguments}),
        )?;
        Ok(rendered(&result))
    }

    /// One request, and its result.
    ///
    /// `&self` rather than `&mut self` everywhere but the handshake would be
    /// tidier and is wrong: the session header is assigned by the server on
    /// the first reply and has to be kept.
    fn request(&mut self, id: u64, method: &str, params: Value) -> Result<Value> {
        let body = json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params});
        let answer = self.post(&body)?.with_context(|| {
            format!(
                "the {} MCP server answered {method} with nothing",
                self.name
            )
        })?;
        if let Some(message) = answer.pointer("/error/message").and_then(Value::as_str) {
            bail!("{} refused {method}: {message}", self.name);
        }
        answer
            .get("result")
            .cloned()
            .with_context(|| format!("the {} MCP server's {method} carried no result", self.name))
    }

    /// A message with no id, which is answered with no body.
    fn notify(&mut self, method: &str) -> Result<()> {
        self.post(&json!({"jsonrpc": "2.0", "method": method}))?;
        Ok(())
    }

    /// Posts one JSON-RPC message and returns whatever came back, which for a
    /// notification is nothing.
    fn post(&mut self, body: &Value) -> Result<Option<Value>> {
        let mut request = ureq::post(&self.url)
            .config()
            .timeout_global(Some(CALL_TIMEOUT))
            .build()
            .header("content-type", "application/json")
            // Both, because which one a server replies with is its choice.
            .header("accept", "application/json, text/event-stream")
            .header("user-agent", USER_AGENT);
        if let Some(session) = &self.session {
            request = request.header("mcp-session-id", session);
        }
        let mut response = request
            .send_json(body)
            .with_context(|| format!("reaching the {} MCP server at {}", self.name, self.url))?;
        if let Some(session) = response
            .headers()
            .get("mcp-session-id")
            .and_then(|value| value.to_str().ok())
        {
            self.session = Some(session.to_string());
        }
        let streamed = response
            .headers()
            .get("content-type")
            .and_then(|value| value.to_str().ok())
            .is_some_and(|value| value.contains("text/event-stream"));
        let text = response
            .body_mut()
            .read_to_string()
            .with_context(|| format!("reading the {} MCP server's answer", self.name))?;
        Ok(first_message(&text, streamed))
    }
}

/// The first JSON-RPC message in a reply body.
///
/// One, not all: every call here carries a single id and the server answers it
/// once. Server-sent events are unwrapped to their payload lines; a plain JSON
/// body is itself.
fn first_message(text: &str, streamed: bool) -> Option<Value> {
    if streamed {
        return text
            .lines()
            .filter_map(|line| line.strip_prefix("data: "))
            .find_map(|payload| serde_json::from_str(payload).ok());
    }
    serde_json::from_str(text.trim()).ok()
}

/// A `tools/call` result as the text a model is given.
///
/// MCP answers with a list of content blocks; the text ones are joined and the
/// rest are named rather than dropped, because a model told nothing came back
/// will try again. `structuredContent` is used when there is no text at all,
/// which is how some servers answer entirely.
fn rendered(result: &Value) -> String {
    let blocks = result.get("content").and_then(Value::as_array);
    let mut parts: Vec<String> = Vec::new();
    for block in blocks.into_iter().flatten() {
        match block.get("type").and_then(Value::as_str) {
            Some("text") => parts.push(
                block
                    .get("text")
                    .and_then(Value::as_str)
                    .unwrap_or_default()
                    .to_string(),
            ),
            Some(kind) => parts.push(format!("[{kind} content, which this session cannot show]")),
            None => {}
        }
    }
    if parts.iter().all(|part| part.trim().is_empty())
        && let Some(structured) = result.get("structuredContent")
    {
        return structured.to_string();
    }
    parts.join("\n")
}

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

    #[test]
    fn an_event_stream_body_is_unwrapped_to_its_payload() {
        let body =
            "event: message\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"ok\":true}}\n\n";
        assert_eq!(
            first_message(body, true),
            Some(json!({"jsonrpc": "2.0", "id": 1, "result": {"ok": true}})),
        );
    }

    #[test]
    fn a_plain_json_body_is_the_message() {
        let body = "  {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}\n";
        assert_eq!(
            first_message(body, false),
            Some(json!({"jsonrpc": "2.0", "id": 1, "result": {}})),
        );
    }

    #[test]
    /// A notification's reply, which is nothing at all.
    fn an_empty_body_is_no_message() {
        assert_eq!(first_message("", false), None);
        assert_eq!(first_message("event: ping\n", true), None);
    }

    #[test]
    fn text_blocks_are_joined_and_other_kinds_are_named() {
        let result = json!({"content": [
            {"type": "text", "text": "first"},
            {"type": "image", "data": "…"},
            {"type": "text", "text": "second"},
        ]});
        assert_eq!(
            rendered(&result),
            "first\n[image content, which this session cannot show]\nsecond",
        );
    }

    #[test]
    /// A server that answers only in structured form. Rendering "" for it
    /// would tell the model the search came back empty, which is a different
    /// fact from the one that is true.
    fn a_result_with_no_text_falls_back_to_its_structured_form() {
        let result = json!({"content": [], "structuredContent": {"hits": 2}});
        assert_eq!(rendered(&result), "{\"hits\":2}");
    }

    #[test]
    /// The real endpoint, which is the only thing that can confirm the
    /// handshake, the session header and the SSE unwrapping all agree with a
    /// server nobody here wrote. Skipped without network rather than failed:
    /// `./run-tests.sh` has to pass on a machine with no route out.
    fn exa_answers_a_search_over_the_real_protocol() {
        let Ok(mut server) = McpServer::connect("exa", super::super::EXA_MCP_URL) else {
            eprintln!("skipping: could not reach Exa");
            return;
        };
        assert!(
            server
                .tools()
                .iter()
                .any(|tool| tool.qualified == "exa_web_search_exa"),
            "Exa offered {:?}",
            server
                .tools()
                .iter()
                .map(|tool| &tool.qualified)
                .collect::<Vec<_>>(),
        );
        let answer = server
            .call(
                "exa_web_search_exa",
                &json!({"query": "llama.cpp server", "numResults": 1}),
            )
            .expect("search");
        assert!(!answer.trim().is_empty(), "a search returned nothing");
    }
}