bctx-nexus 0.1.25

use std::io::{BufRead, Write};

use anyhow::Result;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};

use crate::server::NexusServer;

// ── JSON-RPC 2.0 types ────────────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
struct RawMessage {
    #[serde(default)]
    id: Option<Value>,
    #[serde(default)]
    method: Option<String>,
    #[serde(default)]
    params: Option<Value>,
}

#[derive(Debug, Serialize)]
struct Response {
    jsonrpc: &'static str,
    id: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<RpcError>,
}

#[derive(Debug, Serialize)]
struct RpcError {
    code: i32,
    message: String,
}

impl Response {
    fn ok(id: Value, result: Value) -> Self {
        Self {
            jsonrpc: "2.0",
            id,
            result: Some(result),
            error: None,
        }
    }
    fn err(id: Value, code: i32, message: String) -> Self {
        Self {
            jsonrpc: "2.0",
            id,
            result: None,
            error: Some(RpcError { code, message }),
        }
    }
}

// ── Public API ────────────────────────────────────────────────────────────────

/// Process a single JSON-RPC message string.
/// Returns `Some(encoded_response)` for requests, `None` for notifications.
/// Exposed for testing; `run_stdio` delegates here.
pub fn process_message(server: &NexusServer, line: &str) -> Option<String> {
    let msg = match serde_json::from_str::<RawMessage>(line) {
        Err(e) => {
            let resp = Response::err(Value::Null, -32700, format!("parse error: {e}"));
            return Some(serde_json::to_string(&resp).unwrap_or_default());
        }
        Ok(m) => m,
    };

    // JSON-RPC 2.0 §4: messages without an "id" are notifications — no response.
    if msg.id.is_none() {
        if let Some(ref m) = msg.method {
            handle_notification(server, m);
        }
        return None;
    }

    let id = msg.id.clone().unwrap();
    let method = match msg.method.as_deref() {
        Some(m) => m,
        None => {
            return Some(
                serde_json::to_string(&Response::err(id, -32600, "missing method".into()))
                    .unwrap_or_default(),
            )
        }
    };
    let resp = handle_request(server, id, method, msg.params);
    Some(serde_json::to_string(&resp).unwrap_or_default())
}

/// Run the MCP stdio server: reads newline-delimited JSON-RPC from stdin,
/// writes responses to stdout. Notifications (no id) are silently consumed.
pub fn run_stdio(server: &NexusServer) -> Result<()> {
    let stdin = std::io::stdin();
    let stdout = std::io::stdout();
    let mut out = std::io::BufWriter::new(stdout.lock());

    for line in stdin.lock().lines() {
        // EOF or broken pipe → stdin closed by host; stop reading and flush Vault.
        let line = match line {
            Ok(l) => l,
            Err(_) => break,
        };
        if line.trim().is_empty() {
            continue;
        }
        if let Some(encoded) = process_message(server, &line) {
            writeln!(out, "{encoded}")?;
            out.flush()?;
        }
    }

    atlas::vault_store::flush();
    Ok(())
}

// ── Notification handler (fire-and-forget) ────────────────────────────────────

fn handle_notification(server: &NexusServer, method: &str) {
    match method {
        "notifications/initialized" | "notifications/cancelled" => {}
        _ => tracing::debug!(method, "unhandled MCP notification"),
    }
    let _ = server;
}

// ── Request handler ───────────────────────────────────────────────────────────

fn handle_request(
    server: &NexusServer,
    id: Value,
    method: &str,
    params: Option<Value>,
) -> Response {
    match method {
        // ── MCP lifecycle ──────────────────────────────────────────────────────
        "initialize" => {
            let proto = params
                .as_ref()
                .and_then(|p| p.get("protocolVersion"))
                .and_then(|v| v.as_str())
                .unwrap_or("2024-11-05");
            let negotiated = if proto >= "2025-03-26" {
                "2025-03-26"
            } else if proto >= "2024-11-05" {
                "2024-11-05"
            } else {
                proto
            };
            Response::ok(
                id,
                json!({
                    "protocolVersion": negotiated,
                    "capabilities": { "tools": {} },
                    "serverInfo": { "name": "bctx", "version": env!("CARGO_PKG_VERSION") }
                }),
            )
        }

        "ping" => Response::ok(id, json!({})),

        // ── Tool protocol ──────────────────────────────────────────────────────
        "tools/list" => Response::ok(id, server.gateway.list_tools()),

        "tools/call" => {
            let p = params.unwrap_or(Value::Null);
            let name = match p.get("name").and_then(|v| v.as_str()) {
                Some(n) => n.to_string(),
                None => return Response::err(id, -32602, "missing params.name".into()),
            };
            let args = p
                .get("arguments")
                .cloned()
                .unwrap_or_else(|| Value::Object(Default::default()));
            let caller = p
                .get("_caller")
                .and_then(|v| v.as_str())
                .unwrap_or("mcp-client");

            match server.gateway.dispatch(&name, args, caller) {
                Ok(full) => {
                    // Surface the skill's own payload; the wrapper metadata is noise for the LLM.
                    let payload = full.get("result").unwrap_or(&full);
                    let text = serde_json::to_string_pretty(payload)
                        .unwrap_or_else(|_| payload.to_string());
                    Response::ok(
                        id,
                        json!({
                            "content": [{ "type": "text", "text": text }]
                        }),
                    )
                }
                Err(e) => {
                    // MCP spec §5.9: tool errors → content with isError:true.
                    // JSON-RPC errors are reserved for protocol faults only.
                    Response::ok(
                        id,
                        json!({
                            "content": [{ "type": "text", "text": e.to_string() }],
                            "isError": true
                        }),
                    )
                }
            }
        }

        // ── Optional capability probes ─────────────────────────────────────────
        // Claude Code probes these on connect; returning empty lists prevents it
        // from treating "method not found" as a broken server.
        "prompts/list" => Response::ok(id, json!({ "prompts": [] })),
        "resources/list" => Response::ok(id, json!({ "resources": [] })),
        "logging/setLevel" => Response::ok(id, json!({})),

        other => Response::err(id, -32601, format!("method not found: {other}")),
    }
}