"""AGNT Models client. Standard library only; Python 3.9+.

Any OpenAI-compatible SDK also works: set base_url to
https://models.agnt.gg/models/v1 and model to "agnt-flash".
"""
import json
import urllib.request
import uuid

BASE = "https://models.agnt.gg/models/v1"


class ModelsError(Exception):
    def __init__(self, status, code):
        super().__init__(f"{status} {code}")
        self.status = status
        self.code = code


class ModelsClient:
    def __init__(self, key, base=BASE):
        self.key = key
        self.base = base.rstrip("/")

    def _request(self, path, method="GET", body=None, idempotency_key=None):
        headers = {"Authorization": "Bearer " + self.key}
        data = None
        if body is not None:
            headers["Content-Type"] = "application/json"
            data = json.dumps(body).encode()
        if idempotency_key:
            headers["Idempotency-Key"] = idempotency_key
        request = urllib.request.Request(self.base + path, data=data, method=method, headers=headers)
        try:
            with urllib.request.urlopen(request, timeout=310) as response:
                return json.load(response)
        except urllib.error.HTTPError as error:
            try:
                code = json.load(error).get("error", "request_failed")
            except Exception:
                code = "request_failed"
            raise ModelsError(error.code, code) from None

    def chat(self, messages, idempotency_key=None, **options):
        """Run a completion. Reuse idempotency_key on an uncertain retry to get the receipt, not a second run."""
        body = {"model": "agnt-flash", "messages": messages, "stream": False, **options}
        return self._request("/chat/completions", "POST", body, idempotency_key or str(uuid.uuid4()))

    def stream(self, messages, idempotency_key=None, **options):
        """Yield OpenAI-shaped chunks. The last chunk carries usage and requestId."""
        body = {"model": "agnt-flash", "messages": messages, "stream": True, **options}
        headers = {"Authorization": "Bearer " + self.key, "Content-Type": "application/json", "Idempotency-Key": idempotency_key or str(uuid.uuid4())}
        request = urllib.request.Request(self.base + "/chat/completions", data=json.dumps(body).encode(), method="POST", headers=headers)
        with urllib.request.urlopen(request, timeout=310) as response:
            for raw in response:
                line = raw.decode().strip()
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    return
                chunk = json.loads(data)
                if "error" in chunk:
                    raise ModelsError(500, chunk["error"].get("message", "request_failed"))
                yield chunk

    def quote(self, messages, **options):
        return self._request("/quote", "POST", {"model": "agnt-flash", "messages": messages, **options})

    def usage(self):
        return self._request("/usage")

    def receipts(self):
        return self._request("/requests")

    def receipt(self, request_id):
        return self._request("/requests/" + request_id)