"""Persistent MCP client for the local paper retrieval server."""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any

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

PROJECT_ROOT = Path(__file__).resolve().parents[1]


class PaperMCPClient:
    """Keeps one MCP session open for retrieval tool calls."""

    def __init__(self, project_root: Path | None = None):
        self.project_root = project_root or PROJECT_ROOT
        self._stdio_ctx = None
        self._session_ctx = None
        self._session: ClientSession | None = None

    async def connect(self) -> None:
        if self._session is not None:
            return

        server_script = os.getenv("MCP_SERVER_SCRIPT", "waterscope_retrieval_server.py")
        server_params = StdioServerParameters(
            command=sys.executable,
            args=[str(self.project_root / server_script)],
            cwd=str(self.project_root),
            env={**os.environ},
        )

        self._stdio_ctx = stdio_client(server_params)
        read, write = await self._stdio_ctx.__aenter__()
        self._session_ctx = ClientSession(read, write)
        self._session = await self._session_ctx.__aenter__()
        await self._session.initialize()

    async def close(self) -> None:
        if self._session_ctx is not None:
            await self._session_ctx.__aexit__(None, None, None)
            self._session_ctx = None
            self._session = None
        if self._stdio_ctx is not None:
            await self._stdio_ctx.__aexit__(None, None, None)
            self._stdio_ctx = None

    async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> Any:
        if self._session is None:
            raise RuntimeError("MCP client is not connected.")

        payload = {key: value for key, value in (arguments or {}).items() if value is not None}
        result = await self._session.call_tool(name, payload)

        if result.isError:
            message = result.content[0].text if result.content else "Unknown MCP tool error"
            raise RuntimeError(message)

        if not result.content:
            return {}

        text = result.content[0].text
        return json.loads(text)