"""Manage short-lived Vertex AI access tokens for EVALUATOR users.

Tokens are minted by the proxy (which impersonates a dedicated SA) and
cached locally. The manager handles automatic refresh before expiry so
that LLM and embedding calls always have a valid token.
"""

import logging
import time
from dataclasses import dataclass, field
from typing import Optional

from google.oauth2.credentials import Credentials as OAuth2Credentials

from puli_mcp_server.proxy_client import ProxyClient

logger = logging.getLogger(__name__)

# Refresh the token when it has less than this many seconds left.
_REFRESH_BUFFER_SECONDS = 5 * 60  # 5 minutes


@dataclass
class VertexTokenManager:
    """Fetches and caches short-lived Vertex AI OAuth2 access tokens.

    Usage::

        manager = VertexTokenManager(proxy_client)
        creds = manager.get_credentials()  # google.oauth2.credentials.Credentials
    """

    proxy_client: ProxyClient
    _access_token: Optional[str] = field(default=None, repr=False)
    _expires_at: int = 0

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #

    def get_credentials(self) -> OAuth2Credentials:
        """Return a ``google.oauth2.credentials.Credentials`` with a valid token.

        Automatically refreshes the token if it is expired or about to expire.
        The returned object can be passed directly to:
          - ``google.genai.Client(credentials=...)``
          - ``pydantic_ai.providers.google.GoogleProvider(credentials=...)``
          - ``anthropic.AsyncAnthropicVertex`` (via its ``access_token`` kwarg)
        """
        if self._needs_refresh():
            self._refresh()
        return OAuth2Credentials(token=self._access_token)

    @property
    def access_token(self) -> str:
        """Return the raw access token string (refreshes if needed)."""
        if self._needs_refresh():
            self._refresh()
        assert self._access_token is not None
        return self._access_token

    # ------------------------------------------------------------------ #
    # Internal
    # ------------------------------------------------------------------ #

    def _needs_refresh(self) -> bool:
        if self._access_token is None:
            return True
        return time.time() >= (self._expires_at - _REFRESH_BUFFER_SECONDS)

    def _refresh(self) -> None:
        logger.info("Requesting new Vertex AI access token from proxy…")
        resp = self.proxy_client.request_vertex_token()
        self._access_token = resp["access_token"]
        self._expires_at = int(resp["expires_at"])
        ttl = self._expires_at - int(time.time())
        logger.info("Vertex AI token refreshed (TTL %ds)", ttl)

