"""Agent-side helper for the budget gateway, in one file with no dependencies (Python 3.8+).

The agent keeps the official OpenAI or Anthropic SDK and points it at the gateway; this module adds
what those SDKs cannot do: read the agent's budget, ask for more, wait for the decision, and recognize
a call the gateway refused because a budget could not cover it.

    from budget_agent import BudgetAgent, budget_exceeded
    from anthropic import Anthropic

    agent = BudgetAgent(api_key=os.environ["AGENT_KEY"], base_url="https://prism.example.com")
    client = Anthropic(api_key=agent.api_key, base_url=agent.anthropic_base_url)
    try:
        client.messages.create(model="claude-sonnet-5-5", max_tokens=2000, messages=[...])
    except Exception as e:
        blocked = budget_exceeded(e)
        if not blocked:
            raise
        req = agent.request_budget(amount_usd="20", purpose="Finish the report", justification="…",
                                   expected_outcome="Report delivered", duration_hours=24, idempotency_key="task-42")
        req = agent.wait_for_decision(req["id"], timeout_s=600)
"""
import json
import re
import time
import uuid
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Dict, Optional

__all__ = ["BudgetAgent", "GatewayError", "budget_exceeded"]


class GatewayError(Exception):
    def __init__(self, status: int, type_: str, message: str, body: Any = None):
        super().__init__(message)
        self.status, self.type, self.body = status, type_, body


class BudgetAgent:
    def __init__(self, api_key: str, base_url: str, timeout_s: float = 30.0):
        if not re.match(r"^aib_(live|test)_", api_key or ""):
            raise ValueError("BudgetAgent needs an agent key (aib_live_…)")
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout_s = timeout_s

    # Base URLs to give the provider SDKs.
    @property
    def anthropic_base_url(self) -> str:
        return f"{self.base_url}/gw/anthropic"

    @property
    def openai_base_url(self) -> str:
        return f"{self.base_url}/gw/openai/v1"

    def provider_base_url(self, provider: str) -> str:
        """Base URL for any OpenAI-compatible provider (use it with the OpenAI SDK): mistral, google, xai…"""
        return f"{self.base_url}/gw/{provider}/v1"

    def budget(self) -> Dict[str, Any]:
        """Limits, spent, reserved and remaining for today and this month, and active temporary increases."""
        return self._call("GET", "/v1/agent/me/budget")

    def request_budget(self, amount_usd, purpose: str, justification: str, expected_outcome: str,
                       duration_hours: int, task_id: Optional[str] = None, idempotency_key: Optional[str] = None) -> Dict[str, Any]:
        """Ask for a temporary increase. Pass an idempotency key derived from the task so a retry
        (even after a restart) never files the request twice."""
        body = {"amount_usd": str(amount_usd), "purpose": purpose, "justification": justification,
                "expected_outcome": expected_outcome, "duration_hours": duration_hours}
        if task_id:
            body["task_id"] = task_id
        return self._call("POST", "/v1/agent/budget-requests", body, idempotency_key or str(uuid.uuid4()))

    def get_request(self, request_id: str) -> Dict[str, Any]:
        return self._call("GET", f"/v1/agent/budget-requests/{urllib.parse.quote(request_id)}")

    def cancel_request(self, request_id: str) -> Dict[str, Any]:
        return self._call("POST", f"/v1/agent/budget-requests/{urllib.parse.quote(request_id)}/cancel", {})

    def wait_for_decision(self, request_id: str, timeout_s: float = 600, interval_s: float = 2) -> Dict[str, Any]:
        """Poll until the request is decided, or the timeout passes (it is then returned still pending)."""
        deadline = time.monotonic() + timeout_s
        while True:
            r = self.get_request(request_id)
            if r.get("status") != "pending" or time.monotonic() >= deadline:
                return r
            time.sleep(min(interval_s, max(0.0, deadline - time.monotonic())))
            interval_s = min(interval_s * 1.5, 30)

    def _call(self, method: str, path: str, body: Any = None, idempotency_key: Optional[str] = None) -> Any:
        headers = {"authorization": f"Bearer {self.api_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
        req = urllib.request.Request(self.base_url + path, data=data, method=method, headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=self.timeout_s) as res:
                text = res.read().decode()
                return json.loads(text) if text else None
        except urllib.error.HTTPError as e:
            text = e.read().decode()
            try:
                j = json.loads(text) if text else None
            except ValueError:
                j = None
            err = (j or {}).get("error") or {}
            raise GatewayError(e.code, err.get("type", "error"), err.get("message", f"The gateway answered {e.code}"), j) from None


def budget_exceeded(err: Any) -> Optional[Dict[str, Any]]:
    """What blocked a call, if `err` is the gateway's 402 budget refusal; otherwise None.

    Works with exceptions from the `openai` and `anthropic` Python SDKs (APIStatusError), with
    GatewayError, and with a plain response body dict."""
    status = getattr(err, "status_code", None) or getattr(err, "status", None)
    if status is not None and status != 402:
        return None
    body = getattr(err, "body", None) if not isinstance(err, dict) else err
    if not isinstance(body, dict):
        return None
    error = body.get("error", body)          # Anthropic: {"type":"error","error":{…}}; OpenAI: {"error":{…}} or the error itself
    if isinstance(error, dict) and "error" in error and isinstance(error["error"], dict):
        error = error["error"]
    aib = (error or {}).get("aib") if isinstance(error, dict) else None
    if not aib or error.get("type") != "budget_exceeded":
        return None
    return {"blocking_account": aib.get("blocking_account"), "window": aib.get("window"), "limit_usd": aib.get("limit_usd"),
            "remaining_usd": aib.get("remaining_usd"), "required_usd": aib.get("required_usd"), "transaction_id": aib.get("transaction_id")}