# tools/client.py
import asyncio
from contextlib import AsyncExitStack
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
import logging
from typing import List, Dict, Any

# Настройка логирования
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("mcp-client")


class MCPToolManager:
    """
    Безопасный клиент для управления инструментами MCP.
    Реализует Explicit Allowlist для защиты от инъекций возможностей.
    """

    def __init__(self):
        self.exit_stack = AsyncExitStack()
        self.sessions = {}

        # --- EXPLICIT ALLOWLIST ---
        # Только эти инструменты будут доступны агентам.
        # Любые другие инструменты, предлагаемые серверами, будут игнорироваться.
        self.ALLOWED_TOOLS = {
            "search_plots",  # из vectordb
            "get_movie_details",  # из tmdb
            # "get_credits"       # можно раскомментировать, если нужно
        }

    async def connect_to_server(self, server_name: str, script_path: str):
        """Подключается к локальному Python-скрипту как к MCP серверу"""
        # Используем python для запуска скрипта сервера
        server_params = StdioServerParameters(
            command="python",
            args=[script_path],
            env=None
        )

        try:
            # Инициализация транспорта и сессии
            stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
            read, write = stdio_transport
            session = await self.exit_stack.enter_async_context(ClientSession(read, write))
            await session.initialize()

            self.sessions[server_name] = session
            logger.info(f"✅ Connected to MCP server: {server_name}")

            # Валидация и логирование доступных инструментов
            await self._validate_tools(server_name, session)

        except Exception as e:
            logger.error(f"❌ Failed to connect to {server_name} ({script_path}): {e}")

    async def _validate_tools(self, server_name: str, session: ClientSession):
        """Проверяет инструменты сервера на соответствие белому списку"""
        tools_response = await session.list_tools()

        allowed_count = 0
        blocked_count = 0

        for tool in tools_response.tools:
            if tool.name in self.ALLOWED_TOOLS:
                allowed_count += 1
            else:
                blocked_count += 1
                logger.warning(f"⚠️ BLOCKED tool '{tool.name}' from {server_name} (not in allowlist)")

        logger.info(f"Server {server_name}: {allowed_count} allowed, {blocked_count} blocked.")

    async def list_tools(self) -> List[Any]:
        """Возвращает список всех РАЗРЕШЕННЫХ инструментов со всех серверов"""
        all_tools = []
        for session in self.sessions.values():
            result = await session.list_tools()
            # Фильтруем на лету
            valid_tools = [t for t in result.tools if t.name in self.ALLOWED_TOOLS]
            all_tools.extend(valid_tools)
        return all_tools

    async def call_tool(self, tool_name: str, arguments: dict) -> Any:
        """Вызов инструмента на нужном сервере"""
        # Дополнительная проверка перед вызовом
        if tool_name not in self.ALLOWED_TOOLS:
            raise SecurityError(f"Security Alert: Attempt to call unauthorized tool '{tool_name}'")

        for name, session in self.sessions.items():
            # Проверяем, есть ли инструмент на этом сервере
            tools = await session.list_tools()
            if any(t.name == tool_name for t in tools.tools):
                logger.info(f"📞 Calling {tool_name} on {name} with args: {arguments}")
                result = await session.call_tool(tool_name, arguments)
                return result

        raise ValueError(f"Tool {tool_name} not found on any connected server")

    async def cleanup(self):
        """Закрытие всех соединений"""
        await self.exit_stack.aclose()