//! A Model Context Protocol server over stdio (oxygen#159): JSON-RPC 2.0,
//! one message per line, as MCP's stdio transport has it. This is only the
//! protocol: `initialize`, `ping`, `tools/list` and `tools/call`, with the
//! tools themselves supplied by the caller (`oxy mcp`, and `oh mcp` in the
//! hub, which reaches this module through its pinned oxygen).
//!
//! Nothing is written to stdout but protocol messages: a tool's own
//! output goes into its result.

use serde_json::{json, Value};
use std::io::{BufRead, Write};

/// The protocol versions this server speaks, newest first. A client asking
/// for one of them gets it back; any other gets the newest.
pub const PROTOCOL_VERSIONS: &[&str] = &["2025-06-18", "2025-03-26", "2024-11-05"];

/// One tool as `tools/list` describes it.
#[derive(Clone, Debug)]
pub struct Tool {
    pub name: String,
    pub description: String,
    /// A JSON Schema object for the arguments.
    pub input_schema: Value,
    /// Reads only: no change to the repository, the remote or the hub.
    pub read_only: bool,
    /// May lose work or history (undo, abandon, deletes).
    pub destructive: bool,
}

impl Tool {
    fn describe(&self) -> Value {
        json!({
            "name": self.name,
            "description": self.description,
            "inputSchema": self.input_schema,
            "annotations": {
                "readOnlyHint": self.read_only,
                "destructiveHint": self.destructive,
            },
        })
    }
}

/// What a tool call produced: text for the model to read, the same as JSON
/// when there is some, and whether it failed. A failed tool is a result, not
/// a protocol error, so the model sees why.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ToolOutput {
    pub text: String,
    pub structured: Option<Value>,
    pub is_error: bool,
}

impl ToolOutput {
    pub fn ok(text: impl Into<String>) -> ToolOutput {
        ToolOutput {
            text: text.into(),
            structured: None,
            is_error: false,
        }
    }

    pub fn error(text: impl Into<String>) -> ToolOutput {
        ToolOutput {
            text: text.into(),
            structured: None,
            is_error: true,
        }
    }

    fn to_json(&self) -> Value {
        let mut out = json!({
            "content": [{ "type": "text", "text": self.text }],
            "isError": self.is_error,
        });
        // Structured content must be an object; an array is wrapped.
        if let Some(s) = &self.structured {
            let obj = if s.is_object() {
                s.clone()
            } else {
                json!({ "result": s })
            };
            out["structuredContent"] = obj;
        }
        out
    }
}

/// The tools a server offers.
pub trait Tools {
    fn list(&self) -> Vec<Tool>;
    /// Run a tool. `None` when no tool has that name.
    fn call(&self, name: &str, args: &Value) -> Option<ToolOutput>;
}

/// Who the server is, for `initialize`.
#[derive(Clone, Debug)]
pub struct ServerInfo {
    pub name: String,
    pub version: String,
    /// Shown to the model: what the server is for and how to use it safely.
    pub instructions: String,
}

const PARSE_ERROR: i64 = -32700;
const INVALID_REQUEST: i64 = -32600;
const METHOD_NOT_FOUND: i64 = -32601;
const INVALID_PARAMS: i64 = -32602;

fn error(id: Value, code: i64, message: &str) -> Value {
    json!({ "jsonrpc": "2.0", "id": id, "error": { "code": code, "message": message } })
}

fn result(id: Value, result: Value) -> Value {
    json!({ "jsonrpc": "2.0", "id": id, "result": result })
}

/// Answer one message; `None` for a notification (no reply is due).
pub fn handle(info: &ServerInfo, tools: &dyn Tools, message: &Value) -> Option<Value> {
    let Some(obj) = message.as_object() else {
        return Some(error(
            Value::Null,
            INVALID_REQUEST,
            "a message is a JSON object",
        ));
    };
    let method = obj.get("method").and_then(Value::as_str);
    let id = obj.get("id").cloned();
    let Some(method) = method else {
        // A response to something we never sent, or nonsense: a request
        // without a method gets an error, anything else is ignored.
        return id.map(|id| error(id, INVALID_REQUEST, "a request needs a method"));
    };
    // Notifications (initialized, cancelled, …) carry no id and get no reply.
    let id = id?;
    let params = obj.get("params").cloned().unwrap_or(Value::Null);
    Some(match method {
        "initialize" => {
            let asked = params.get("protocolVersion").and_then(Value::as_str);
            let version = asked
                .filter(|v| PROTOCOL_VERSIONS.contains(v))
                .unwrap_or(PROTOCOL_VERSIONS[0]);
            result(
                id,
                json!({
                    "protocolVersion": version,
                    "capabilities": { "tools": { "listChanged": false } },
                    "serverInfo": { "name": info.name, "version": info.version },
                    "instructions": info.instructions,
                }),
            )
        }
        "ping" => result(id, json!({})),
        "tools/list" => result(
            id,
            json!({ "tools": tools.list().iter().map(Tool::describe).collect::<Vec<_>>() }),
        ),
        "tools/call" => {
            let Some(name) = params.get("name").and_then(Value::as_str) else {
                return Some(error(
                    id,
                    INVALID_PARAMS,
                    "tools/call needs the tool's name",
                ));
            };
            let args = params
                .get("arguments")
                .cloned()
                .unwrap_or_else(|| json!({}));
            if !args.is_object() {
                return Some(error(
                    id,
                    INVALID_PARAMS,
                    "a tool's arguments are a JSON object",
                ));
            }
            match tools.call(name, &args) {
                Some(out) => result(id, out.to_json()),
                None => error(id, INVALID_PARAMS, &format!("no tool named {name}")),
            }
        }
        other => error(id, METHOD_NOT_FOUND, &format!("{other} is not supported")),
    })
}

/// Serve until the input ends: read a message per line, write a reply per
/// request. A line that is not JSON gets a parse error; blank lines are
/// skipped.
pub fn serve(
    info: &ServerInfo,
    tools: &dyn Tools,
    input: impl BufRead,
    mut output: impl Write,
) -> std::io::Result<()> {
    for line in input.lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        let reply = match serde_json::from_str::<Value>(&line) {
            Ok(Value::Array(batch)) => {
                let replies: Vec<Value> = batch
                    .iter()
                    .filter_map(|m| handle(info, tools, m))
                    .collect();
                (!replies.is_empty()).then_some(Value::Array(replies))
            }
            Ok(message) => handle(info, tools, &message),
            Err(e) => Some(error(Value::Null, PARSE_ERROR, &format!("not JSON: {e}"))),
        };
        if let Some(reply) = reply {
            writeln!(output, "{reply}")?;
            output.flush()?;
        }
    }
    Ok(())
}

/// A string argument, if given and a string.
pub fn arg_str<'a>(args: &'a Value, key: &str) -> Option<&'a str> {
    args.get(key).and_then(Value::as_str)
}

/// A list of strings, if given; a single string counts as a list of one.
pub fn arg_strings(args: &Value, key: &str) -> Result<Vec<String>, String> {
    match args.get(key) {
        None | Some(Value::Null) => Ok(vec![]),
        Some(Value::String(s)) => Ok(vec![s.clone()]),
        Some(Value::Array(a)) => a
            .iter()
            .map(|v| {
                v.as_str()
                    .map(str::to_string)
                    .ok_or_else(|| format!("{key} must be a list of strings"))
            })
            .collect(),
        Some(_) => Err(format!("{key} must be a list of strings")),
    }
}

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

    struct Echo;

    impl Tools for Echo {
        fn list(&self) -> Vec<Tool> {
            vec![Tool {
                name: "echo".into(),
                description: "says it back".into(),
                input_schema: json!({"type": "object"}),
                read_only: true,
                destructive: false,
            }]
        }
        fn call(&self, name: &str, args: &Value) -> Option<ToolOutput> {
            (name == "echo").then(|| ToolOutput {
                text: args.to_string(),
                structured: Some(json!([1, 2])),
                is_error: false,
            })
        }
    }

    fn info() -> ServerInfo {
        ServerInfo {
            name: "t".into(),
            version: "1".into(),
            instructions: "be careful".into(),
        }
    }

    fn talk(lines: &str) -> Vec<Value> {
        let mut out = Vec::new();
        serve(&info(), &Echo, lines.as_bytes(), &mut out).unwrap();
        String::from_utf8(out)
            .unwrap()
            .lines()
            .map(|l| serde_json::from_str(l).unwrap())
            .collect()
    }

    #[test]
    fn a_session_initializes_lists_and_calls() {
        let r = talk(concat!(
            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"c","version":"0"}}}"#,
            "\n",
            r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#,
            "\n\n",
            r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#,
            "\n",
            r#"{"jsonrpc":"2.0","id":"three","method":"tools/call","params":{"name":"echo","arguments":{"a":1}}}"#,
            "\n",
        ));
        assert_eq!(r.len(), 3, "the notification gets no reply: {r:?}");
        assert_eq!(
            r[0]["result"]["protocolVersion"], "2025-03-26",
            "a version we speak is kept"
        );
        assert_eq!(r[0]["result"]["instructions"], "be careful");
        assert_eq!(r[1]["result"]["tools"][0]["name"], "echo");
        assert_eq!(
            r[1]["result"]["tools"][0]["annotations"]["readOnlyHint"],
            true
        );
        assert_eq!(r[2]["id"], "three");
        assert_eq!(r[2]["result"]["content"][0]["text"], r#"{"a":1}"#);
        assert_eq!(
            r[2]["result"]["structuredContent"],
            json!({"result": [1, 2]}),
            "an array is wrapped in an object"
        );
        assert_eq!(r[2]["result"]["isError"], false);
    }

    #[test]
    fn errors_are_protocol_errors_and_unknown_versions_get_the_newest() {
        let r = talk(concat!(
            "not json\n",
            r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"1999-01-01"}}"#,
            "\n",
            r#"{"jsonrpc":"2.0","id":2,"method":"resources/list"}"#,
            "\n",
            r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"nope"}}"#,
            "\n",
            r#"{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"echo","arguments":[1]}}"#,
            "\n",
            r#"[{"jsonrpc":"2.0","id":5,"method":"ping"},{"jsonrpc":"2.0","method":"notifications/cancelled"}]"#,
            "\n",
        ));
        assert_eq!(r[0]["error"]["code"], PARSE_ERROR);
        assert_eq!(r[1]["result"]["protocolVersion"], PROTOCOL_VERSIONS[0]);
        assert_eq!(r[2]["error"]["code"], METHOD_NOT_FOUND);
        assert_eq!(r[3]["error"]["code"], INVALID_PARAMS);
        assert_eq!(r[4]["error"]["code"], INVALID_PARAMS);
        assert_eq!(
            r[5],
            json!([{"jsonrpc": "2.0", "id": 5, "result": {}}]),
            "a batch is answered as a batch"
        );
    }

    #[test]
    fn string_arguments_and_lists() {
        let a = json!({"one": "x", "many": ["a", "b"], "bad": [1]});
        assert_eq!(arg_str(&a, "one"), Some("x"));
        assert_eq!(arg_strings(&a, "one").unwrap(), vec!["x"]);
        assert_eq!(arg_strings(&a, "many").unwrap(), vec!["a", "b"]);
        assert_eq!(arg_strings(&a, "missing").unwrap(), Vec::<String>::new());
        assert!(arg_strings(&a, "bad").is_err());
    }
}