import asyncio
import uuid
import json
from typing import Any, Dict, Union, List
from starlette.responses import Response
from tools.utils import ToolError

class MCPServer:
    def __init__(self):
        self.sessions: Dict[str, asyncio.Queue] = {}
        self.running = True

    def create_session(self) -> str:
        sid = str(uuid.uuid4())
        self.sessions[sid] = asyncio.Queue()
        return sid

    async def send_to_session(self, sid: str, message: Dict[str, Any]):
        if sid in self.sessions:
            await self.sessions[sid].put(json.dumps(message))

    async def handle_request(self, method: str, params: Dict[str, Any], req_id: Any, tool_registry: list) -> Union[Dict[str, Any], Response]:
        if method == "initialize":
            return {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": "MooseCP", "version": "0.4.1"}
            }
        
        if method == "tools/list":
            return {"tools": [t.to_mcp_dict() for t in tool_registry]}
        
        if method == "tools/call":
            name = params.get("name")
            args = params.get("arguments", {})
            
            tool = next((t for t in tool_registry if t.name == name), None)
            if not tool:
                return Response(
                    content=json.dumps({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Tool {name} not found"}}),
                    media_type="application/json"
                )
            
            try:
                result_data = await tool.handler(args)
            except ToolError as e:
                # Convert tool errors into a text response so the LLM can understand and potentially fix the input
                result_data = {"text": str(e)}
            except Exception as e:
                # Catch-all for unexpected crashes to prevent server death
                result_data = {"text": f"Unexpected internal error: {str(e)}"}
            
            # Handle results that are already lists of content (for multimodal/multi-part responses)
            if isinstance(result_data, list):
                content = result_data
            elif isinstance(result_data, dict) and "text" in result_data:
                content = [{"type": "text", "text": result_data["text"]}]
            elif isinstance(result_data, dict) and result_data.get("type") == "image":
                content = [result_data]
            else:
                content = [result_data]
                
            return {"content": content}

        return Response(
            content=json.dumps({"jsonrpc": "2.0", "id": req_id, "error": {"code": -32601, "message": f"Method {method} not found"}}),
            media_type="application/json"
        )