from typing import List, Dict
import asyncio

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client
from contextlib import AsyncExitStack

async def connect_to_stdio_server() -> None:
    """Connect to a single MCP server."""

    sessions: List[ClientSession] = []
    exit_stack = AsyncExitStack()

    try:
        server_params = StdioServerParameters(
                command="npx",
                args=[
                    "mcp-remote",
                    "https://jarisko-mcp-server-deployment-sample.hf.space/gradio_api/mcp/sse",
                    "--transport",
                    "sse-only"
                ]
        )

        stdio_transport = await exit_stack.enter_async_context(
            stdio_client(server_params)
        )

        read, write = stdio_transport
        session = await exit_stack.enter_async_context(
            ClientSession(read, write)
        )

        await session.initialize()
        sessions.append(session)

        print("3")

        # List available tools for this session
        response = await session.list_tools()
        tools = response.tools
        print(f"{tools=}")

    except Exception as e:
        print(f"Failed to connect: {e}")


async def connect_to_sse_server():

    async with sse_client("https://jarisko-mcp-server-deployment-sample.hf.space/gradio_api/mcp/sse") as streams:
        async with ClientSession(*streams) as session:
            await session.initialize()
            tools_result = await session.list_tools()
            print(f"{tools_result=}")


asyncio.run(connect_to_sse_server())