"""Minimal MCP Streamable-HTTP JSON-RPC client.

We do **not** import the official ``mcp`` Python SDK here. The SDK
ships an asyncio client that opens a long-lived bidirectional session
and is overkill for a one-shot CLI invocation: we want to make a
single HTTP request, parse one SSE response, and exit. A 200-line
hand-rolled client is simpler to audit, has no import-time DB stack
side effects, and doesn't require an event loop.

The on-the-wire shape we implement is exactly the contract the
existing test suite (`tests/test_auth/test_mcp_endpoint.py`) asserts:

* ``POST /api/v1/mcp/`` with
  ``Authorization: Bearer helix_pat_...``,
  ``Content-Type: application/json``,
  ``Accept: application/json, text/event-stream``.
* Body: a JSON-RPC 2.0 envelope.
* Response: ``Content-Type: text/event-stream`` with one or more
  ``data:`` lines that, when concatenated, form a JSON-RPC response.

We hand-roll the SSE parser per the spec section 9.2 of
https://html.spec.whatwg.org/multipage/server-sent-events.html — only
the pieces we actually need: ``data`` lines, blank-line dispatch,
``event``/``id``/``retry``/comments ignored.
"""
from __future__ import annotations

import json
import secrets
from dataclasses import dataclass
from typing import Any

import httpx

from helix_cli.config import PAT_PREFIX, CliConfig
from helix_cli.http_client import (
    DEFAULT_TIMEOUT_SECONDS,
    USER_AGENT,
    CliRestError,
)

#: MCP protocol version we declare in ``initialize``. Picked from the
#: same constant the test suite uses; the server responds with its own
#: version and we don't currently negotiate (one-shot calls don't need
#: feature flags).
MCP_PROTOCOL_VERSION = "2025-06-18"


# ── SSE parser ────────────────────────────────────────────────────────


class SseParseError(Exception):
    """Raised when an SSE body cannot be turned into a JSON payload."""


def parse_sse_payload(body: str) -> dict:
    """Extract a single JSON-RPC response from an SSE body.

    Streamable-HTTP always frames a JSON-RPC reply as one logical
    message. The wire format may be:

    * a single ``data: {...}`` line (the common case, what FastMCP
      emits today);
    * multiple ``data:`` lines that, joined with ``\\n``, form the JSON
      payload (per the SSE spec — also valid);
    * an ``event: message`` line followed by a ``data:`` line, which
      we treat the same.

    Lines starting with ``:`` are comments per spec and ignored. The
    spec terminates a message at the first blank line; subsequent
    messages are ignored here because Streamable-HTTP single-shot
    requests only return one.
    """
    if body is None or body == "":
        raise SseParseError("empty SSE body")

    data_lines: list[str] = []
    for raw in body.splitlines():
        # Per SSE spec: a blank line after one or more data lines
        # dispatches the message; we exit on the first dispatch since
        # Streamable-HTTP single-shot only carries one message.
        if raw == "":
            if data_lines:
                break
            continue
        if raw.startswith(":"):
            continue  # comment
        if ":" not in raw:
            field, value = raw, ""
        else:
            field, _, value = raw.partition(":")
            # Spec: strip a single leading space after the colon.
            if value.startswith(" "):
                value = value[1:]
        if field == "data":
            data_lines.append(value)
        # event / id / retry: ignored — single-shot Streamable HTTP
        # never depends on them.
    if not data_lines:
        raise SseParseError("no `data:` line in SSE response")
    payload = "\n".join(data_lines)
    try:
        result = json.loads(payload)
    except json.JSONDecodeError as exc:
        raise SseParseError(
            f"SSE data is not valid JSON: {exc.msg} at offset {exc.pos}"
        ) from exc
    if not isinstance(result, dict):
        raise SseParseError("SSE data did not decode to a JSON object")
    return result


# ── JSON-RPC client ───────────────────────────────────────────────────


@dataclass(frozen=True)
class McpToolResult:
    """Decoded ``tools/call`` result. ``payload`` is the parsed JSON
    that came back from the tool's ``content[0].text`` text-block,
    or ``None`` if the tool returned a non-text content block (e.g. a
    structured ``result`` only). ``raw`` is the full JSON-RPC ``result``
    object so callers can fish out fields the wrapper drops.
    """

    payload: Any
    is_error: bool
    raw: dict


class McpJsonRpcError(CliRestError):
    """Server-side MCP error (matches a JSON-RPC error envelope)."""

    def __init__(self, code: int, message: str):
        super().__init__(f"MCP JSON-RPC error {code}: {message}", exit_code=4, code="JSONRPC")
        self.rpc_code = code
        self.rpc_message = message


class McpClient:
    """One-shot Streamable-HTTP JSON-RPC client over httpx.

    Stateless from the protocol's point of view: every call opens a
    new HTTP connection, sends one request, parses one SSE response,
    and exits. ``stateless_http=True`` on the server side matches.

    The client never auto-retries — failures (network, timeout, 4xx,
    5xx, JSON-RPC error, ``isError=True`` tool body) propagate to the
    caller as :class:`CliRestError` with an exit code. Auto-retry on a
    mutating tool would risk double-submitting an approval / ticket;
    auto-retry on a read would be a quiet way to mask transient state
    that ought to surface to the operator. So: fail once, surface
    clearly, let the human decide.
    """

    def __init__(
        self,
        config: CliConfig,
        *,
        timeout: float = DEFAULT_TIMEOUT_SECONDS,
    ) -> None:
        if not config.token:
            raise CliRestError(
                "no Helix PAT configured. Run "
                "`helix auth token set --token-stdin` or set HELIX_PAT.",
                exit_code=1,
                code="NO_TOKEN",
            )
        if not config.token.startswith(PAT_PREFIX):
            raise CliRestError(
                f"configured token is not a personal access token "
                f"(must start with {PAT_PREFIX}).",
                exit_code=1,
                code="BAD_TOKEN",
            )
        self._config = config
        self._timeout = timeout
        self._next_id = 1

    # ── internals ─────────────────────────────────────────────────────

    def _next_request_id(self) -> str:
        # JSON-RPC ids must be unique within a session; ours are
        # one-shot, but we still want a fresh id per call so logs
        # are easy to disambiguate. ``token_hex`` keeps it printable.
        suffix = secrets.token_hex(4)
        idx = self._next_id
        self._next_id += 1
        return f"helix-cli-{idx}-{suffix}"

    def _post(self, payload: dict) -> dict:
        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
            "Authorization": f"Bearer {self._config.token}",
            "User-Agent": USER_AGENT,
        }
        try:
            with httpx.Client(
                timeout=httpx.Timeout(self._timeout),
                headers=headers,
                trust_env=False,
                follow_redirects=False,
                # Explicit TLS-verify even though httpx's default is True
                # — match the defence-in-depth posture in
                # ``http_client.py:69``. MCP is the PRIMARY PAT-bearing
                # surface; without this, a future httpx default flip
                # would silently downgrade TLS on every MCP call but
                # not on ``health``/``ready``. (R3 security finding
                # 2026-05-28.)
                verify=True,
            ) as client:
                r = client.post(self._config.mcp_url, json=payload)
        except httpx.TimeoutException as exc:
            raise CliRestError(
                f"MCP request timed out after {self._timeout}s",
                exit_code=5,
                code="TIMEOUT",
            ) from exc
        except httpx.HTTPError as exc:
            raise CliRestError(
                f"network error calling MCP: {type(exc).__name__}",
                exit_code=5,
                code="NETWORK",
            ) from exc

        return self._decode_response(r)

    def _decode_response(self, r: httpx.Response) -> dict:
        # Auth / authz failures: REST-style envelope from the PAT
        # middleware. Use exit codes 2 / 3 to keep the contract
        # documented in the task brief.
        if r.status_code == 401:
            raise CliRestError(
                _extract_envelope_message(r) or "unauthorized",
                exit_code=2,
                code="UNAUTHORIZED",
            )
        if r.status_code == 403:
            raise CliRestError(
                _extract_envelope_message(r) or "forbidden",
                exit_code=3,
                code="FORBIDDEN",
            )
        if r.status_code == 413:
            raise CliRestError(
                "request body too large",
                exit_code=4,
                code="PAYLOAD_TOO_LARGE",
            )
        if r.status_code == 429:
            retry_after = r.headers.get("retry-after", "?")
            raise CliRestError(
                f"rate-limited by server (retry after {retry_after}s)",
                exit_code=4,
                code="RATE_LIMITED",
            )
        if r.status_code == 426:
            # Server-side min-version kill-switch: this CLI is too old.
            raise CliRestError(
                _extract_envelope_message(r)
                or "this helix CLI is too old; please download the latest",
                exit_code=4,
                code="CLI_UPGRADE_REQUIRED",
            )
        if r.status_code >= 500:
            raise CliRestError(
                f"server error (status {r.status_code})",
                exit_code=5,
                code="SERVER_ERROR",
            )
        if r.status_code != 200:
            raise CliRestError(
                _extract_envelope_message(r)
                or f"unexpected MCP status {r.status_code}",
                exit_code=4,
                code="UNEXPECTED_STATUS",
            )

        ctype = (r.headers.get("content-type") or "").lower()
        if ctype.startswith("application/json"):
            # Some MCP servers (or future Streamable-HTTP modes) reply
            # with a plain JSON envelope when ``Accept`` is satisfied;
            # handle both.
            try:
                envelope = r.json()
            except ValueError as exc:
                raise CliRestError(
                    "MCP returned malformed JSON",
                    exit_code=5,
                    code="BAD_RESPONSE",
                ) from exc
        elif "text/event-stream" in ctype:
            try:
                envelope = parse_sse_payload(r.text)
            except SseParseError as exc:
                raise CliRestError(
                    f"failed to parse MCP SSE response: {exc}",
                    exit_code=5,
                    code="BAD_RESPONSE",
                ) from exc
        else:
            raise CliRestError(
                f"MCP returned unexpected content-type: {ctype!r}",
                exit_code=5,
                code="BAD_RESPONSE",
            )

        if not isinstance(envelope, dict):
            raise CliRestError(
                "MCP response was not a JSON-RPC object",
                exit_code=5,
                code="BAD_RESPONSE",
            )
        if "error" in envelope and envelope["error"] is not None:
            err = envelope["error"]
            if isinstance(err, dict):
                raise McpJsonRpcError(
                    int(err.get("code", -1)),
                    str(err.get("message", "unknown error")),
                )
            raise CliRestError(
                f"MCP returned malformed error: {err!r}",
                exit_code=5,
                code="BAD_RESPONSE",
            )
        return envelope

    # ── public RPCs ───────────────────────────────────────────────────

    def initialize(
        self,
        *,
        client_name: str = "helix-cli",
        client_version: str = "1.0",
    ) -> dict:
        """Send the JSON-RPC ``initialize`` handshake.

        Returns the ``result`` object (containing ``serverInfo``,
        ``capabilities``, ``protocolVersion``).
        """
        envelope = self._post({
            "jsonrpc": "2.0",
            "id": self._next_request_id(),
            "method": "initialize",
            "params": {
                "protocolVersion": MCP_PROTOCOL_VERSION,
                "capabilities": {},
                "clientInfo": {"name": client_name, "version": client_version},
            },
        })
        return envelope.get("result", {})

    def list_tools(self) -> list[dict]:
        """Return the list of registered MCP tools."""
        envelope = self._post({
            "jsonrpc": "2.0",
            "id": self._next_request_id(),
            "method": "tools/list",
            "params": {},
        })
        result = envelope.get("result", {})
        tools = result.get("tools", [])
        if not isinstance(tools, list):
            raise CliRestError(
                "MCP tools/list returned a non-list payload",
                exit_code=5,
                code="BAD_RESPONSE",
            )
        return tools

    def call_tool(
        self,
        name: str,
        arguments: dict | None = None,
    ) -> McpToolResult:
        """Invoke ``tools/call`` and return the unwrapped payload."""
        envelope = self._post({
            "jsonrpc": "2.0",
            "id": self._next_request_id(),
            "method": "tools/call",
            "params": {"name": name, "arguments": arguments or {}},
        })
        result = envelope.get("result")
        if not isinstance(result, dict):
            raise CliRestError(
                "MCP tools/call returned no result",
                exit_code=5,
                code="BAD_RESPONSE",
            )
        is_error = bool(result.get("isError", False))
        payload = _extract_text_block(result)
        if is_error:
            # FastMCP funnels tool exceptions into ``isError=True``
            # with a text content-block describing the failure. Treat
            # that as an exit-code-4 error from the user's POV.
            msg = payload if isinstance(payload, str) else json.dumps(payload)
            raise CliRestError(
                f"tool {name!r} failed: {msg}",
                exit_code=4,
                code="TOOL_ERROR",
            )
        return McpToolResult(payload=payload, is_error=is_error, raw=result)


def _extract_text_block(result: dict) -> Any:
    """Pull the first text content-block out of an MCP tool result.

    Helix tools return ``serialize(...)`` JSON inside a single
    ``{"type": "text", "text": "<json>"}`` block. We try to JSON-decode
    the text; if it isn't JSON we hand the raw string back so the
    caller can still display it.

    When neither a text content-block nor a ``structuredContent`` field
    is present we return ``None`` rather than the whole JSON-RPC result
    envelope — surfacing ``isError`` / ``_meta`` / future-protocol keys
    to the operator would leak server internals into the CLI output and
    break the "machine-readable JSON or human summary" contract.
    """
    content = result.get("content")
    if isinstance(content, list) and content:
        for block in content:
            if not isinstance(block, dict):
                continue
            if block.get("type") == "text":
                text = block.get("text", "")
                try:
                    return json.loads(text)
                except (TypeError, ValueError):
                    return text
        # Non-text content blocks (image / resource / future types) fall
        # through to ``structuredContent`` rather than dumping the raw
        # block list, which contains base64 blobs the user can't read.
    structured = result.get("structuredContent")
    if structured is not None:
        return structured
    return None


def _extract_envelope_message(r: httpx.Response) -> str | None:
    """Best-effort extraction of an error message from a Helix REST
    envelope returned by the PAT middleware on auth failures.
    """
    try:
        body = r.json()
    except ValueError:
        return None
    if not isinstance(body, dict):
        return None
    errors = body.get("errors") or []
    if errors and isinstance(errors[0], dict):
        return errors[0].get("message")
    return None


__all__ = [
    "MCP_PROTOCOL_VERSION",
    "SseParseError",
    "parse_sse_payload",
    "McpToolResult",
    "McpJsonRpcError",
    "McpClient",
]