from __future__ import annotations
from typing import Any, Callable
from .registry import get_registry
from .exports import McpServerExport

class AielMcpServer:
    def __init__(self, name: str):
        self.name = name
        reg = get_registry()
        reg.mcp_servers.setdefault(name, McpServerExport(name=name, tools=[]))

    def tool(self, name: str | None = None):
        def decorator(fn: Callable[..., Any]):
            reg = get_registry()
            tool_name = name or fn.__name__

            srv = reg.mcp_servers.setdefault(self.name, McpServerExport(name=self.name, tools=[]))
            if tool_name not in srv.tools:
                srv.tools.append(tool_name)

            reg.mcp_tool_fns[(self.name, tool_name)] = fn
            return fn
        return decorator

def mcp_server(name: str, tools: list[str] | None = None) -> AielMcpServer:
    reg = get_registry()
    srv = reg.mcp_servers.setdefault(name, McpServerExport(name=name, tools=[]))
    if tools:
        for t in tools:
            if t not in srv.tools:
                srv.tools.append(t)
    return AielMcpServer(name)
