app/services/google_oauth_service.py

from __future__ import annotations

import asyncio
import time
from typing import Any, Dict, Optional
from urllib.parse import urlencode

import httpx
import jwt

from app.config import get_settings
from app.core.logger import get_logger
from app.core.thread_pool import run_in_executor
from app.models.schemas import GoogleOAuthUserInfo
from app.utils.http_utils import SharedAsyncClient

_logger = get_logger(__name__)
_settings = get_settings()

# Retryable transient HTTP status codes when talking to Google.
_RETRYABLE_STATUS = frozenset({429, 500, 502, 503, 504})


class GoogleOAuthError(Exception):
    """Raised for upstream Google OAuth failures that map to a client-facing error."""

    def __init__(self, message: str, status_code: int = 400) -> None:
        super().__init__(message)
        self.message = message
        self.status_code = status_code


class GoogleJWKSCache:
    """Asynchronously fetches and caches Google's public JWKS keys.

    Keys are cached for a short TTL to avoid hammering Google's cert endpoint,
    and re-fetched on demand if the requested ``kid`` is missing (key rotation).
    """

    def __init__(self) -> None:
        self._keys: Dict[str, Dict[str, Any]] = {}
        self._last_fetched: float = 0.0
        self._lock = asyncio.Lock()

    def _is_expired(self) -> bool:
        return (
            not self._keys
            or time.time() - self._last_fetched >= _settings.google_oauth_jwks_ttl_seconds
        )

    async def get_public_key(self, client: httpx.AsyncClient, kid: str) -> Any:
        """Return the RSA public key for the given key id (``kid``)."""
        if self._is_expired():
            await self._refresh(client)

        jwk = self._keys.get(kid)
        if jwk is None:
            # Key may have rotated; force a refresh before giving up.
            await self._refresh(client)
            jwk = self._keys.get(kid)

        if jwk is None:
            raise GoogleOAuthError(
                f"Google signing key not found for kid '{kid}'", status_code=401
            )

        return jwt.algorithms.RSAAlgorithm.from_jwk(jwk)

    async def _refresh(self, client: httpx.AsyncClient) -> None:
        async with self._lock:
            if not self._is_expired():
                return
            try:
                response = await client.get(
                    _settings.google_oauth_jwks_url, timeout=10.0
                )
                response.raise_for_status()
                data = response.json()
                self._keys = {
                    key["kid"]: key
                    for key in data.get("keys", [])
                    if isinstance(key, dict) and "kid" in key
                }
                self._last_fetched = time.time()
                _logger.info("Google JWKS refreshed (%d keys)", len(self._keys))
            except httpx.HTTPError as exc:
                _logger.error("Failed to refresh Google JWKS: %s", exc)
                raise GoogleOAuthError(
                    "Unable to fetch Google signing keys", status_code=502
                ) from exc


class GoogleOAuthService:
    """Abstraction over the Google OAuth 2.0 / OpenID Connect flow."""

    def __init__(self) -> None:
        self._http = SharedAsyncClient(
            timeout=httpx.Timeout(_settings.google_oauth_timeout),
            follow_redirects=False,
        )
        self._jwks = GoogleJWKSCache()

    # ------------------------------------------------------------------
    # HTTP client management
    # ------------------------------------------------------------------

    async def _get_client(self) -> httpx.AsyncClient:
        return await self._http.get()

    async def close(self) -> None:
        await self._http.close()

    # ------------------------------------------------------------------
    # Authorization URL (Step 1)
    # ------------------------------------------------------------------

    def build_auth_url(
        self,
        *,
        client_id: str,
        redirect_uri: str,
        state: str,
        scope: str,
        prompt: Optional[str] = None,
        access_type: Optional[str] = None,
        login_hint: Optional[str] = None,
        include_granted_scopes: bool = False,
    ) -> str:
        params: Dict[str, Any] = {
            "client_id": client_id,
            "redirect_uri": redirect_uri,
            "response_type": "code",
            "scope": scope,
            "state": state,
        }
        if prompt:
            params["prompt"] = prompt
        if access_type:
            params["access_type"] = access_type
        if login_hint:
            params["login_hint"] = login_hint
        if include_granted_scopes:
            params["include_granted_scopes"] = "true"

        return f"{_settings.google_oauth_auth_url}?{urlencode(params)}"

    # ------------------------------------------------------------------
    # Token endpoint helpers
    # ------------------------------------------------------------------

    async def _post_token(
        self, data: Dict[str, str], *, context: str
    ) -> Dict[str, Any]:
        client = await self._get_client()
        last_error: Optional[str] = None

        for attempt in range(1 + _settings.google_oauth_max_retries):
            try:
                response = await client.post(_settings.google_oauth_token_url, data=data)
                if response.status_code == 200:
                    return response.json()

                if response.status_code in _RETRYABLE_STATUS:
                    last_error = (
                        f"Google OAuth upstream error HTTP {response.status_code}"
                    )
                    _logger.warning(
                        "%s: transient HTTP %s (attempt %d/%d)",
                        context,
                        response.status_code,
                        attempt + 1,
                        1 + _settings.google_oauth_max_retries,
                    )
                else:
                    return self._parse_token_error(response, context=context)

            except httpx.TimeoutException:
                last_error = "Google OAuth request timed out"
                _logger.warning(
                    "%s: timeout (attempt %d/%d)",
                    context,
                    attempt + 1,
                    1 + _settings.google_oauth_max_retries,
                )
            except httpx.RequestError as exc:
                last_error = f"Google OAuth request failed: {exc}"
                _logger.warning(
                    "%s: %s (attempt %d/%d)",
                    context,
                    last_error,
                    attempt + 1,
                    1 + _settings.google_oauth_max_retries,
                )

            if attempt < _settings.google_oauth_max_retries:
                await asyncio.sleep(2 ** attempt)

        raise GoogleOAuthError(last_error or "Unknown Google OAuth error", status_code=502)

    @staticmethod
    def _parse_token_error(
        response: httpx.Response, *, context: str
    ) -> Dict[str, Any]:
        try:
            body = response.json()
            error: str = body.get("error", "") or ""
            error_description: str = body.get("error_description", "") or ""
        except Exception:
            error = ""
            error_description = ""
            body = {}

        _logger.warning("%s: Google token error '%s': %s", context, error, error_description)

        if error == "invalid_grant":
            message = error_description or (
                "The provided code or refresh token is invalid, expired, or has been revoked."
            )
            raise GoogleOAuthError(message, status_code=401)
        if error == "invalid_client":
            raise GoogleOAuthError(
                error_description or "Invalid client_id or client_secret.", status_code=401
            )
        if error == "invalid_request":
            raise GoogleOAuthError(
                error_description or "Malformed OAuth request.", status_code=400
            )
        if error == "access_denied":
            raise GoogleOAuthError(
                error_description or "Access denied by the user.", status_code=403
            )

        # Google does not return a structured error body for some failures.
        if response.status_code >= 500:
            raise GoogleOAuthError(
                "Google OAuth service is temporarily unavailable.", status_code=502
            )
        if response.status_code == 429:
            raise GoogleOAuthError(
                "Google OAuth rate limit exceeded. Please retry later.", status_code=429
            )

        message = error_description or f"Google OAuth error (HTTP {response.status_code})."
        return {"error": True, "message": message, "raw": body}

    # ------------------------------------------------------------------
    # Code exchange (Step 2)
    # ------------------------------------------------------------------

    async def exchange_code(
        self,
        *,
        client_id: str,
        client_secret: str,
        code: str,
        redirect_uri: str,
    ) -> Dict[str, Any]:
        data = {
            "client_id": client_id,
            "client_secret": client_secret,
            "code": code,
            "grant_type": "authorization_code",
            "redirect_uri": redirect_uri,
        }
        tokens = await self._post_token(data, context="authorization code exchange")
        self._raise_for_token_error(tokens)
        return tokens

    async def refresh_access_token(
        self,
        *,
        client_id: str,
        client_secret: str,
        refresh_token: str,
    ) -> Dict[str, Any]:
        data = {
            "client_id": client_id,
            "client_secret": client_secret,
            "refresh_token": refresh_token,
            "grant_type": "refresh_token",
        }
        tokens = await self._post_token(data, context="refresh token exchange")
        self._raise_for_token_error(tokens)
        return tokens

    @staticmethod
    def _raise_for_token_error(tokens: Dict[str, Any]) -> None:
        if tokens.get("error"):
            raise GoogleOAuthError(tokens.get("message", "Unknown OAuth error"), status_code=400)
        if not tokens.get("access_token"):
            raise GoogleOAuthError(
                "Google did not return an access token.", status_code=502
            )

    # ------------------------------------------------------------------
    # ID token verification
    # ------------------------------------------------------------------

    async def verify_id_token(
        self, id_token: str, client_id: str
    ) -> GoogleOAuthUserInfo:
        client = await self._get_client()

        try:
            unverified = jwt.get_unverified_header(id_token)
            kid = unverified.get("kid")
            alg = unverified.get("alg")
        except jwt.DecodeError as exc:
            raise GoogleOAuthError("Malformed ID token.", status_code=400) from exc

        if not kid or alg != "RS256":
            raise GoogleOAuthError(
                "ID token does not use an RS256 signature.", status_code=401
            )

        try:
            public_key = await self._jwks.get_public_key(client, kid)
        except GoogleOAuthError:
            raise

        try:
            payload = await run_in_executor(
                jwt.decode,
                id_token,
                key=public_key,
                algorithms=["RS256"],
                audience=client_id,
                issuer=["accounts.google.com", "https://accounts.google.com"],
                options={
                    "verify_exp": True,
                    "verify_iat": True,
                    "verify_aud": True,
                    "verify_iss": True,
                    "require": ["exp", "iat", "sub", "email"],
                },
            )
        except jwt.ExpiredSignatureError as exc:
            raise GoogleOAuthError("ID token has expired.", status_code=401) from exc
        except jwt.InvalidAudienceError as exc:
            raise GoogleOAuthError(
                "ID token audience does not match the provided client_id.", status_code=401
            ) from exc
        except jwt.InvalidIssuerError as exc:
            raise GoogleOAuthError(
                "ID token issuer is not Google.", status_code=401
            ) from exc
        except jwt.PyJWTError as exc:
            _logger.warning("ID token verification failed: %s", exc)
            raise GoogleOAuthError("Invalid ID token.", status_code=401) from exc

        return self._user_info_from_claims(payload)

    # ------------------------------------------------------------------
    # Userinfo (fallback profile source)
    # ------------------------------------------------------------------

    @staticmethod
    def _user_info_from_claims(payload: Dict[str, Any]) -> GoogleOAuthUserInfo:
        return GoogleOAuthUserInfo(
            sub=payload.get("sub", ""),
            email=payload.get("email", ""),
            email_verified=bool(payload.get("email_verified", False)),
            name=payload.get("name"),
            given_name=payload.get("given_name"),
            family_name=payload.get("family_name"),
            picture=payload.get("picture"),
            locale=payload.get("locale"),
        )