//! MCP server-side skeleton: accept JSON-RPC over stdio, dispatch to registered tools.

use async_trait::async_trait;
use serde_json::{json, Value};
use std::collections::HashMap;
use thulp_core::{ToolDefinition, ToolResult};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};

/// Trait implemented by tool handlers that the server dispatches to.
#[async_trait]
pub trait ToolHandler: Send + Sync {
    /// Execute the tool with the given JSON arguments.
    async fn call(&self, arguments: Value) -> ToolResult;
}

/// MCP server that accepts JSON-RPC requests and dispatches to registered tools.
pub struct McpServer {
    name: String,
    version: String,
    tools: HashMap<String, (ToolDefinition, Box<dyn ToolHandler>)>,
}

/// Builder for [`McpServer`].
pub struct McpServerBuilder {
    name: String,
    version: String,
    tools: HashMap<String, (ToolDefinition, Box<dyn ToolHandler>)>,
}

impl McpServerBuilder {
    fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            version: version.into(),
            tools: HashMap::new(),
        }
    }

    /// Register a tool with its definition and handler.
    pub fn tool(
        mut self,
        name: impl Into<String>,
        definition: ToolDefinition,
        handler: Box<dyn ToolHandler>,
    ) -> Self {
        self.tools.insert(name.into(), (definition, handler));
        self
    }

    /// Build the server.
    pub fn build(self) -> McpServer {
        McpServer {
            name: self.name,
            version: self.version,
            tools: self.tools,
        }
    }
}

impl McpServer {
    /// Create a builder for an MCP server.
    pub fn builder(name: impl Into<String>, version: impl Into<String>) -> McpServerBuilder {
        McpServerBuilder::new(name, version)
    }

    /// Handle a single JSON-RPC request and return a JSON-RPC response.
    pub async fn handle_request(&self, request: Value) -> Value {
        let id = request.get("id").cloned().unwrap_or(Value::Null);
        let method = request
            .get("method")
            .and_then(|m| m.as_str())
            .unwrap_or("");

        match method {
            "initialize" => self.handle_initialize(id),
            "tools/list" => self.handle_tools_list(id),
            "tools/call" => self.handle_tools_call(id, &request).await,
            _ => json_rpc_error(id, -32601, &format!("Method not found: {method}")),
        }
    }

    fn handle_initialize(&self, id: Value) -> Value {
        json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": self.name,
                    "version": self.version
                }
            }
        })
    }

    fn handle_tools_list(&self, id: Value) -> Value {
        let tools: Vec<Value> = self
            .tools
            .values()
            .map(|(def, _)| {
                json!({
                    "name": def.name,
                    "description": def.description,
                    "inputSchema": {
                        "type": "object",
                        "properties": {},
                    }
                })
            })
            .collect();

        json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": { "tools": tools }
        })
    }

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

        match self.tools.get(tool_name) {
            Some((_, handler)) => {
                let result = handler.call(arguments).await;
                let content = if result.success {
                    let text = result
                        .data
                        .map(|d| d.to_string())
                        .unwrap_or_default();
                    vec![json!({"type": "text", "text": text})]
                } else {
                    let text = result.error.unwrap_or_else(|| "unknown error".into());
                    vec![json!({"type": "text", "text": text})]
                };
                json!({
                    "jsonrpc": "2.0",
                    "id": id,
                    "result": {
                        "content": content,
                        "isError": !result.success
                    }
                })
            }
            None => json_rpc_error(id, -32602, &format!("Unknown tool: {tool_name}")),
        }
    }

    /// Serve MCP over stdio. Reads newline-delimited JSON-RPC from stdin,
    /// dispatches each request, writes responses to stdout.
    pub async fn serve_stdio(&self) -> std::io::Result<()> {
        let stdin = tokio::io::stdin();
        let mut stdout = tokio::io::stdout();
        let mut reader = BufReader::new(stdin);
        let mut line = String::new();

        loop {
            line.clear();
            let n = reader.read_line(&mut line).await?;
            if n == 0 {
                break; // EOF
            }
            let trimmed = line.trim();
            if trimmed.is_empty() {
                continue;
            }
            let request: Value = match serde_json::from_str(trimmed) {
                Ok(v) => v,
                Err(_) => {
                    let err = json_rpc_error(Value::Null, -32700, "Parse error");
                    let mut out = serde_json::to_string(&err).unwrap();
                    out.push('\n');
                    stdout.write_all(out.as_bytes()).await?;
                    stdout.flush().await?;
                    continue;
                }
            };

            let response = self.handle_request(request).await;
            let mut out = serde_json::to_string(&response).unwrap();
            out.push('\n');
            stdout.write_all(out.as_bytes()).await?;
            stdout.flush().await?;
        }
        Ok(())
    }
}

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

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

    struct EchoHandler;

    #[async_trait]
    impl ToolHandler for EchoHandler {
        async fn call(&self, arguments: Value) -> ToolResult {
            ToolResult::success(arguments)
        }
    }

    fn test_server() -> McpServer {
        let def = ToolDefinition::new("echo");
        McpServer::builder("test-server", "0.1.0")
            .tool("echo", def, Box::new(EchoHandler))
            .build()
    }

    #[tokio::test]
    async fn test_initialize() {
        let server = test_server();
        let req = json!({"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}});
        let resp = server.handle_request(req).await;
        assert_eq!(resp["id"], 1);
        assert_eq!(resp["result"]["serverInfo"]["name"], "test-server");
        assert!(resp["result"]["capabilities"]["tools"].is_object());
    }

    #[tokio::test]
    async fn test_tools_list() {
        let server = test_server();
        let req = json!({"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}});
        let resp = server.handle_request(req).await;
        let tools = resp["result"]["tools"].as_array().unwrap();
        assert_eq!(tools.len(), 1);
        assert_eq!(tools[0]["name"], "echo");
    }

    #[tokio::test]
    async fn test_tools_call() {
        let server = test_server();
        let req = json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "echo",
                "arguments": {"msg": "hello"}
            }
        });
        let resp = server.handle_request(req).await;
        assert_eq!(resp["id"], 3);
        assert_eq!(resp["result"]["isError"], false);
        let content = resp["result"]["content"].as_array().unwrap();
        assert!(content[0]["text"].as_str().unwrap().contains("hello"));
    }

    #[tokio::test]
    async fn test_unknown_method() {
        let server = test_server();
        let req = json!({"jsonrpc": "2.0", "id": 4, "method": "unknown/method", "params": {}});
        let resp = server.handle_request(req).await;
        assert!(resp.get("error").is_some());
        assert_eq!(resp["error"]["code"], -32601);
    }

    #[tokio::test]
    async fn test_unknown_tool() {
        let server = test_server();
        let req = json!({
            "jsonrpc": "2.0",
            "id": 5,
            "method": "tools/call",
            "params": {"name": "nonexistent", "arguments": {}}
        });
        let resp = server.handle_request(req).await;
        assert!(resp.get("error").is_some());
        assert_eq!(resp["error"]["code"], -32602);
    }
}
