from __future__ import annotations

import os
import subprocess

from agents.mcp import MCPServerSse, MCPServerStdio
from mcp.client.stdio import stdio_client

from .config import Settings


class SessionMCPServerSse(MCPServerSse):
    """SSE MCP server that stays connected for the whole app session."""

    async def cleanup(self) -> None:
        # Runner triggers cleanup around each run; suppress to keep login session alive.
        return

    async def __aexit__(self, *args) -> None:
        return

    async def close_session(self) -> None:
        await super().cleanup()


class SessionMCPServerStdio(MCPServerStdio):
    """Stdio MCP server that keeps subprocess alive between chat turns."""

    def create_streams(self):
        return stdio_client(self.params, errlog=subprocess.DEVNULL)

    async def cleanup(self) -> None:
        return

    async def __aexit__(self, *args) -> None:
        return

    async def close_session(self) -> None:
        await super().cleanup()


class MCPRuntime:
    """Holds and manages all MCP connections used by OpenKite."""

    def __init__(self, settings: Settings) -> None:
        self.kite_server = SessionMCPServerSse(
            name="kite",
            params={"url": settings.kite_sse_url},
            cache_tools_list=True,
            client_session_timeout_seconds=30,
        )
        self.firecrawl_server = SessionMCPServerStdio(
            name="firecrawl",
            params={
                "command": "npx",
                "args": ["-y", "firecrawl-mcp"],
                "env": {
                    **os.environ.copy(),
                    "FIRECRAWL_API_KEY": settings.firecrawl_api_key,
                },
            },
            cache_tools_list=True,
            client_session_timeout_seconds=30,
        )
        self._connected = False

    @property
    def servers(self) -> list:
        return [self.kite_server, self.firecrawl_server]

    async def connect(self) -> None:
        if self._connected:
            return
        await self.kite_server.connect()
        await self.firecrawl_server.connect()
        self._connected = True

    async def close(self) -> None:
        if not self._connected:
            return
        await self.firecrawl_server.close_session()
        await self.kite_server.close_session()
        self._connected = False
