from typing import Callable

from mcp.server.fastmcp import FastMCP

from atopile.logging import get_logger

logger = get_logger(__name__)


MCP_DECORATOR = Callable[[FastMCP], Callable]


class MCPTools:
    def __init__(self):
        self._tools: dict[Callable, MCP_DECORATOR] = {}

    def register(self, decorator: MCP_DECORATOR = lambda mcp: mcp.tool()):
        def decorator_wrapper(func: Callable):
            self._tools[func] = decorator

            return func

        return decorator_wrapper

    def install(self, mcp: FastMCP):
        for func, decorator in self._tools.items():
            d = decorator(mcp)
            d(func)
