"""Small synchronous MCP client. pip install requests"""
import itertools
import json
import os
import random
import time

import requests


class MCP:
    def __init__(self, token, endpoint="https://mcp.amazingcat.net/mcp", rpm=200):
        self.endpoint = endpoint
        self.rpm = rpm
        self.last = 0.0
        self.ids = itertools.count(1)
        self.http = requests.Session()
        self.http.headers.update({
            "Authorization": f"Bearer {token}",
            "Accept": "application/json, text/event-stream",
            "Content-Type": "application/json",
        })

    def rpc(self, method, params=None, notification=False):
        payload = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            payload["params"] = params
        if not notification:
            payload["id"] = next(self.ids)
        for attempt in range(5):
            time.sleep(max(0, 60 / self.rpm - (time.monotonic() - self.last)))
            self.last = time.monotonic()
            try:
                with self.http.post(self.endpoint, json=payload, stream=True,
                                    timeout=(15, 120)) as response:
                    # This API also signals tool-call quota exhaustion with HTTP 200.
                    wait = response.headers.get("Retry-After")
                    if wait or response.status_code in (429, 502, 503, 504):
                        delay = int(wait) if wait and wait.isdigit() else min(30, 2 ** attempt)
                        if attempt == 4:
                            raise RuntimeError(f"Retry budget exhausted; HTTP {response.status_code}")
                        # A daily limit must be scheduled for later, not kept in a retry loop.
                        if delay > 120:
                            raise RuntimeError(f"Quota exhausted; retry after {delay} seconds")
                        time.sleep(delay + random.uniform(0, 0.5))
                        continue
                    if response.status_code == 404 and method != "initialize" and self.http.headers.get("Mcp-Session-Id") and attempt < 4:
                        self.http.headers.pop("Mcp-Session-Id", None)
                        self.initialize()
                        continue
                    response.raise_for_status()  # 401: replace/refresh credentials, do not retry.
                    sid = response.headers.get("Mcp-Session-Id")
                    if sid:
                        self.http.headers["Mcp-Session-Id"] = sid
                    if response.status_code == 202:
                        return None
                    if "text/event-stream" in response.headers.get("Content-Type", ""):
                        response.encoding = "utf-8"
                        data = []
                        message = None
                        for line in response.iter_lines(decode_unicode=True, delimiter="\n"):
                            line = line.rstrip("\r")
                            if line.startswith("data:"):
                                data.append(line[5:].lstrip(" "))
                            elif not line and data:
                                event = json.loads("\n".join(data))
                                data = []
                                if event.get("id") == payload.get("id"):
                                    message = event
                                    break
                        if message is None:
                            raise RuntimeError("SSE ended without the matching JSON-RPC response")
                    else:
                        message = response.json()
            except (requests.ConnectionError, requests.Timeout):
                if attempt == 4:
                    raise
                time.sleep(min(30, 2 ** attempt) + random.uniform(0, 0.5))
                continue
            if "error" in message:
                raise RuntimeError(message["error"])
            return message["result"]
        raise RuntimeError("Retry budget exhausted")

    def initialize(self):
        result = self.rpc("initialize", {
            "protocolVersion": "2025-06-18", "capabilities": {},
            "clientInfo": {"name": "amazingcat-api-example", "version": "1.0"},
        })
        self.http.headers["MCP-Protocol-Version"] = result["protocolVersion"]
        self.rpc("notifications/initialized", notification=True)
        return result

    def call(self, name, arguments):
        for attempt in range(5):
            result = self.rpc("tools/call", {"name": name, "arguments": arguments})
            if result.get("isError"):
                raise RuntimeError(result.get("content"))
            blocks = [b["text"] for b in result.get("content", []) if b.get("type") == "text"]
            if not blocks:
                raise RuntimeError("No text content; inspect the MCP result")
            try:
                value = json.loads(blocks[0])
            except json.JSONDecodeError:
                # Plain text can be a quota message, never interpret it as zero results.
                raise RuntimeError(blocks[0]) from None
            if isinstance(value, dict) and value.get("error") == "server_busy":
                if attempt < 4:
                    time.sleep(value.get("retry_after_seconds", 5) + random.uniform(0, 0.5))
                    continue
            if isinstance(value, dict) and (value.get("error") or value.get("ошибка")):
                raise RuntimeError(value)
            return value
        raise RuntimeError("Server busy; retry later")

    def close(self):
        if "Mcp-Session-Id" in self.http.headers:
            self.http.delete(self.endpoint, timeout=(15, 30)).close()
        self.http.close()


def main():
    client = MCP(os.environ["AMAZINGCAT_TOKEN"])
    try:
        client.initialize()
        print("tools:", len(client.rpc("tools/list")["tools"]))
        result = client.call("kas_search", {"text": "принятие SENTENCE наследства", "court": 0, "kind": "civil", "limit": 1, "detail": "brief"})
        print(json.dumps(result, ensure_ascii=False))
    finally:
        client.close()


if __name__ == "__main__":
    main()