"""Minimal MCP client over the Streamable HTTP transport (JSON-RPC 2.0).

Supports the handshake plus ``tools/list`` and ``tools/call``. Each high-level
call opens a fresh session (initialize -> notifications/initialized -> op) which
keeps the client stateless and robust against session expiry. Responses may come
back as either ``application/json`` or an SSE (``text/event-stream``) frame; both
are handled. The HTTP client comes from the shared connector runtime so tests can
inject a ``MockTransport`` and exercise the real code with no network.
"""

from __future__ import annotations

import contextlib
import json

import httpx

from legate.connectors.runtime import new_async_client
from legate.security.ssrf import SSRFError, validate_public_url

_PROTOCOL_VERSION = "2025-06-18"
_CLIENT_INFO = {"name": "legate-agent", "version": "0.1"}


class MCPError(RuntimeError):
    """Raised when an MCP request fails or returns an error."""


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


def _extract_message(resp: httpx.Response) -> dict:
    """Return the JSON-RPC message from a JSON or SSE response."""
    content_type = resp.headers.get("content-type", "")
    if "text/event-stream" in content_type:
        for block in resp.text.split("\n\n"):
            data = "".join(
                line[len("data:") :].strip()
                for line in block.splitlines()
                if line.startswith("data:")
            )
            if not data:
                continue
            try:
                obj = json.loads(data)
            except json.JSONDecodeError:
                continue
            if isinstance(obj, dict) and ("result" in obj or "error" in obj):
                return obj
        raise MCPError("no JSON-RPC message found in SSE response")
    try:
        return resp.json()
    except ValueError as exc:
        raise MCPError(f"invalid JSON response: {exc}") from exc


def _result(resp: httpx.Response, what: str) -> dict:
    if resp.status_code >= 400:
        raise MCPError(f"{what} failed: HTTP {resp.status_code}")
    message = _extract_message(resp)
    if isinstance(message.get("error"), dict):
        raise MCPError(message["error"].get("message", f"{what} returned an error"))
    result = message.get("result")
    return result if isinstance(result, dict) else {}


async def _open_session(client: httpx.AsyncClient, url: str, token: str | None) -> str | None:
    init = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "initialize",
        "params": {
            "protocolVersion": _PROTOCOL_VERSION,
            "capabilities": {},
            "clientInfo": _CLIENT_INFO,
        },
    }
    try:
        resp = await client.post(url, json=init, headers=_headers(token))
    except httpx.HTTPError as exc:
        raise MCPError(f"connection failed: {exc}") from exc
    _result(resp, "initialize")
    session_id = resp.headers.get("mcp-session-id")
    # A well-behaved client must confirm initialization before other requests.
    # The notification is best-effort; ignore transport hiccups here.
    with contextlib.suppress(httpx.HTTPError):
        await client.post(
            url,
            json={"jsonrpc": "2.0", "method": "notifications/initialized"},
            headers=_headers(token, session_id),
        )
    return session_id


def _guard(url: str) -> str:
    try:
        return validate_public_url(url)
    except SSRFError as exc:
        raise MCPError(str(exc)) from exc


async def list_tools(url: str, token: str | None = None) -> list[dict]:
    """Discover the tools a server offers. Returns raw MCP tool definitions."""
    url = _guard(url)
    async with new_async_client() as client:
        session = await _open_session(client, url, token)
        try:
            resp = await client.post(
                url,
                json={"jsonrpc": "2.0", "id": 2, "method": "tools/list"},
                headers=_headers(token, session),
            )
        except httpx.HTTPError as exc:
            raise MCPError(f"tools/list failed: {exc}") from exc
        result = _result(resp, "tools/list")
        tools = result.get("tools", [])
        return tools if isinstance(tools, list) else []


async def call_tool(url: str, token: str | None, name: str, arguments: dict) -> dict:
    """Invoke a tool on the server and return the raw MCP result."""
    url = _guard(url)
    async with new_async_client() as client:
        session = await _open_session(client, url, token)
        try:
            resp = await client.post(
                url,
                json={
                    "jsonrpc": "2.0",
                    "id": 3,
                    "method": "tools/call",
                    "params": {"name": name, "arguments": arguments or {}},
                },
                headers=_headers(token, session),
            )
        except httpx.HTTPError as exc:
            raise MCPError(f"tools/call failed: {exc}") from exc
        return _result(resp, "tools/call")


def content_text(result: dict) -> str:
    """Flatten an MCP tool result's content blocks into a single string."""
    parts = [
        block.get("text", "")
        for block in result.get("content", [])
        if isinstance(block, dict) and block.get("type") == "text"
    ]
    return "\n".join(p for p in parts if p)
