"""MCP client — Streamable HTTP transport.

An McpSession is opened once per prompt run and kept alive for the whole
tool-calling loop (list tools, then call tools repeatedly), mirroring how a
real MCP client (e.g. Claude Code) holds one session per conversation.
"""

import json
from typing import Any

import httpx

_PROTOCOL_VERSION = "2025-06-18"
_CLIENT_INFO = {"name": "ows-ai-eval-runner", "version": "1.0.0"}
_TIMEOUT_SECONDS = 30.0


def _headers(token: str, session_id: str | None) -> dict[str, str]:
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json, text/event-stream",
        "Authorization": f"Bearer {token}",
    }
    if session_id:
        headers["Mcp-Session-Id"] = session_id
    return headers


def _parse_body(response: httpx.Response) -> dict[str, Any]:
    # Streamable HTTP transport may reply as plain JSON or as a single SSE event.
    if "text/event-stream" in response.headers.get("content-type", ""):
        for line in response.text.splitlines():
            if line.startswith("data:"):
                return json.loads(line[len("data:") :].strip())
        raise ValueError("MCP response was SSE but contained no data event")
    return response.json()


class McpError(Exception):
    """Raised when the target MCP returns a JSON-RPC error."""


class McpSession:
    """One MCP client session: initialize once, then list/call tools any number of times."""

    def __init__(self, endpoint: str, token: str):
        self._endpoint = endpoint
        self._token = token
        self._client = httpx.Client(timeout=_TIMEOUT_SECONDS)
        self._session_id: str | None = None
        self._next_id = 1

    def __enter__(self) -> "McpSession":
        self._session_id = self._initialize()
        return self

    def __exit__(self, *exc_info: object) -> None:
        self._client.close()

    def _rpc_id(self) -> int:
        request_id = self._next_id
        self._next_id += 1
        return request_id

    def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
        response = self._client.post(
            self._endpoint,
            headers=_headers(self._token, self._session_id),
            json=payload,
        )
        response.raise_for_status()
        return _parse_body(response)

    def _initialize(self) -> str | None:
        response = self._client.post(
            self._endpoint,
            headers=_headers(self._token, session_id=None),
            json={
                "jsonrpc": "2.0",
                "id": self._rpc_id(),
                "method": "initialize",
                "params": {
                    "protocolVersion": _PROTOCOL_VERSION,
                    "capabilities": {},
                    "clientInfo": _CLIENT_INFO,
                },
            },
        )
        response.raise_for_status()
        session_id = response.headers.get("mcp-session-id")

        response = self._client.post(
            self._endpoint,
            headers=_headers(self._token, session_id),
            json={"jsonrpc": "2.0", "method": "notifications/initialized"},
        )
        response.raise_for_status()
        return session_id

    def list_tools(self) -> list[dict[str, Any]]:
        """Return the target MCP's tool manifest as Bedrock-shaped tool specs."""
        body = self._post(
            {
                "jsonrpc": "2.0",
                "id": self._rpc_id(),
                "method": "tools/list",
                "params": {},
            }
        )
        if "error" in body:
            raise McpError(f"MCP tools/list error: {body['error']}")

        return [
            {
                "name": tool["name"],
                "description": tool.get("description", ""),
                "inputSchema": tool.get(
                    "inputSchema", {"type": "object", "properties": {}}
                ),
            }
            for tool in body["result"]["tools"]
        ]

    def call_tool(self, name: str, arguments: dict[str, Any]) -> str:
        """Call a tool and return its result as plain text (concatenated text content blocks)."""
        body = self._post(
            {
                "jsonrpc": "2.0",
                "id": self._rpc_id(),
                "method": "tools/call",
                "params": {"name": name, "arguments": arguments},
            }
        )
        if "error" in body:
            raise McpError(f"MCP tools/call error for '{name}': {body['error']}")

        result = body["result"]
        content = result.get("content", [])
        text = "".join(
            block.get("text", "") for block in content if block.get("type") == "text"
        )
        if result.get("isError"):
            raise McpError(f"Tool '{name}' returned an error: {text}")
        return text