"""agy provider adapter — Google Antigravity accounts behind ModelProviderPort.

``agy`` is Google's Antigravity: a Google-AI-subscription-billed route to the
Gemini 3 / Claude 4.5 / GPT-OSS families. There is no public API key; instead
each *account* is a Google OAuth identity whose refresh token is exchanged for
short-lived access tokens. This module implements that protocol directly, so
the platform never shells out to the CLI binary:

* :class:`AgyAuthClient` — PKCE OAuth in two shapes: **listener mode** (a
  loopback HTTP server captures the redirect; used when the operator's browser
  runs on the same machine) and **paste mode** (the operator copies the
  callback URL from the address bar and hands it back; the headless/remote
  case — the page after redirect never finishes loading, and that is expected).
* :class:`AgyModelAdapter` — the ``ModelProviderPort`` implementation. Converts
  our message format to the ``/v1internal:generateContent`` envelope, calls
  Google's cloudcode-pa backend, parses the Gemini-shaped response.
* :class:`AgyAccountPool` state inside the adapter — round-robin across the
  logged-in accounts with automatic 429 benching: a rate-limited account is
  skipped until its cooldown expires, both in-process (fast) and in the
  database (so background workers respect a cooldown set by the API).

Nothing here ever returns or logs a refresh token: :func:`load_agy_accounts`
hands the adapter decrypted material in memory only, and admin endpoints see
email/project/status, never ciphertext.
"""

from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import logging
import secrets
import time
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from urllib.parse import parse_qs, urlencode, urlparse

import httpx

from app.application.ports import ChatMessage, ModelHealth, ModelResponse
from app.core.config import settings
from app.core.crypto import decrypt
from app.core.errors import ProviderUnavailableError, ValidationError
from app.domain.ai_control.models import AiModel, ModelCredential
from app.domain.base import utcnow

logger = logging.getLogger("agy.adapter")

# ---------------------------------------------------------------------------
# Protocol constants (the public Antigravity CLI OAuth application and the
# cloudcode-pa backend it talks to).
# ---------------------------------------------------------------------------

OAUTH_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token"
USERINFO_URL = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json"

#: ``daily-cloudcode-pa`` first: it carries the free daily quota and 429s more
#: gracefully. ``cloudcode-pa`` is the fallback.
BASE_URLS: tuple[str, ...] = (
    "https://daily-cloudcode-pa.sandbox.googleapis.com",
    "https://cloudcode-pa.googleapis.com",
)
PATH_GENERATE = "/v1internal:generateContent"
PATH_MODELS = "/v1internal:fetchAvailableModels"
PATH_LOAD_PROJECT = "/v1internal:loadCodeAssist"

OAUTH_SCOPES: tuple[str, ...] = (
    "https://www.googleapis.com/auth/cloud-platform",
    "https://www.googleapis.com/auth/userinfo.email",
    "https://www.googleapis.com/auth/userinfo.profile",
    "https://www.googleapis.com/auth/cclog",
    "https://www.googleapis.com/auth/experimentsandconfigs",
)

#: Same all-OFF matrix the Antigravity client sends. Without it Google applies
#: its defaults and refuses coding prompts it considers dangerous.
SAFETY_SETTINGS: list[dict[str, str]] = [
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
    {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"},
]

#: How long a started login stays redeemable. Two OAuth round-trips plus a
#: manual paste must fit inside it comfortably.
CHALLENGE_TTL_SECONDS = 600

#: agy ``reasoning_effort`` → thinking budget (tokens), mirroring the Antigravity
#: client. An unknown value falls back to ``medium`` server-side, and the UI
#: only ever offers these five.
EFFORT_THINKING_BUDGETS: dict[str, int] = {
    "minimal": 1024,
    "low": 4096,
    "medium": 8192,
    "high": 16384,
    "xhigh": 32768,
}

#: Where thinking is honoured. A request must never carry ``thinkingConfig`` to
#: a model that does not support it — that is a hard 400 from the backend.
_THINKING_MARKERS: tuple[str, ...] = (
    "thinking",
    "gemini-3",
    "gemini-2.5-pro",
    "gemini-2.5-flash",
)


def _supports_thinking(model_name: str) -> bool:
    lowered = model_name.lower()
    if "image" in lowered:
        return False
    return any(marker in lowered for marker in _THINKING_MARKERS)

_StatusFn = Callable[..., Awaitable[None]]
_LoaderFn = Callable[[], Awaitable[list[ModelCredential]]]


# ---------------------------------------------------------------------------
# Small helpers
# ---------------------------------------------------------------------------


def _aware(value: datetime | None) -> datetime | None:
    """SQLite hands back naive datetimes; comparisons need aware ones."""
    if value is None:
        return None
    if value.tzinfo is None:
        return value.replace(tzinfo=UTC)
    return value


def _pkce_pair() -> tuple[str, str]:
    verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode()
    digest = hashlib.sha256(verifier.encode()).digest()
    challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
    return verifier, challenge


def _random_session_id() -> str:
    return f"-{secrets.randbelow(9_000_000_000_000_000_000)}"


def _fallback_project_id() -> str:
    """Ad-hoc project when ``loadCodeAssist`` did not report one."""
    adjectives = ("useful", "bright", "swift", "calm", "bold")
    nouns = ("fuze", "wave", "spark", "flow", "core")
    adj = secrets.choice(adjectives)
    noun = secrets.choice(nouns)
    return f"{adj}-{noun}-{uuid.uuid4().hex[:5].lower()}"


def _retry_after_seconds(value: str | None) -> int | None:
    if not value:
        return None
    try:
        seconds = int(value.strip())
    except ValueError:
        return None
    return max(1, min(seconds, 4 * 3600))


# ---------------------------------------------------------------------------
# Internal failure taxonomy (never crosses the port boundary)
# ---------------------------------------------------------------------------


class _AgyRateLimited(Exception):
    def __init__(self, retry_after: int | None) -> None:
        super().__init__("agy account rate limited")
        self.retry_after = retry_after


class _AgyAuthError(Exception):
    """Refresh token rejected / account not entitled — needs a fresh login."""


class _AgyTransientError(Exception):
    """5xx or transport failure on one backend URL — try the next account."""


# ---------------------------------------------------------------------------
# OAuth login
# ---------------------------------------------------------------------------


@dataclass(slots=True, kw_only=True)
class AgyLoginChallenge:
    login_id: str
    auth_url: str
    redirect_uri: str
    verifier: str
    state: str
    created_at: float


@dataclass(slots=True, kw_only=True)
class AgyLoginResult:
    """Decrypted account material, produced by a completed OAuth exchange."""

    email: str
    project_id: str
    refresh_token: str


class AgyOAuthError(ValidationError):
    """The pasted URL or the token exchange was rejected by Google."""


class AgyAuthClient:
    """Google OAuth (PKCE, installed-app client) for agy accounts."""

    def __init__(
        self,
        *,
        client_id: str | None = None,
        client_secret: str | None = None,
        redirect_uri: str | None = None,
        timeout: float = 30.0,
    ) -> None:
        self.client_id = client_id or settings.agy_client_id
        self.client_secret = client_secret or settings.agy_client_secret
        self.redirect_uri = redirect_uri or settings.agy_redirect_uri
        self.timeout = timeout

    # -- flow -------------------------------------------------------------

    def begin(self) -> AgyLoginChallenge:
        """Open a login: returns the URL the operator must open in a browser.

        The verifier stays in the returned challenge — it is the proof that
        whoever redeems the login is whoever started it, and it is never
        serialised out of the process.
        """
        verifier, challenge = _pkce_pair()
        login_id = uuid.uuid4().hex
        state = uuid.uuid4().hex
        params = urlencode(
            {
                "client_id": self.client_id,
                "response_type": "code",
                "redirect_uri": self.redirect_uri,
                "scope": " ".join(OAUTH_SCOPES),
                "code_challenge": challenge,
                "code_challenge_method": "S256",
                "state": state,
                "access_type": "offline",
                "prompt": "consent",
            }
        )
        return AgyLoginChallenge(
            login_id=login_id,
            auth_url=f"{OAUTH_AUTH_URL}?{params}",
            redirect_uri=self.redirect_uri,
            verifier=verifier,
            state=state,
            created_at=time.monotonic(),
        )

    async def exchange(self, challenge: AgyLoginChallenge, callback_url: str) -> AgyLoginResult:
        """Redeem a started login with the callback URL from the browser.

        The operator may paste the full ``http://localhost:51121/oauth-callback?
        …`` URL (it never finishes loading — that is the documented shape of
        this flow, the parameters are already in the address bar).
        """
        parsed = urlparse(callback_url.strip())
        query = parse_qs(parsed.query)
        if query.get("error"):
            raise AgyOAuthError(f"Google خطای احراز هویت داد: {query['error'][0]}")
        codes = query.get("code") or []
        states = query.get("state") or []
        if not codes or not states:
            raise AgyOAuthError(
                "URL بازگشت معتبر نیست؛ باید پارامترهای code و state را داشته باشد."
            )
        if states[0] != challenge.state:
            raise AgyOAuthError("state بررسی نشد؛ این URL به این درخواست لاگین تعلق ندارد.")
        return await self._exchange_code(challenge, codes[0])

    async def _exchange_code(self, challenge: AgyLoginChallenge, code: str) -> AgyLoginResult:
        tokens = await self._post_token(
            {
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "code": code,
                "grant_type": "authorization_code",
                "redirect_uri": challenge.redirect_uri,
                "code_verifier": challenge.verifier,
            }
        )
        refresh_token = tokens.get("refresh_token") or ""
        if not refresh_token:
            # Happens when the consent screen was already granted for these
            # scopes: Google omits the refresh token. A re-run with
            # ``prompt=consent`` (which begin() always sends) fixes it, so the
            # operator message says exactly that.
            raise AgyOAuthError(
                "Google refresh_token برنگرداند. یک بار دیگر لاگین کنید و در صفحهٔ "
                "Google حتماً «Block/Allow» را تأیید کنید."
            )
        email, project_id = await asyncio.gather(
            self._fetch_email(tokens.get("access_token", "")),
            self._fetch_project_id(tokens.get("access_token", "")),
        )
        return AgyLoginResult(
            email=email, project_id=project_id, refresh_token=refresh_token
        )

    async def refresh(self, refresh_token: str) -> str:
        tokens = await self._post_token(
            {
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "refresh_token": refresh_token,
                "grant_type": "refresh_token",
            }
        )
        access_token = tokens.get("access_token")
        if not access_token:
            raise _AgyAuthError("Google refresh را نپذیرفت (توکن منقضی یا باطل شده).")
        return access_token

    # -- Google endpoints --------------------------------------------------

    async def _post_token(self, data: dict[str, str]) -> dict:
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    OAUTH_TOKEN_URL,
                    data=data,
                    headers={"Content-Type": "application/x-www-form-urlencoded"},
                )
        except httpx.HTTPError as exc:
            raise ProviderUnavailableError("اتصال به سرور OAuth گوگل برقرار نشد.") from exc
        if response.status_code >= 400:
            body = response.text[:400]
            # A rejected refresh token is an account problem, not an outage.
            if "invalid_grant" in body:
                raise _AgyAuthError(f"Google refresh را نپذیرفت: {body}")
            raise AgyOAuthError(f"تبادل توکن ناموفق بود ({response.status_code}): {body}")
        return response.json()

    async def _fetch_email(self, access_token: str) -> str:
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.get(
                    USERINFO_URL, headers={"Authorization": f"Bearer {access_token}"}
                )
            if response.is_success:
                return str(response.json().get("email", ""))
        except httpx.HTTPError:
            logger.debug("agy userinfo fetch failed", exc_info=True)
        return ""

    async def _fetch_project_id(self, access_token: str) -> str:
        for base in BASE_URLS:
            try:
                async with httpx.AsyncClient(timeout=self.timeout) as client:
                    response = await client.post(
                        f"{base}{PATH_LOAD_PROJECT}",
                        headers=self._api_headers(access_token),
                        json={
                            "metadata": {
                                "ideType": "IDE_UNSPECIFIED",
                                "platform": "PLATFORM_UNSPECIFIED",
                                "pluginType": "GEMINI",
                            }
                        },
                    )
                if response.is_success:
                    project = response.json().get("cloudaicompanionProject")
                    if isinstance(project, str) and project:
                        return project
                    if isinstance(project, dict) and project.get("id"):
                        return str(project["id"])
            except httpx.HTTPError:
                continue
        return ""

    def _api_headers(self, access_token: str) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
            "User-Agent": settings.agy_user_agent,
        }


# ---------------------------------------------------------------------------
# Pending-login store (process-local, bounded, TTL'd)
# ---------------------------------------------------------------------------

_CHALLENGES: dict[str, AgyLoginChallenge] = {}
_CHALLENGES_LOCK = asyncio.Lock()


def _prune_challenges(now: float) -> None:
    expired = [
        login_id
        for login_id, challenge in _CHALLENGES.items()
        if now - challenge.created_at > CHALLENGE_TTL_SECONDS
    ]
    for login_id in expired:
        _CHALLENGES.pop(login_id, None)


async def begin_login() -> tuple[str, str]:
    """Start a headless login. Returns ``(login_id, auth_url)``.

    ``auth_url`` is what the admin console shows the platform owner; the owner
    opens it, signs in with Google, and the resulting callback URL (or, in
    listener mode, the captured redirect) is redeemed against ``login_id``.
    """
    challenge = AgyAuthClient().begin()
    now = time.monotonic()
    async with _CHALLENGES_LOCK:
        _prune_challenges(now)
        if len(_CHALLENGES) >= 32:
            # Bound memory: pending logins are transient by nature.
            oldest = min(_CHALLENGES, key=lambda k: _CHALLENGES[k].created_at)
            _CHALLENGES.pop(oldest, None)
        _CHALLENGES[challenge.login_id] = challenge
    return challenge.login_id, challenge.auth_url


async def redeem_login(login_id: str, callback_url: str) -> AgyLoginResult:
    challenge = _CHALLENGES.pop(login_id, None)
    if challenge is None:
        raise AgyOAuthError("این درخواست لاگین منقضی شده است؛ دوباره آغاز کنید.")
    return await AgyAuthClient().exchange(challenge, callback_url)


async def redeem_login_by_waiting(
    login_id: str, *, timeout_seconds: float = 300
) -> AgyLoginResult:
    """Listener mode: capture the redirect on the loopback callback port.

    Used by the operator-side script (and by a same-machine deployment): Google
    redirects the browser to ``redirect_uri``; this server catches it, closes,
    and redeems the login without any copy-paste.
    """
    challenge = _CHALLENGES.get(login_id)
    if challenge is None:
        raise AgyOAuthError("این درخواست لاگین منقضی شده است؛ دوباره آغاز کنید.")
    callback_path = urlparse(challenge.redirect_uri).path or "/"
    port = int(urlparse(challenge.redirect_uri).port or 51121)
    loop = asyncio.get_running_loop()
    future: asyncio.Future[str] = loop.create_future()

    async def _handle(
        reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        try:
            request_line = await reader.readline()
            parts = request_line.decode("latin-1").split(" ")
            target = parts[1] if len(parts) > 1 else "/"
            if urlparse(target).path != callback_path:
                writer.write(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n")
                await writer.drain()
                return
            writer.write(
                b"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\n"
                b"Content-Length: 46\r\nConnection: close\r\n\r\n"
                b"<html><body>OK - you may close this tab.</body></html>"
            )
            await writer.drain()
            if not future.done():
                future.set_result(target)
        except Exception:  # noqa: BLE001 - the future carries the real error
            if not future.done():
                future.set_exception(AgyOAuthError("دریافت بازگشت OAuth ناموفق بود."))
        finally:
            writer.close()

    server = await asyncio.start_server(_handle, "127.0.0.1", port)
    try:
        async with asyncio.timeout(timeout_seconds):
            target = await future
    except TimeoutError as exc:
        raise AgyOAuthError("مهلت لاگین تمام شد؛ دوباره آغاز کنید.") from exc
    finally:
        server.close()
        await server.wait_closed()
    _CHALLENGES.pop(login_id, None)
    redirect_uri = challenge.redirect_uri
    return await AgyAuthClient(redirect_uri=redirect_uri).exchange(
        challenge, f"http://127.0.0.1:{port}{target}"
    )


# ---------------------------------------------------------------------------
# Account pool
# ---------------------------------------------------------------------------


@dataclass(slots=True, kw_only=True)
class AgyAccount:
    """One usable account: decrypted material plus its rotation identity."""

    credential_id: str
    email: str
    project_id: str
    refresh_token: str


@dataclass(slots=True)
class _PoolState:
    #: Monotonic rotation cursor — spreads load across the account set.
    cursor: int = 0
    #: credential_id -> epoch seconds until which the account is benched.
    benched_until: dict[str, float] = field(default_factory=dict)
    #: credential_id -> (access token, expiry epoch) — avoids a Google
    #: round-trip on every single completion.
    tokens: dict[str, tuple[str, float]] = field(default_factory=dict)


def _account_of(credential: ModelCredential) -> AgyAccount:
    """Decrypt one credential row into usable login material.

    The ciphertext is a JSON envelope written at login time; older key-shaped
    credentials (empty email) are treated as unusable for agy.
    """
    try:
        payload = json.loads(decrypt(credential.encrypted_api_key))
        return AgyAccount(
            credential_id=credential.id,
            email=str(payload.get("email") or credential.email or ""),
            project_id=str(payload.get("project_id") or credential.project_id or ""),
            refresh_token=str(payload.get("refresh_token") or ""),
        )
    except Exception:
        return AgyAccount(
            credential_id=credential.id,
            email=credential.email,
            project_id=credential.project_id,
            refresh_token="",
        )


class AgyModelAdapter:
    """ModelProviderPort over the agy account pool.

    One instance per process (the adapter registry caches it), so the rotation
    cursor and the in-memory bench are shared by every worker and every
    request. Hard state (cooldown / error / last-used) is mirrored to the
    database through the injected ``report`` callback so the background
    workers see the same picture.
    """

    def __init__(
        self,
        *,
        loader: _LoaderFn | None = None,
        report: _StatusFn | None = None,
        timeout: float = 240.0,
    ) -> None:
        self._loader = loader
        self._report = report
        self._timeout = timeout
        self._state = _PoolState()
        self._lock = asyncio.Lock()

    # -- port ---------------------------------------------------------------

    async def complete(
        self,
        *,
        model: AiModel,
        credential: ModelCredential | None = None,
        credentials: list[ModelCredential] | None = None,
        messages: list[ChatMessage],
        system: str = "",
        temperature: float = 0.2,
        max_tokens: int = 2048,
    ) -> ModelResponse:
        pool = credentials if credentials is not None else (
            [credential] if credential is not None else await self._load()
        )
        accounts = [a for a in (_account_of(c) for c in pool) if a.refresh_token]
        if not accounts:
            raise ProviderUnavailableError(
                "هیچ اکانت agy فعالی ثبت نشده است. مالک سامانه باید ابتدا از "
                "کنسول مدیریت یک اکانت Google اضافه کند."
            )
        ordered = self._order_accounts(accounts)
        failures: list[str] = []
        for account in ordered:
            try:
                response = await self._complete_with(
                    account, model=model, messages=messages, system=system,
                    temperature=temperature, max_tokens=max_tokens,
                )
            except _AgyRateLimited as exc:
                failures.append(f"{account.email or account.credential_id}: محدودیت نرخ")
                await self._bench(account, exc.retry_after)
                continue
            except _AgyAuthError as exc:
                failures.append(f"{account.email or account.credential_id}: {exc}")
                await self._bench(account, 3600, status="error", error=str(exc))
                continue
            except _AgyTransientError as exc:
                failures.append(f"{account.email or account.credential_id}: {exc}")
                continue
            await self._succeed(account)
            return response
        raise ProviderUnavailableError(
            "همهٔ اکانت‌های agy در دسترس نبودند: " + "; ".join(failures[:4])
        )

    async def health(self) -> ModelHealth:
        accounts = [a for a in (_account_of(c) for c in await self._load()) if a.refresh_token]
        if not accounts:
            return ModelHealth(ok=False, detail="هیچ اکانت agy ثبت نشده است.")
        try:
            await self._access_token(accounts[0])
        except _AgyAuthError as exc:
            return ModelHealth(ok=False, detail=str(exc))
        except Exception as exc:  # noqa: BLE001 - health must not raise
            return ModelHealth(ok=False, detail=str(exc))
        return ModelHealth(ok=True, detail=f"{len(accounts)} اکانت فعال")

    async def available_models(self) -> list[str]:
        """Models the logged-in subscription exposes (admin console helper)."""
        accounts = [a for a in (_account_of(c) for c in await self._load()) if a.refresh_token]
        if not accounts:
            return []
        return await self._fetch_models(accounts[0])

    async def available_models_for(self, credential: ModelCredential) -> list[str]:
        """Live model catalog of exactly one account.

        Called per account selection in the admin console, so the list is
        fetched fresh from Google every time (the access token is cached, the
        ``fetchAvailableModels`` round-trip is not) and reflects what *this*
        subscription is entitled to.
        """
        account = _account_of(credential)
        if not account.refresh_token:
            raise ProviderUnavailableError(
                "این اکانت اطلاعات ورود ندارد؛ یک بار دیگر لاگین کنید."
            )
        return await self._fetch_models(account)

    async def _fetch_models(self, account: AgyAccount) -> list[str]:
        token = await self._access_token(account)
        for base in BASE_URLS:
            try:
                async with httpx.AsyncClient(timeout=30.0) as client:
                    response = await client.post(
                        f"{base}{PATH_MODELS}",
                        headers={
                            "Authorization": f"Bearer {token}",
                            "Content-Type": "application/json",
                            "User-Agent": settings.agy_user_agent,
                        },
                        json={},
                    )
                if response.is_success:
                    models = response.json().get("models") or {}
                    return sorted(models)
            except httpx.HTTPError:
                continue
        return []

    # -- internals ----------------------------------------------------------

    async def _load(self) -> list[ModelCredential]:
        if self._loader is None:
            return []
        return await self._loader()

    async def _report_state(
        self,
        credential_id: str,
        *,
        status: str | None = None,
        cooldown_until: datetime | None = None,
        error: str | None = None,
        used: bool = False,
    ) -> None:
        if self._report is None:
            return
        try:
            await self._report(
                credential_id,
                status=status,
                cooldown_until=cooldown_until,
                error=error,
                used=used,
            )
        except Exception:  # noqa: BLE001 - reporting must never kill a call
            logger.debug("agy state report failed", exc_info=True)

    def _order_accounts(self, accounts: list[AgyAccount]) -> list[AgyAccount]:
        """Round-robin, skipping accounts benched in this process.

        The DB-shaped cooldown is already applied by the query that produced
        the pool; this pass additionally honours recent 429s that have not been
        persisted yet (or a reporter that is not wired).
        """
        now = time.time()
        usable = [a for a in accounts if self._state.benched_until.get(a.credential_id, 0) <= now]
        if not usable:
            # Everything is benched in memory — fall through to the DB's view
            # rather than failing a request the pool could still serve.
            usable = list(accounts)
        n = len(usable)
        self._state.cursor += 1
        start = self._state.cursor % n
        return usable[start:] + usable[:start]

    async def _bench(
        self,
        account: AgyAccount,
        retry_after: int | None,
        *,
        status: str = "cooldown",
        error: str = "",
    ) -> None:
        seconds = retry_after or settings.agy_429_cooldown_seconds
        until = datetime.now(UTC) + timedelta(seconds=seconds)
        async with self._lock:
            self._state.benched_until[account.credential_id] = until.timestamp()
        await self._report_state(
            account.credential_id, status=status, cooldown_until=until, error=error
        )

    async def _succeed(self, account: AgyAccount) -> None:
        async with self._lock:
            self._state.benched_until.pop(account.credential_id, None)
        await self._report_state(account.credential_id, status="active", used=True)

    async def _access_token(self, account: AgyAccount) -> str:
        """Cached refresh: Google tokens live ~1h; one call serves the burst."""
        cached = self._state.tokens.get(account.credential_id)
        if cached and cached[1] > time.time() + 30:
            return cached[0]
        token = await AgyAuthClient().refresh(account.refresh_token)
        async with self._lock:
            self._state.tokens[account.credential_id] = (token, time.time() + 3300)
        return token

    async def _complete_with(
        self,
        account: AgyAccount,
        *,
        model: AiModel,
        messages: list[ChatMessage],
        system: str,
        temperature: float,
        max_tokens: int,
    ) -> ModelResponse:
        token = await self._access_token(account)
        payload = _build_generate_payload(
            account, model.internal_name, messages, system, temperature, max_tokens,
            reasoning_effort=model.reasoning_effort,
        )
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
            "User-Agent": settings.agy_user_agent,
            "Accept": "application/json",
        }
        started = time.perf_counter()
        last_transient: Exception | None = None
        retry_after: int | None = None
        for base in BASE_URLS:
            try:
                async with httpx.AsyncClient(timeout=self._timeout) as client:
                    response = await client.post(
                        f"{base}{PATH_GENERATE}", json=payload, headers=headers
                    )
            except httpx.HTTPError:
                last_transient = _AgyTransientError(f"اتصال به {base} برقرار نشد.")
                continue
            if response.status_code == 429:
                # Both backends share one quota, but a 429 can also be a
                # backend hiccup — try the other base URL before benching the
                # whole account.
                retry_after = _retry_after_seconds(
                    response.headers.get("Retry-After")
                ) or retry_after
                continue
            if response.status_code in (401, 403):
                raise _AgyAuthError(
                    f"اکانت {account.email or 'agy'} دسترسی ندارد یا نشست منقضی شده است."
                )
            if response.status_code >= 500:
                last_transient = _AgyTransientError(f"{base} خطای {response.status_code} داد.")
                continue
            if response.status_code >= 400:
                # Our request shape is wrong — another account will fail the
                # same way, so surface it instead of burning the pool.
                raise ProviderUnavailableError(
                    f"agy درخواست را رد کرد ({response.status_code}): "
                    f"{response.text[:300]}"
                )
            return _parse_generate_response(
                response.json(), model=model, started=started
            )
        if retry_after is not None:
            raise _AgyRateLimited(retry_after)
        raise last_transient or _AgyTransientError("هیچ بک‌اند agy پاسخ نداد.")


# ---------------------------------------------------------------------------
# Message / payload translation (portable, pure functions — unit-tested)
# ---------------------------------------------------------------------------


def _build_generate_payload(
    account: AgyAccount,
    model_name: str,
    messages: list[ChatMessage],
    system: str,
    temperature: float,
    max_tokens: int,
    reasoning_effort: str = "",
) -> dict:
    """Translate the port's OpenAI-shaped call into the Antigravity envelope."""
    system_parts: list[str] = [system] if system else []
    contents: list[dict] = []
    for message in messages:
        if message.role == "system":
            system_parts.append(message.content)
            continue
        role = "model" if message.role == "assistant" else "user"
        contents.append({"role": role, "parts": [{"text": message.content}]})

    generation: dict = {
        "temperature": temperature,
        "maxOutputTokens": max_tokens,
    }
    # Reasoning effort is only sent to models that can think; the backend
    # rejects ``thinkingConfig`` everywhere else with a 400.
    if reasoning_effort and _supports_thinking(model_name):
        budget = EFFORT_THINKING_BUDGETS.get(reasoning_effort, 8192)
        generation["thinkingConfig"] = {
            "thinkingBudget": budget,
            "includeThoughts": True,
        }
        if max_tokens <= budget:
            # The budget must fit inside the answer window or the thinking
            # model has nowhere to emit its final answer.
            generation["maxOutputTokens"] = budget + 8000

    request: dict = {
        "contents": contents,
        "sessionId": _random_session_id(),
        "generationConfig": generation,
        "toolConfig": {"functionCallingConfig": {"mode": "VALIDATED"}},
        "safetySettings": SAFETY_SETTINGS,
    }
    if system_parts:
        request["systemInstruction"] = {
            "role": "user",
            "parts": [{"text": "\n\n".join(p for p in system_parts if p)}],
        }
    return {
        "project": account.project_id or _fallback_project_id(),
        "model": model_name,
        "userAgent": "antigravity",
        "requestId": f"agent-{uuid.uuid4()}",
        "request": request,
    }


def _parse_generate_response(
    data: dict, *, model: AiModel, started: float
) -> ModelResponse:
    candidates = data.get("candidates") or [{}]
    candidate = candidates[0] if isinstance(candidates[0], dict) else {}
    text_parts: list[str] = []
    for part in (candidate.get("content") or {}).get("parts") or []:
        if not isinstance(part, dict):
            continue
        text = part.get("text")
        if text and not part.get("thought"):
            text_parts.append(text)
    usage = data.get("usageMetadata") or {}
    latency_ms = int((time.perf_counter() - started) * 1000)
    return ModelResponse(
        content="".join(text_parts),
        model_internal_name=model.internal_name,
        input_tokens=int(usage.get("promptTokenCount", 0) or 0),
        output_tokens=int(usage.get("candidatesTokenCount", 0) or 0),
        latency_ms=latency_ms,
        raw={
            "finish_reason": candidate.get("finishReason", ""),
            "thought_tokens": int(usage.get("thoughtsTokenCount", 0) or 0),
        },
    )


# ---------------------------------------------------------------------------
# Database wiring for the registry (each call opens its own short session, so
# the adapter never outlives or depends on an HTTP request's session).
# ---------------------------------------------------------------------------


async def load_agy_accounts() -> list[ModelCredential]:
    """Every usable agy credential row, ordered least-recently-used first."""
    from app.core.database import SessionFactory
    from app.infrastructure.repositories.ai_control import (
        ModelCredentialRepository,
        ProviderRepository,
    )

    async with SessionFactory() as session:
        providers = await ProviderRepository(session).list()
        provider = next(
            (p for p in providers if p.type.value == "agy" and p.is_active), None
        )
        if provider is None:
            return []
        return await ModelCredentialRepository(session).list_active_for_provider(
            provider.id
        )


async def report_agy_state(
    credential_id: str,
    *,
    status: str | None = None,
    cooldown_until: datetime | None = None,
    error: str | None = None,
    used: bool = False,
) -> None:
    """Mirror rotation state into the DB for the workers' benefit."""
    from app.core.database import SessionFactory
    from app.infrastructure.repositories.ai_control import ModelCredentialRepository

    async with SessionFactory() as session:
        repo = ModelCredentialRepository(session)
        if status is not None:
            await repo.set_status(
                credential_id,
                status=status,
                cooldown_until=cooldown_until,
                error=error or "",
            )
        if used:
            await repo.mark_used(credential_id, when=utcnow())
        await session.commit()


__all__ = [
    "AgyAccount",
    "AgyAuthClient",
    "AgyLoginChallenge",
    "AgyLoginResult",
    "AgyModelAdapter",
    "AgyOAuthError",
    "begin_login",
    "load_agy_accounts",
    "redeem_login",
    "redeem_login_by_waiting",
    "report_agy_state",
]