collet 0.1.1

use serde::{Deserialize, Serialize};
use serde_json::Value;

// ---------------------------------------------------------------------------
// JSON-RPC 2.0 envelope
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcRequest {
    pub jsonrpc: String,
    pub id: u64,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

impl JsonRpcRequest {
    pub fn new(id: u64, method: &str, params: Option<Value>) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            method: method.to_string(),
            params,
        }
    }
}

/// JSON-RPC 2.0 notification — no `id` field per spec.
///
/// MCP requires the client to send `notifications/initialized` after the
/// `initialize` handshake completes and before issuing any other request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcNotification {
    pub jsonrpc: String,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

impl JsonRpcNotification {
    pub fn new(method: &str, params: Option<Value>) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            method: method.to_string(),
            params,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcError {
    pub code: i64,
    pub message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonRpcResponse {
    pub jsonrpc: String,
    pub id: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<JsonRpcError>,
}

// ---------------------------------------------------------------------------
// MCP-specific types
// ---------------------------------------------------------------------------

/// A tool definition returned by the MCP server via `tools/list`.
///
/// Wire format uses camelCase per MCP spec — `inputSchema` is the canonical
/// field name. Without the rename, the JSON Schema would silently deserialize
/// to `Value::Null`, leaving the LLM unable to construct valid arguments and
/// causing constant `-32602` "Invalid arguments" errors from the server.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpToolDefinition {
    pub name: String,
    #[serde(default)]
    pub description: String,
    /// JSON Schema describing the tool's input parameters.
    #[serde(default)]
    pub input_schema: Value,
}

/// A single content block inside a tool result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpContent {
    #[serde(rename = "type")]
    pub type_: String,
    #[serde(default)]
    pub text: String,
}

/// The result payload returned by `tools/call`.
///
/// Per MCP spec, `isError` (camelCase) signals that the tool itself reported
/// failure — the JSON-RPC envelope is still `result`, not `error`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpToolResult {
    #[serde(default)]
    pub content: Vec<McpContent>,
    #[serde(default)]
    pub is_error: bool,
}

/// Server info returned by the `initialize` handshake.
///
/// The MCP `initialize` result has the shape:
/// ```json
/// {
///   "protocolVersion": "...",
///   "capabilities": { ... },
///   "serverInfo": { "name": "...", "version": "..." },
///   "instructions": "..."
/// }
/// ```
/// `name` and `version` live under `serverInfo`, not at the root — the client
/// extracts them manually so this struct stays a flat, ergonomic view.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct McpServerInfo {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub version: String,
    #[serde(default)]
    pub protocol_version: String,
    #[serde(default)]
    pub capabilities: Value,
    /// Optional server-level instructions describing the server's purpose
    /// and how its tools should be used.
    #[serde(default)]
    pub instructions: Option<String>,
}

// ---------------------------------------------------------------------------
// Well-known MCP method names
// ---------------------------------------------------------------------------

pub const METHOD_INITIALIZE: &str = "initialize";
pub const METHOD_TOOLS_LIST: &str = "tools/list";
pub const METHOD_TOOLS_CALL: &str = "tools/call";
pub const METHOD_NOTIFICATIONS_INITIALIZED: &str = "notifications/initialized";

/// MCP protocol version this client speaks.
///
/// Per spec, the server may negotiate down to an earlier version it supports.
/// We advertise the latest stable revision; older servers will reply with
/// their own `protocolVersion` and we proceed regardless.
pub const MCP_PROTOCOL_VERSION: &str = "2025-06-18";

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn test_json_rpc_request_roundtrip() {
        let req = JsonRpcRequest::new(
            1,
            METHOD_TOOLS_CALL,
            Some(json!({"name": "read_file", "arguments": {"path": "/tmp/x"}})),
        );
        let serialized = serde_json::to_string(&req).unwrap();
        let deserialized: JsonRpcRequest = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized.id, 1);
        assert_eq!(deserialized.method, "tools/call");
        assert_eq!(deserialized.jsonrpc, "2.0");
        assert!(deserialized.params.is_some());
    }

    #[test]
    fn test_json_rpc_response_roundtrip() {
        let resp = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: 42,
            result: Some(json!({"tools": []})),
            error: None,
        };
        let serialized = serde_json::to_string(&resp).unwrap();
        let deserialized: JsonRpcResponse = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized.id, 42);
        assert!(deserialized.result.is_some());
        assert!(deserialized.error.is_none());
    }

    #[test]
    fn test_json_rpc_error_response_roundtrip() {
        let resp = JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id: 7,
            result: None,
            error: Some(JsonRpcError {
                code: -32601,
                message: "Method not found".to_string(),
                data: None,
            }),
        };
        let serialized = serde_json::to_string(&resp).unwrap();
        let deserialized: JsonRpcResponse = serde_json::from_str(&serialized).unwrap();
        assert!(deserialized.error.is_some());
        let err = deserialized.error.unwrap();
        assert_eq!(err.code, -32601);
        assert_eq!(err.message, "Method not found");
    }

    #[test]
    fn test_mcp_tool_definition_roundtrip() {
        let tool = McpToolDefinition {
            name: "read_file".to_string(),
            description: "Read a file from disk".to_string(),
            input_schema: json!({
                "type": "object",
                "properties": {
                    "path": { "type": "string" }
                },
                "required": ["path"]
            }),
        };
        let serialized = serde_json::to_string(&tool).unwrap();
        let deserialized: McpToolDefinition = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized.name, "read_file");
        assert_eq!(deserialized.description, "Read a file from disk");
    }

    #[test]
    fn test_mcp_tool_result_roundtrip() {
        let result = McpToolResult {
            content: vec![
                McpContent {
                    type_: "text".to_string(),
                    text: "hello world".to_string(),
                },
                McpContent {
                    type_: "text".to_string(),
                    text: "second block".to_string(),
                },
            ],
            is_error: false,
        };
        let serialized = serde_json::to_string(&result).unwrap();
        let deserialized: McpToolResult = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized.content.len(), 2);
        assert_eq!(deserialized.content[0].text, "hello world");
        assert!(!deserialized.is_error);
    }

    /// Real-world MCP servers send `inputSchema` in camelCase. The struct
    /// must rename to it — otherwise the schema is dropped and tool calls
    /// fail with `-32602 Invalid arguments`.
    #[test]
    fn test_mcp_tool_definition_camelcase_input_schema() {
        let wire = json!({
            "name": "search",
            "description": "Search docs",
            "inputSchema": {
                "type": "object",
                "properties": { "query": { "type": "string" } },
                "required": ["query"]
            }
        });
        let parsed: McpToolDefinition = serde_json::from_value(wire).unwrap();
        assert_eq!(parsed.name, "search");
        assert!(
            parsed.input_schema.get("properties").is_some(),
            "inputSchema must round-trip — was {:?}",
            parsed.input_schema
        );
    }

    /// `isError` (camelCase) on a tool result must surface the failure flag.
    #[test]
    fn test_mcp_tool_result_is_error_camelcase() {
        let wire = json!({
            "content": [{ "type": "text", "text": "boom" }],
            "isError": true
        });
        let parsed: McpToolResult = serde_json::from_value(wire).unwrap();
        assert!(parsed.is_error);
        assert_eq!(parsed.content[0].text, "boom");
    }

    #[test]
    fn test_mcp_server_info_roundtrip() {
        let info = McpServerInfo {
            name: "test-server".to_string(),
            version: "0.1.0".to_string(),
            protocol_version: MCP_PROTOCOL_VERSION.to_string(),
            capabilities: json!({"tools": true}),
            instructions: None,
        };
        let serialized = serde_json::to_string(&info).unwrap();
        let deserialized: McpServerInfo = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized.name, "test-server");
        assert_eq!(deserialized.version, "0.1.0");
    }

    #[test]
    fn test_notification_has_no_id() {
        let n = JsonRpcNotification::new(METHOD_NOTIFICATIONS_INITIALIZED, None);
        let s = serde_json::to_string(&n).unwrap();
        assert!(!s.contains("\"id\""), "notification must omit id: {s}");
        assert!(s.contains("notifications/initialized"));
    }
}