peeker 1.1.0

use anyhow::Result;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::io::{BufRead, BufReader, Write};

use crate::parser;

const SERVER_NAME: &str = "peeker";
const SERVER_VERSION: &str = env!("CARGO_PKG_VERSION");

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

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

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

impl JsonRpcResponse {
    fn success(id: Value, result: Value) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: Some(result),
            error: None,
        }
    }

    fn error(id: Value, code: i32, message: String) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: None,
            error: Some(JsonRpcError {
                code,
                message,
                data: None,
            }),
        }
    }
}

pub fn run_mcp_server() -> Result<()> {
    let stdin = std::io::stdin();
    let mut stdout = std::io::stdout();
    let reader = BufReader::new(stdin.lock());

    for line in reader.lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }

        let response = match serde_json::from_str::<JsonRpcRequest>(&line) {
            Ok(request) => handle_request(request),
            Err(e) => Some(JsonRpcResponse::error(
                Value::Null,
                -32700,
                format!("Parse error: {}", e),
            )),
        };

        // Only send response if it's not a notification
        if let Some(response) = response {
            let response_json = serde_json::to_string(&response)?;
            writeln!(stdout, "{}", response_json)?;
            stdout.flush()?;
        }
    }

    Ok(())
}

fn handle_request(request: JsonRpcRequest) -> Option<JsonRpcResponse> {
    // Notifications (no id) don't get responses
    let is_notification = request.id.is_none();
    let id = request.id.unwrap_or(Value::Null);

    // Handle notifications silently
    if is_notification || request.method.starts_with("notifications/") {
        return None;
    }

    Some(match request.method.as_str() {
        "initialize" => handle_initialize(id),
        "initialized" => JsonRpcResponse::success(id, json!({})),
        "shutdown" => JsonRpcResponse::success(id, json!(null)),
        "tools/list" => handle_tools_list(id),
        "tools/call" => handle_tools_call(id, request.params),
        _ => JsonRpcResponse::error(id, -32601, format!("Method not found: {}", request.method)),
    })
}

fn handle_initialize(id: Value) -> JsonRpcResponse {
    JsonRpcResponse::success(
        id,
        json!({
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": {}
            },
            "serverInfo": {
                "name": SERVER_NAME,
                "version": SERVER_VERSION
            }
        }),
    )
}

fn handle_tools_list(id: Value) -> JsonRpcResponse {
    JsonRpcResponse::success(
        id,
        json!({
            "tools": [
                {
                    "name": "peek",
                    "description": "Analyze source code structure using tree-sitter. Extracts imports, structs/classes, functions, traits/interfaces, and enums from source files.",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "file": {
                                "type": "string",
                                "description": "Path to the source file to analyze"
                            },
                            "exports_only": {
                                "type": "boolean",
                                "description": "Show only public/exported items",
                                "default": false
                            }
                        },
                        "required": ["file"]
                    }
                }
            ]
        }),
    )
}

fn handle_tools_call(id: Value, params: Value) -> JsonRpcResponse {
    let tool_name = params.get("name").and_then(|v| v.as_str()).unwrap_or("");
    let arguments = params.get("arguments").cloned().unwrap_or(json!({}));

    match tool_name {
        "peek" => handle_peek_tool(id, arguments),
        _ => JsonRpcResponse::error(id, -32602, format!("Unknown tool: {}", tool_name)),
    }
}

fn handle_peek_tool(id: Value, arguments: Value) -> JsonRpcResponse {
    let file_path = match arguments.get("file").and_then(|v| v.as_str()) {
        Some(path) => path,
        None => {
            return JsonRpcResponse::error(
                id,
                -32602,
                "Missing required parameter: file".to_string(),
            );
        }
    };

    let exports_only = arguments
        .get("exports_only")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    // Read the file
    let source = match std::fs::read_to_string(file_path) {
        Ok(content) => content,
        Err(e) => {
            return JsonRpcResponse::error(
                id,
                -32000,
                format!("Failed to read file '{}': {}", file_path, e),
            );
        }
    };

    // Get the file extension
    let extension = std::path::Path::new(file_path)
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("");

    // Parse the file
    let structure = match parser::parse_file(&source, extension) {
        Ok(s) => s,
        Err(e) => {
            return JsonRpcResponse::error(id, -32000, format!("Failed to parse file: {}", e));
        }
    };

    let filtered = if exports_only {
        structure.exports_only()
    } else {
        structure
    };

    // Convert to JSON
    let result_json = match serde_json::to_value(&filtered) {
        Ok(v) => v,
        Err(e) => {
            return JsonRpcResponse::error(
                id,
                -32000,
                format!("Failed to serialize result: {}", e),
            );
        }
    };

    JsonRpcResponse::success(
        id,
        json!({
            "content": [
                {
                    "type": "text",
                    "text": serde_json::to_string_pretty(&result_json).unwrap_or_default()
                }
            ]
        }),
    )
}