from src.Llms.MCPs.McpClient import MCPClient
from typing import Any

class ClientRegistry:
    def __init__(self):
        self._clients: dict[str, MCPClient] = {}

    def add(self, name: str, wrapper: MCPClient) -> None:
        self._clients[name] = wrapper

    def get(self, name: str) -> MCPClient:
        if name not in self._clients:
            raise KeyError(f"Cliente '{name}' no encontrado en el registro.")
        return self._clients[name]

    def all(self) -> dict[str, MCPClient]:
        return self._clients

    def list_names(self) -> list[str]:
        return list(self._clients.keys())

    def __contains__(self, name: str) -> bool:
        return name in self._clients
    
    async def listar_tools_por_cliente(self) -> dict[str, Any]:
        resultado = {"tools": {}, "errores": {}}
        for name, wrapper in self._clients.items():
            try:
                async with wrapper:
                    resultado["tools"][name] = await wrapper.list_tools()
            except Exception as e:
                resultado["errores"][name] = str(e)
                resultado["tools"][name] = []
        return resultado

    async def listar_prompts_por_cliente(self) -> dict[str, Any]:
        resultado = {"prompts": {}, "errores": {}}
        for name, wrapper in self._clients.items():
            try:
                async with wrapper:
                    resultado["prompts"][name] = await wrapper.list_prompts()
            except Exception as e:
                resultado["errores"][name] = str(e)
                resultado["prompts"][name] = []
        return resultado

    async def listar_resources_por_cliente(self) -> dict[str, Any]:
        resultado = {"resources": {}, "errores": {}}
        for name, wrapper in self._clients.items():
            try:
                async with wrapper:
                    resultado["resources"][name] = await wrapper.list_resources()
            except Exception as e:
                resultado["errores"][name] = str(e)
                resultado["resources"][name] = []
        return resultado