"""Reference client for the documented text webhook; not a published platform SDK."""

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import httpx
from jsonschema import Draft202012Validator


class CheckUnavailable(Exception):
    """Do not persist, return or send the unchecked text downstream."""


class PolicyDenied(Exception):
    """A completed check rejected the content; retrying is not a recovery action."""


@dataclass(frozen=True)
class CheckedText:
    outcome: Literal["pass", "mask"]
    messages: list[dict[str, str]]


def _unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("duplicate JSON property")
        result[key] = value
    return result


class GuardrailsClient:
    """Use one instance per worker lifecycle, with deployment-owned configuration.

    base_url must be an approved HTTPS service address; plain HTTP is only for
    the loopback quickstart. Pass a verified TLS context through verify for an
    internal CA. Never use verify=False or derive configuration from user input.
    """

    def __init__(
        self,
        *,
        base_url: str,
        token_file: Path,
        policy_id: str,
        revision: str,
        digest: str,
        schema_file: Path,
        timeout_seconds: float = 1.0,
        max_body_bytes: int = 32768,
        max_messages: int = 32,
        verify=True,
    ):
        url = httpx.URL(base_url)
        local = url.scheme == "http" and url.host in ("127.0.0.1", "::1")
        if (
            not (url.scheme == "https" or local)
            or url.userinfo
            or url.query
            or url.fragment
            or url.path not in ("", "/")
        ):
            raise ValueError("Use an approved service origin")
        schema = json.loads(schema_file.read_text())
        self._validators = {
            name: Draft202012Validator(
                {
                    "$ref": f"#/components/schemas/{name}",
                    "components": schema["components"],
                }
            )
            for name in ("Request", "Response", "RequestResult", "ResponseResult")
        }
        self._max_body_bytes = max_body_bytes
        self._max_messages = max_messages
        token = token_file.read_text().rstrip("\r\n")
        if not 32 <= len(token) <= 4096 or any(not 33 <= ord(c) <= 126 for c in token):
            raise ValueError("Invalid service token file")
        self._http = httpx.Client(
            base_url=str(url),
            timeout=timeout_seconds,
            verify=verify,
            trust_env=False,
            follow_redirects=False,
            headers={
                "Authorization": f"Bearer {token}",
                "x-ap-gr-policy-id": policy_id,
                "x-ap-gr-policy-revision": revision,
                "x-ap-gr-bundle-digest": digest,
                "x-ap-gr-content-profile": "agentgateway-1.4.1-text-webhook-v1",
                "Accept-Encoding": "identity",
            },
        )

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self._http.close()

    def check(
        self, stage: Literal["request", "response"], messages: list[dict[str, str]]
    ) -> CheckedText:
        if stage not in ("request", "response"):
            raise ValueError("Unknown stage")
        schema_name = "Request" if stage == "request" else "Response"
        body = (
            {"messages": messages}
            if stage == "request"
            else {"choices": [{"message": m} for m in messages]}
        )
        payload = {"body": body}
        if not self._validators[schema_name].is_valid(payload):
            raise CheckUnavailable("invalid_local_payload")
        encoded = json.dumps(payload, ensure_ascii=False, allow_nan=False).encode()
        if len(messages) > self._max_messages or len(encoded) > self._max_body_bytes:
            raise CheckUnavailable("local_limit_exceeded")
        try:
            with self._http.stream(
                "POST", f"/{stage}", content=encoded, headers={"Content-Type": "application/json"}
            ) as response:
                if response.status_code != 200:
                    raise CheckUnavailable(f"guardrails_http_{response.status_code}")
                if response.headers.get("content-encoding", "identity") != "identity":
                    raise CheckUnavailable("unexpected_encoding")
                if response.headers.get("content-type", "").split(";")[0] != "application/json":
                    raise CheckUnavailable("unexpected_media_type")
                data = bytearray()
                for chunk in response.iter_raw(chunk_size=8192):
                    data.extend(chunk)
                    if len(data) > self._max_body_bytes:
                        raise CheckUnavailable("response_limit_exceeded")
                result = json.loads(data, object_pairs_hook=_unique_object)
                json.dumps(result, ensure_ascii=False, allow_nan=False).encode("utf-8")
        except httpx.HTTPError, ValueError, UnicodeError, RecursionError:
            raise CheckUnavailable("guardrails_transport_or_json_error") from None
        if not self._validators[schema_name + "Result"].is_valid(result):
            raise CheckUnavailable("invalid_action")
        action = result["action"]
        if action["reason"] == "policy_denied":
            raise PolicyDenied("policy_denied")
        if action["reason"] == "policy_pass":
            return CheckedText("pass", [dict(message) for message in messages])
        changed = (
            action["body"]["messages"]
            if stage == "request"
            else [c["message"] for c in action["body"]["choices"]]
        )
        if len(changed) != len(messages) or [m["role"] for m in changed] != [
            m["role"] for m in messages
        ]:
            raise CheckUnavailable("invalid_mask_mapping")
        return CheckedText("mask", changed)