import asyncio
import os
from typing import Optional, List, Dict
from contextlib import AsyncExitStack
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.types import Tool

class MCPStdioServer:
    def __init__(
        self,
        name: str,
        command: str,
        args: List[str],
        env: Optional[Dict[str, str]] = None
    ):
        self.name = name
        self.command = command
        self.args = args
        self.env = env or os.environ.copy()
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None
        self.tools: List[Tool] = []

    async def start(self):
        # Configurar el bucle de eventos Proactor en Windows si es necesario
        if os.name == "nt":
            asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

        server_params = StdioServerParameters(
            command=self.command,
            args=self.args,
            env=self.env
        )

        # Iniciar el transporte y establecer la sesión
        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        read, write = stdio_transport

        self.session = await self.exit_stack.enter_async_context(ClientSession(read, write))
        await self.session.initialize()
        response = await self.session.list_tools()
        self.tools = response.tools

        if self.tools:
            print(f"[{self.name}] Servidor iniciado con herramientas:")
            for tool in self.tools:
                nombre = getattr(tool, "name", "[sin nombre]")
                descripcion = getattr(tool, "description", "[sin descripción]")
                print(f"  - {nombre} - {descripcion}")
        else:
            print(f"[{self.name}] Servidor iniciado, pero no se detectaron herramientas.")

    async def call_tool(self, tool_name: str, arguments: Dict):
        if not self.session:
            raise RuntimeError("La sesión no está inicializada.")
        result = await self.session.call_tool(tool_name, arguments)
        return result.content

    def get_tool_names(self) -> List[str]:
        return [tool.name for tool in self.tools]

    async def stop(self):
        await self.exit_stack.aclose()
        print(f"[{self.name}] Servidor detenido.")