import json
import sys
from typing import Dict, Any, Optional, List
from .handlers import ToolHandlers


class MCPServer:
    def __init__(self, handlers: ToolHandlers):
        self.handlers = handlers

    def sendResponse(self, result: Any = None, error: Any = None, id: Optional[Any] = None) -> None:
        """Send formatted JSON-RPC 2.0 response to stdout."""
        response = {"jsonrpc": "2.0", "id": id}
        if error:
            response["error"] = error
        else:
            response["result"] = result
        sys.stdout.write(json.dumps(response) + "\n")
        sys.stdout.flush()

    def sendError(self, code: int, message: str, id: Optional[Any] = None) -> None:
        """Helper to send JSON-RPC errors."""
        self.sendResponse(error={"code": code, "message": message}, id=id)

    def run(self) -> None:
        """Main loop reading JSON-RPC messages from stdin."""
        for line in sys.stdin:
            if not line.strip():
                continue
            try:
                request = json.loads(line)
            except json.JSONDecodeError:
                self.sendError(-32700, "Parse error")
                continue

            # Standard JSON-RPC validation
            if not isinstance(request, dict) or request.get("jsonrpc") != "2.0":
                self.sendError(-32600, "Invalid Request", id=request.get("id"))
                continue

            method = request.get("method")
            params = request.get("params", {})
            req_id = request.get("id")

            try:
                self.handleRequest(method, params, req_id)
            except Exception as e:
                # Catch-all for tool execution errors or internal issues
                self.sendError(-32603, f"Internal error: {str(e)}", id=req_id)

    def handleRequest(self, method: str, params: Dict[str, Any], req_id: Any) -> None:
        """Dispatcher for MCP methods."""
        if method == "initialize":
            # Respond with server capabilities
            self.sendResponse(
                result={
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {}
                    },
                    "serverInfo": {
                        "name": "trilium-mcp-server",
                        "version": "0.1.0"
                    }
                },
                id=req_id
            )
        elif method == "notifications/initialized":
            # Ignore acknowledgment for now
            pass
        elif method == "tools/list":
            # List all tools and their schemas
            tools = self.handlers.listTools()
            self.sendResponse(result={"tools": tools}, id=req_id)
        elif method == "tools/call":
            # Call a specific tool with params
            tool_name = params.get("name")
            tool_params = params.get("arguments", {})
            if not tool_name:
                self.sendError(-32602, "Invalid params: Missing tool name", id=req_id)
                return
            
            result = self.handlers.callTool(tool_name, tool_params)
            
            # Format the result correctly for MCP
            mcp_result = self.formatToolResult(result)
            self.sendResponse(result={"content": mcp_result}, id=req_id)
        else:
            self.sendError(-32601, f"Method not found: {method}", id=req_id)

    def formatToolResult(self, result: Any) -> List[Dict[str, str]]:
        """Format arbitrary result into MCP TextContent list."""
        if isinstance(result, (dict, list)):
            text = json.dumps(result, indent=2)
        else:
            text = str(result)
        
        return [{"type": "text", "text": text}]