unshackle-SeFree/unshackle/core/providers/tmdb.py

from __future__ import annotations

from difflib import SequenceMatcher
from typing import Optional, Union

import requests

from unshackle.core.config import config
from unshackle.core.providers._base import ExternalIds, MetadataProvider, MetadataResult, _clean, _strip_year


class TMDBProvider(MetadataProvider):
    """TMDB (The Movie Database) metadata provider."""

    NAME = "tmdb"
    REQUIRES_KEY = True
    BASE_URL = "https://api.themoviedb.org/3"

    def is_available(self) -> bool:
        return bool(config.tmdb_api_key)

    @property
    def _api_key(self) -> str:
        return config.tmdb_api_key

    def search(self, title: str, year: Optional[int], kind: str) -> Optional[MetadataResult]:
        search_title = _strip_year(title)
        self.log.debug("Searching TMDB for %r (%s, %s)", search_title, kind, year)

        params: dict[str, str | int] = {"api_key": self._api_key, "query": search_title}
        if year is not None:
            params["year" if kind == "movie" else "first_air_date_year"] = year

        try:
            r = self.session.get(f"{self.BASE_URL}/search/{kind}", params=params, timeout=30)
            r.raise_for_status()
            results = r.json().get("results") or []
            self.log.debug("TMDB returned %d results", len(results))
            if not results:
                return None
        except requests.RequestException as exc:
            self.log.warning("Failed to search TMDB for %s: %s", title, exc)
            return None

        best_ratio = 0.0
        best_id: Optional[int] = None
        best_title: Optional[str] = None
        for result in results:
            candidates = [
                result.get("title"),
                result.get("name"),
                result.get("original_title"),
                result.get("original_name"),
            ]
            candidates = [c for c in candidates if c]

            for candidate in candidates:
                ratio = SequenceMatcher(None, _clean(search_title), _clean(candidate)).ratio()
                if ratio > best_ratio:
                    best_ratio = ratio
                    best_id = result.get("id")
                    best_title = candidate

        self.log.debug("Best candidate ratio %.2f for %r (ID %s)", best_ratio, best_title, best_id)

        if best_id is None:
            first = results[0]
            best_id = first.get("id")
            best_title = first.get("title") or first.get("name")

        if best_id is None:
            return None

        # Fetch full detail for caching
        detail = self._fetch_detail(best_id, kind)
        ext_raw = self._fetch_external_ids_raw(best_id, kind)

        date = (detail or {}).get("release_date") or (detail or {}).get("first_air_date")
        result_year = int(date[:4]) if date and len(date) >= 4 and date[:4].isdigit() else None

        ext = ExternalIds(
            imdb_id=ext_raw.get("imdb_id") if ext_raw else None,
            tmdb_id=best_id,
            tmdb_kind=kind,
            tvdb_id=ext_raw.get("tvdb_id") if ext_raw else None,
        )

        return MetadataResult(
            title=best_title,
            year=result_year,
            kind=kind,
            external_ids=ext,
            source="tmdb",
            raw={"detail": detail or {}, "external_ids": ext_raw or {}},
        )

    def get_by_id(self, provider_id: Union[int, str], kind: str) -> Optional[MetadataResult]:
        detail = self._fetch_detail(int(provider_id), kind)
        if not detail:
            return None

        title = detail.get("title") or detail.get("name")
        date = detail.get("release_date") or detail.get("first_air_date")
        year = int(date[:4]) if date and len(date) >= 4 and date[:4].isdigit() else None

        return MetadataResult(
            title=title,
            year=year,
            kind=kind,
            external_ids=ExternalIds(tmdb_id=int(provider_id), tmdb_kind=kind),
            source="tmdb",
            raw=detail,
        )

    def get_external_ids(self, provider_id: Union[int, str], kind: str) -> ExternalIds:
        raw = self._fetch_external_ids_raw(int(provider_id), kind)
        if not raw:
            return ExternalIds(tmdb_id=int(provider_id), tmdb_kind=kind)
        return ExternalIds(
            imdb_id=raw.get("imdb_id"),
            tmdb_id=int(provider_id),
            tmdb_kind=kind,
            tvdb_id=raw.get("tvdb_id"),
        )

    def find_by_imdb_id(self, imdb_id: str, kind: str) -> Optional[ExternalIds]:
        """Look up TMDB/TVDB IDs from an IMDB ID using TMDB's /find endpoint."""
        self.log.debug("Looking up IMDB ID %s on TMDB", imdb_id)
        try:
            r = self.session.get(
                f"{self.BASE_URL}/find/{imdb_id}",
                params={"api_key": self._api_key, "external_source": "imdb_id"},
                timeout=30,
            )
            r.raise_for_status()
            data = r.json()
        except requests.RequestException as exc:
            self.log.debug("TMDB find by IMDB ID failed: %s", exc)
            return None

        # Check movie_results or tv_results based on kind
        if kind == "movie":
            results = data.get("movie_results") or []
        else:
            results = data.get("tv_results") or []

        if not results:
            # Try the other type as fallback
            fallback_key = "tv_results" if kind == "movie" else "movie_results"
            results = data.get(fallback_key) or []
            if results:
                kind = "tv" if kind == "movie" else "movie"

        if not results:
            self.log.debug("No TMDB results found for IMDB ID %s", imdb_id)
            return None

        match = results[0]
        tmdb_id = match.get("id")
        if not tmdb_id:
            return None

        self.log.debug("TMDB find -> ID %s (%s) for IMDB %s", tmdb_id, kind, imdb_id)

        # Now fetch the full external IDs from TMDB to get TVDB etc.
        ext_raw = self._fetch_external_ids_raw(tmdb_id, kind)

        return ExternalIds(
            imdb_id=imdb_id,
            tmdb_id=tmdb_id,
            tmdb_kind=kind,
            tvdb_id=ext_raw.get("tvdb_id") if ext_raw else None,
        )

    def _fetch_detail(self, tmdb_id: int, kind: str) -> Optional[dict]:
        try:
            r = self.session.get(
                f"{self.BASE_URL}/{kind}/{tmdb_id}",
                params={"api_key": self._api_key},
                timeout=30,
            )
            r.raise_for_status()
            return r.json()
        except requests.RequestException as exc:
            self.log.debug("Failed to fetch TMDB detail: %s", exc)
            return None

    def _fetch_external_ids_raw(self, tmdb_id: int, kind: str) -> Optional[dict]:
        try:
            r = self.session.get(
                f"{self.BASE_URL}/{kind}/{tmdb_id}/external_ids",
                params={"api_key": self._api_key},
                timeout=30,
            )
            r.raise_for_status()
            return r.json()
        except requests.RequestException as exc:
            self.log.debug("Failed to fetch TMDB external IDs: %s", exc)
            return None