# jarvis_mcp/client.py
import asyncio
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from util.logger import get_logger

logger = get_logger("MCP")


class MCPClient:
    def __init__(self, url: str):
        self.url = url.rstrip('/') + "/mcp/"   # WICHTIG: /mcp Endpunkt
        self.session = None
        self._client_context = None

    async def connect(self):
        if self.session is not None:
            return

        logger.info(f"Verbinde mit OpenHAB MCP: {self.url}")

        try:
            self._client_context = streamablehttp_client(
                self.url,
                terminate_on_close=True
            )

            streams = await self._client_context.__aenter__()

            # Manche Versionen geben 2, manche 3 Werte zurück
            if isinstance(streams, tuple):
                if len(streams) == 3:
                    read, write, _ = streams
                else:
                    read, write = streams
            else:
                raise ValueError("Unerwartetes Format von streamablehttp_client")

            self.session = ClientSession(read, write)
            
            # Mit Timeout
            await asyncio.wait_for(self.session.initialize(), timeout=8.0)

            logger.info("✅ OpenHAB MCP Verbindung erfolgreich")

        except asyncio.TimeoutError:
            logger.error("❌ Timeout beim MCP Connect")
            raise
        except Exception as e:
            logger.error(f"❌ MCP Connect Fehler: {e}")
            raise

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

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

    async def close(self):
        """Sauberes Aufräumen"""
        closed = False
        if self.session:
            try:
                if hasattr(self.session, "close"):
                    await self.session.close()
                closed = True
            except Exception as e:
                logger.warning(f"Session close Warning: {e}")

        if self._client_context:
            try:
                await self._client_context.__aexit__(None, None, None)
            except Exception as e:
                logger.warning(f"Context close Warning: {e}")

        self.session = None
        self._client_context = None
        if closed:
            logger.info("MCP Verbindung geschlossen")