from __future__ import annotations

import asyncio
import threading
from typing import Any

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


class StdioMcpSession:
    """同步封装 MCP stdio ClientSession（后台线程维持单 task 生命周期）。"""

    def __init__(self) -> None:
        self._loop = asyncio.new_event_loop()
        self._thread: threading.Thread | None = None
        self._ready = threading.Event()
        self._stop = threading.Event()
        self._session: ClientSession | None = None

    def connect(self, command: str, args: list[str], *, env: dict[str, str] | None = None) -> None:
        if self._thread is not None:
            return

        def runner() -> None:
            asyncio.set_event_loop(self._loop)
            self._loop.run_until_complete(self._lifecycle(command, args, env))

        self._thread = threading.Thread(target=runner, daemon=True)
        self._thread.start()
        if not self._ready.wait(timeout=30):
            raise TimeoutError("MCP stdio connect timeout")

    async def _lifecycle(self, command: str, args: list[str], env: dict[str, str] | None) -> None:
        params = StdioServerParameters(command=command, args=args, env=env)
        async with stdio_client(params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                self._session = session
                self._ready.set()
                while not self._stop.is_set():
                    await asyncio.sleep(0.05)

    def _run(self, coro: Any) -> Any:
        if self._loop.is_closed():
            raise RuntimeError("MCP stdio session closed")
        return asyncio.run_coroutine_threadsafe(coro, self._loop).result(timeout=30)

    def list_tools(self) -> list[str]:
        async def _list() -> list[str]:
            assert self._session is not None
            res = await self._session.list_tools()
            return [t.name for t in res.tools]

        return self._run(_list())

    def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
        async def _call() -> dict[str, Any]:
            assert self._session is not None
            res = await self._session.call_tool(name, arguments)
            parts = [c.text for c in res.content if hasattr(c, "text") and c.text]
            return {"ok": True, "content": "".join(parts), "raw": res.model_dump()}

        return self._run(_call())

    def close(self) -> None:
        if self._thread is None:
            return
        self._stop.set()
        self._thread.join(timeout=10)
        if not self._loop.is_closed():
            self._loop.close()
        self._session = None
        self._thread = None
