Aurora-Core/app/providers/tmdb.py

"""TMDb metadata provider adapter.

Requires the `TMDB_API_KEY` environment variable (a TMDb v3 API key —
https://www.themoviedb.org/settings/api). If it is not configured, every
method raises `ProviderNotConfiguredError` rather than returning fake data.
"""

from __future__ import annotations

from typing import Any

import httpx
from tenacity import (
    RetryCallState,
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from app.core.config import get_settings
from app.providers.base import (
    MetadataProvider,
    ProviderError,
    ProviderNotConfiguredError,
    ProviderNotFoundError,
    ProviderRateLimitedError,
)

_MOVIE_APPEND = "credits,images,videos,keywords,alternative_titles,external_ids,release_dates,recommendations"
_SERIES_APPEND = "credits,images,videos,keywords,alternative_titles,external_ids,content_ratings"
_PERSON_APPEND = "external_ids,images,combined_credits"

_DEFAULT_EXPONENTIAL_WAIT = wait_exponential(multiplier=0.5, max=5)


def _tmdb_wait(retry_state: RetryCallState) -> float:
    """A 429's own `Retry-After` (carried on ProviderRateLimitedError) takes
    priority over guessing via exponential backoff -- TMDb told us exactly
    how long to wait, so use that instead of a generic schedule."""

    exc = retry_state.outcome.exception() if retry_state.outcome else None
    if isinstance(exc, ProviderRateLimitedError):
        return exc.retry_after
    return _DEFAULT_EXPONENTIAL_WAIT(retry_state)


class TMDbProvider(MetadataProvider):
    slug = "tmdb"
    name = "The Movie Database (TMDb)"
    base_url = "https://api.themoviedb.org/3"

    def __init__(self, api_key: str | None = None, client: httpx.AsyncClient | None = None) -> None:
        settings = get_settings()
        self._api_key = api_key if api_key is not None else settings.TMDB_API_KEY
        self._client = client or httpx.AsyncClient(base_url=self.base_url, timeout=15.0)
        self._owns_client = client is None

    def _require_api_key(self) -> None:
        if not self._api_key:
            raise ProviderNotConfiguredError(
                "TMDB_API_KEY is not set. Configure it in the environment to use the TMDb "
                "provider — see .env.example."
            )

    @retry(
        reraise=True,
        stop=stop_after_attempt(5),
        wait=_tmdb_wait,
        retry=retry_if_exception_type((httpx.TransportError, ProviderRateLimitedError)),
    )
    async def _get(self, path: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
        self._require_api_key()
        query = {"api_key": self._api_key, **(params or {})}
        try:
            response = await self._client.get(path, params=query)
        except httpx.TransportError as exc:
            raise ProviderError(f"TMDb request failed: {exc}") from exc

        if response.status_code == 404:
            raise ProviderNotFoundError(f"TMDb has no record at {path}")
        if response.status_code == 401:
            raise ProviderNotConfiguredError("TMDb rejected the configured TMDB_API_KEY.")
        if response.status_code == 429:
            try:
                # Retry-After is defined (RFC 9110) as either a delay in
                # seconds or an HTTP-date; TMDb sends seconds, but fall back
                # rather than crash if that ever changes.
                retry_after = float(response.headers.get("Retry-After", "1"))
            except ValueError:
                retry_after = 1.0
            raise ProviderRateLimitedError(
                f"TMDb rate-limited {path}", retry_after=retry_after
            )
        if response.status_code >= 400:
            raise ProviderError(f"TMDb error {response.status_code} at {path}: {response.text[:300]}")
        return response.json()

    async def get_movie(self, external_id: str) -> dict[str, Any]:
        return await self._get(f"/movie/{external_id}", {"append_to_response": _MOVIE_APPEND})

    async def get_series(self, external_id: str) -> dict[str, Any]:
        return await self._get(f"/tv/{external_id}", {"append_to_response": _SERIES_APPEND})

    async def get_season(self, series_external_id: str, season_number: int) -> dict[str, Any]:
        return await self._get(
            f"/tv/{series_external_id}/season/{season_number}",
            {"append_to_response": "credits,images,external_ids"},
        )

    async def get_person(self, external_id: str) -> dict[str, Any]:
        return await self._get(f"/person/{external_id}", {"append_to_response": _PERSON_APPEND})

    async def get_collection(self, external_id: str) -> dict[str, Any]:
        return await self._get(f"/collection/{external_id}")

    async def search_movie(self, query: str) -> list[dict[str, Any]]:
        data = await self._get("/search/movie", {"query": query})
        return data.get("results", [])

    async def aclose(self) -> None:
        if self._owns_client:
            await self._client.aclose()