"""Small dependency-free MCP client for OpenClasp Shield benchmark runs."""

from __future__ import annotations

import json

import uuid

from typing import Any

from urllib.error import HTTPError, URLError

from urllib.request import Request, urlopen

class OpenClaspMcpError(RuntimeError):

    pass

class OpenClaspMcpClient:

    def __init__(self, base_url: str, token: str, timeout_seconds: float = 58.0):

        if not base_url.startswith(("http://", "https://")):

            raise ValueError("OPENCLASP_URL must be an HTTP(S) URL")

        if not token.startswith("oc_at_"):

            raise ValueError("OPENCLASP_AGENT_TOKEN must be an oc_at_ token")

        self.endpoint = (

            base_url.rstrip("/")

            if base_url.rstrip("/").endswith("/mcp")

            else f"{base_url.rstrip('/')}/mcp"

        )

        self._token = token

        self._timeout_seconds = timeout_seconds

    def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> Any:

        payload = json.dumps(

            {

                "jsonrpc": "2.0",

                "id": str(uuid.uuid4()),

                "method": "tools/call",

                "params": {"name": name, "arguments": arguments or {}},

            }

        ).encode("utf-8")

        request = Request(

            self.endpoint,

            data=payload,

            method="POST",

            headers={

                "Authorization": f"Bearer {self._token}",

                "Accept": "application/json, text/event-stream",

                "Content-Type": "application/json",

                "MCP-Protocol-Version": "2025-06-18",

            },

        )

        try:

            with urlopen(request, timeout=self._timeout_seconds) as response:

                body = response.read().decode("utf-8")

        except HTTPError as error:

            body = error.read().decode("utf-8", errors="replace")

            raise OpenClaspMcpError(

                f"OpenClasp MCP returned HTTP {error.code}: {body[:500]}"

            ) from error

        except (URLError, TimeoutError) as error:

            raise OpenClaspMcpError(f"OpenClasp MCP request failed: {error}") from error

        rpc = _parse_rpc_body(body)

        if rpc.get("error"):

            message = rpc["error"].get("message", "Unknown MCP error")

            raise OpenClaspMcpError(str(message))

        result = rpc.get("result")

        if not isinstance(result, dict):

            raise OpenClaspMcpError("OpenClasp returned no MCP result")

        if result.get("isError"):

            raise OpenClaspMcpError(_content_text(result) or "OpenClasp tool failed")

        text = _content_text(result)

        if text is None:

            return result

        try:

            return json.loads(text)

        except json.JSONDecodeError:

            return text

    def open_case(

        self,

        *,

        title: str,

        goal: str,

        domain_policy: str,

    ) -> dict[str, Any]:

        policies = [

            {

                "title": f"Visible τ³ domain policy {index + 1}",

                "statement": chunk,

            }

            for index, chunk in enumerate(_split_policy(domain_policy))

        ]

        value = self.call_tool(

            "openclasp_shield_open_case",

            {

                "title": title[:160],

                "goal": goal[:2000],

                "brief": (

                    "Benchmark evaluation case. Shield receives only benchmark-visible policy, "

                    "messages, tool results, and the protected agent's proposed next step."

                ),

                "counterparty": {"type": "human"},

                "facts": [],

                "evidence": [],

                "policies": policies,

            },

        )

        if not isinstance(value, dict) or not value.get("caseId"):

            raise OpenClaspMcpError("Shield did not return a caseId")

        return value

    def consult(

        self,

        *,

        case_id: str,

        draft: str,

        situation_context: str,

    ) -> dict[str, Any]:

        value = self.call_tool(

            "openclasp_shield_consult",

            {

                "caseId": case_id,

                "message": (

                    "Review the protected agent's proposed next step. Check it against the "

                    "visible policy and verified tool results. Identify unsupported user claims, "

                    "pressure tactics, missing evidence, and the safest correct response.\n\n"

                    f"PROPOSED NEXT STEP\n{draft}"

                )[:4000],

                "situationContext": situation_context[:8000],

                "proposedAction": draft[:1000],

                "facts": [],

                "evidence": [],

                "policies": [],

            },

        )

        if not isinstance(value, dict) or not isinstance(value.get("consultation"), dict):

            raise OpenClaspMcpError("Shield returned an invalid consultation")

        return value

    def close_case(

        self,

        *,

        case_id: str,

        result: str,

        accepted_advice: bool,

        action_taken: str,

        observed_impact: str,

    ) -> dict[str, Any]:

        value = self.call_tool(

            "openclasp_shield_close_case",

            {

                "caseId": case_id,

                "result": result,

                "acceptedAdvice": accepted_advice,

                "actionTaken": action_taken[:2000],

                "observedImpact": observed_impact[:2000],

            },

        )

        if not isinstance(value, dict):

            raise OpenClaspMcpError("Shield returned an invalid outcome")

        return value

def _split_policy(policy: str, limit: int = 1900) -> list[str]:

    chunks: list[str] = []

    current = ""

    for line in policy.splitlines():

        remaining = line

        while len(remaining) > limit:

            if current:

                chunks.append(current)

                current = ""

            chunks.append(remaining[:limit])

            remaining = remaining[limit:]

        candidate = f"{current}\n{remaining}".strip() if current else remaining

        if len(candidate) > limit:

            chunks.append(current)

            current = remaining

        else:

            current = candidate

    if current:

        chunks.append(current)

    if len(chunks) > 50:

        raise ValueError("τ³ domain policy is too large for a Shield case")

    return chunks

def _parse_rpc_body(body: str) -> dict[str, Any]:

    candidates = [

        line[5:].strip()

        for line in body.splitlines()

        if line.startswith("data:") and line[5:].strip()

    ]

    raw = candidates[-1] if candidates else body

    try:

        value = json.loads(raw)

    except json.JSONDecodeError as error:

        raise OpenClaspMcpError("OpenClasp returned invalid JSON") from error

    if not isinstance(value, dict):

        raise OpenClaspMcpError("OpenClasp returned an invalid JSON-RPC envelope")

    return value

def _content_text(result: dict[str, Any]) -> str | None:

    content = result.get("content")

    if not isinstance(content, list):

        return None

    for item in content:

        if isinstance(item, dict) and item.get("type") == "text":

            text = item.get("text")

            if isinstance(text, str):

                return text

    return None