from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client


class MCPClient:
    def __init__(self, url: str):
        self.url = url
        self.session = None
        self._client = None

    async def connect(self):
        self._client = streamable_http_client(self.url)
        read, write = await self._client.__aenter__()

        self.session = ClientSession(read, write)
        await self.session.initialize()

    async def list_tools(self):
        if not self.session:
            await self.connect()
        return await self.session.list_tools()

    async def call_tool(self, name: str, arguments: dict):
        if not self.session:
            await self.connect()
        return await self.session.call_tool(name, arguments)

    async def close(self):
        if self._client:
            await self._client.__aexit__(None, None, None)