"""Call a governed agent like a typed function. Python 3.12+, standard library only.

    triage = AgentFunction(client, agent_id, schema)
    accepted = triage.submit("Triage this ticket: ...", idempotency_key=key)   # 202 in milliseconds
    value = triage.result(accepted["request_id"])                            # validated JSON

`submit` is safe to retry with the same key and body: after a lost response or
a crash, the platform returns the original acceptance instead of new work.
`result` waits with a revision cursor; giving up never cancels the agent's work.
"""

import json
import time
from urllib.error import HTTPError
from urllib.request import Request, urlopen


USER_AGENT = "dt-typed-agent-calls-example/1.0"
TERMINAL_RESULT_STATES = {"validated", "validation_failed", "generation_failed", "no_result", "purged"}


class AgentCallError(RuntimeError):
    """The request reached a state where no typed value will arrive without someone acting."""

    def __init__(self, code, request):
        super().__init__(code)
        self.code, self.request = code, request


class ApiError(RuntimeError):
    def __init__(self, status, payload, headers):
        self.status, self.payload, self.headers = status, payload, headers
        super().__init__(f"HTTP {status}: {payload.get('code') or payload.get('detail') or payload}")


class Client:
    def __init__(self, api_base, api_key):
        self.api_base, self.api_key = api_base.rstrip("/"), api_key

    def request(self, method, path, body=None, headers=None, timeout=45):
        data = json.dumps(body).encode() if body is not None else None
        all_headers = {"X-Api-Key": self.api_key, "Accept": "application/json", "User-Agent": USER_AGENT,
                       **({"Content-Type": "application/json"} if data else {}), **(headers or {})}
        try:
            with urlopen(Request(f"{self.api_base}{path}", data=data, headers=all_headers, method=method),
                         timeout=timeout) as response:
                raw = response.read()
                return (json.loads(raw) if raw else {}), dict(response.headers)
        except HTTPError as exc:
            raw = exc.read()
            try:
                payload = json.loads(raw)
            except ValueError:
                payload = {"detail": raw[:300].decode(errors="replace")}
            raise ApiError(exc.code, payload if isinstance(payload, dict) else {"detail": payload},
                           dict(exc.headers)) from None


def _retry_after(error, default=1.0):
    try:
        return max(float(error.headers.get("Retry-After", default)), 0.5)
    except ValueError:
        return default


class AgentFunction:
    def __init__(self, client, agent_id, result_schema):
        self.client, self.agent_id, self.result_schema = client, agent_id, result_schema
        self.base = f"/agents/{agent_id}/requests/"

    def submit(self, body, *, idempotency_key, client_reference="", attempts=5):
        """Durably hand work to the agent. Returns (acceptance, replayed)."""
        payload = {"body": body, "result_schema": self.result_schema}
        if client_reference:
            payload["client_reference"] = client_reference
        for attempt in range(attempts):
            try:
                accepted, headers = self.client.request(
                    "POST", self.base, payload, headers={"Idempotency-Key": idempotency_key})
                return accepted, headers.get("Idempotency-Replayed") == "true"
            except ApiError as exc:
                # Backpressure and transient unavailability: same key, same body, later.
                if exc.status in {429, 503} and exc.payload.get("code") != "structured_results_disabled" \
                        and attempt + 1 < attempts:
                    time.sleep(_retry_after(exc))
                    continue
                raise
            except OSError:
                # Ambiguous: the request may have been accepted. Retrying the same key is safe.
                if attempt + 1 == attempts:
                    raise
                time.sleep(1 + attempt)

    def detail(self, request_id):
        return self.client.request("GET", f"{self.base}{request_id}/")[0]

    def wait(self, request_id, *, timeout_seconds=600, on_change=None):
        """Follow revisions until the typed result settles or needs a person."""
        snapshot = self.detail(request_id)
        cursor, deadline = snapshot["cursor"], time.monotonic() + timeout_seconds
        while True:
            if on_change:
                on_change(snapshot)
            self._raise_if_stuck(snapshot)
            if snapshot["result"]["state"] in TERMINAL_RESULT_STATES:
                return snapshot
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TimeoutError(f"{request_id} is still running; the agent keeps working. Call wait() again later.")
            try:
                waited, _ = self.client.request("POST", f"{self.base}{request_id}/wait/",
                                                {"cursor": cursor, "timeout_seconds": int(min(30, remaining))})
            except ApiError as exc:
                if exc.status in {429, 503}:
                    time.sleep(_retry_after(exc))
                    continue
                raise
            cursor = waited["cursor"]
            if waited["changed"]:
                snapshot = waited["request"]

    @staticmethod
    def _raise_if_stuck(snapshot):
        if snapshot["execution_state"] == "awaiting_action":
            raise AgentCallError("awaiting_human_decision", snapshot)  # See its pending_actions.
        if snapshot["publication_state"] == "blocked" or snapshot["execution_state"] in {"blocked", "failed"}:
            raise AgentCallError(snapshot.get("blocked_reason") or snapshot.get("failure_code") or "blocked", snapshot)

    def result(self, request_id, **wait_options):
        """Wait, then return the validated native JSON value."""
        snapshot = self.wait(request_id, **wait_options)
        state = snapshot["result"]["state"]
        if state != "validated":
            raise AgentCallError(f"result_{state}:{snapshot['result'].get('error_code') or ''}".rstrip(":"), snapshot)
        typed, _ = self.client.request("GET", f"{self.base}{request_id}/result/")
        if typed.get("availability") != "available" or not typed.get("has_value"):
            raise AgentCallError(f"result_{typed.get('availability')}", snapshot)
        return typed["value"]  # Native JSON. A validated null is a real value.