"""Core MCP server implementation with tool registry."""

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 with automatic schema generation."""
    
    def __init__(self):
        self._tools: dict[str, dict[str, Any]] = {}
    
    def tool(
        self,
        name: str,
        description: str,
        schema: type[BaseModel] | None = None
    ):
        """Decorator to register a tool with automatic schema validation.
        
        Usage:
            @registry.tool("tool_name", "Description", InputSchema)
            async def my_tool(args: InputSchema) -> str:
                return "result"
        """
        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]:
        """Generate MCP tool definitions from registered tools."""
        tools = []
        for name, config in self._tools.items():
            input_schema = {}
            if config["schema"]:
                input_schema = config["schema"].model_json_schema()
            
            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]:
        """Execute a registered tool with validation."""
        if name not in self._tools:
            raise ValueError(f"Tool '{name}' not found")
        
        config = self._tools[name]
        
        # Validate and parse arguments if schema exists
        if config["schema"]:
            args = config["schema"](**arguments)
        else:
            args = arguments
        
        # Execute handler
        result = await config["handler"](args)
        
        return [TextContent(type="text", text=str(result))]


class MCPServer:
    """Main 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):
        """Configure MCP protocol handlers."""
        
        @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):
        """Register a tool (delegates to registry)."""
        return self.registry.tool(name, description, schema)
    
    async def run(self):
        """Start the MCP server."""
        async with stdio_server() as (read_stream, write_stream):
            await self.server.run(
                read_stream,
                write_stream,
                self.server.create_initialization_options()
            )
