"""MCP stdio server for Claude Desktop."""

from typing import Any, Callable
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from pydantic import BaseModel


class ToolRegistry:
    """Manages tool registration and execution."""

    def __init__(self):
        self._tools: dict[str, dict[str, Any]] = {}

    def tool(self, name: str, description: str, schema: type[BaseModel] | None = None):
        def decorator(func: Callable) -> Callable:
            self._tools[name] = {
                "description": description,
                "schema": schema,
                "handler": func,
            }
            return func

        return decorator

    def get_tool_definitions(self) -> list[Tool]:
        tools = []
        for name, config in self._tools.items():
            input_schema = (
                config["schema"].model_json_schema() if config["schema"] else {}
            )
            tools.append(
                Tool(
                    name=name,
                    description=config["description"],
                    inputSchema=input_schema,
                )
            )
        return tools

    async def execute(self, name: str, arguments: dict[str, Any]) -> list[TextContent]:
        if name not in self._tools:
            raise ValueError(f"Tool '{name}' not found")
        config = self._tools[name]
        args = config["schema"](**arguments) if config["schema"] else arguments
        result = await config["handler"](args)
        return [TextContent(type="text", text=str(result))]


class MCPServer:
    """MCP server with tool registry."""

    def __init__(self, name: str, version: str):
        self.server = Server(name)
        self.registry = ToolRegistry()
        self.name = name
        self.version = version
        self._setup_handlers()

    def _setup_handlers(self):
        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            return self.registry.get_tool_definitions()

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
            return await self.registry.execute(name, arguments)

    def tool(self, name: str, description: str, schema: type[BaseModel] | None = None):
        return self.registry.tool(name, description, schema)

    async def run(self):
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream, write_stream, self.server.create_initialization_options()
            )
