archived-MoviePilot-Plugins/plugins.v2/imdbsource/officialapi.py

import re
from textwrap import dedent
from typing import Any, Dict, List, Optional, Final, AsyncGenerator

import httpx
import requests
from pydantic import ValidationError

from app.core.cache import cached
from app.log import logger
from app.utils.common import retry
from app.utils.http import RequestUtils, AsyncRequestUtils

from .schema.imdbtypes import ImdbType
from .schema import VerticalList, AdvancedTitleSearchResponse, AdvancedTitleSearch, TitleEdge, SearchParams

INTERESTS_ID: Final[Dict[str, Dict[str, str]]] = {
    "Action": {
        "Action": "in0000001",
        "Action Epic": "in0000002",
        "B-Action": "in0000003",
        "Car Action": "in0000004",
        "Disaster": "in0000005",
        "Gun Fu": "in0000197",
        "Kung Fu": "in0000198",
        "Martial Arts": "in0000006",
        "One-Person Army Action": "in0000007",
        "Samurai": "in0000199",
        "Superhero": "in0000008",
        "Sword & Sandal": "in0000009",
        "War": "in0000010",
        "War Epic": "in0000011",
        "Wuxia": "in0000200"
    },
    "Adventure": {
        "Adventure": "in0000012",
        "Adventure Epic": "in0000015",
        "Desert Adventure": "in0000013",
        "Dinosaur Adventure": "in0000014",
        "Globetrotting Adventure": "in0000016",
        "Jungle Adventure": "in0000017",
        "Mountain Adventure": "in0000018",
        "Quest": "in0000019",
        "Road Trip": "in0000020",
        "Sea Adventure": "in0000021",
        "Swashbuckler": "in0000022",
        "Teen Adventure": "in0000023",
        "Urban Adventure": "in0000024"
    },
    "Animation": {
        "Adult Animation": "in0000025",
        "Animation": "in0000026",
        "Computer Animation": "in0000028",
        "Hand-Drawn Animation": "in0000029",
        "Stop Motion Animation": "in0000030"
    },
    "Anime": {
        "Anime": "in0000027",
        "Isekai": "in0000201",
        "Iyashikei": "in0000202",
        "Josei": "in0000203",
        "Mecha": "in0000204",
        "Seinen": "in0000205",
        "Shōjo": "in0000207",
        "Shōnen": "in0000206",
        "Slice of Life": "in0000208"
    },
    "Comedy": {
        "Body Swap Comedy": "in0000031",
        "Buddy Comedy": "in0000032",
        "Buddy Cop": "in0000033",
        "Comedy": "in0000034",
        "Dark Comedy": "in0000035",
        "Farce": "in0000036",
        "High-Concept Comedy": "in0000037",
        "Mockumentary": "in0000038",
        "Parody": "in0000039",
        "Quirky Comedy": "in0000040",
        "Raunchy Comedy": "in0000041",
        "Satire": "in0000042",
        "Screwball Comedy": "in0000043",
        "Sitcom": "in0000044",
        "Sketch Comedy": "in0000045",
        "Slapstick": "in0000046",
        "Stand-Up": "in0000047",
        "Stoner Comedy": "in0000048",
        "Teen Comedy": "in0000049"
    },
    "Crime": {
        "Caper": "in0000050",
        "Cop Drama": "in0000051",
        "Crime": "in0000052",
        "Drug Crime": "in0000053",
        "Film Noir": "in0000054",
        "Gangster": "in0000055",
        "Heist": "in0000056",
        "Police Procedural": "in0000057",
        "True Crime": "in0000058"
    },
    "Documentary": {
        "Crime Documentary": "in0000059",
        "Documentary": "in0000060",
        "Docuseries": "in0000061",
        "Faith & Spirituality Documentary": "in0000062",
        "Food Documentary": "in0000063",
        "History Documentary": "in0000064",
        "Military Documentary": "in0000065",
        "Music Documentary": "in0000066",
        "Nature Documentary": "in0000067",
        "Political Documentary": "in0000068",
        "Science & Technology Documentary": "in0000069",
        "Sports Documentary": "in0000070",
        "Travel Documentary": "in0000071"
    },
    "Drama": {
        "Biography": "in0000072",
        "Coming-of-Age": "in0000073",
        "Costume Drama": "in0000074",
        "Docudrama": "in0000075",
        "Drama": "in0000076",
        "Epic": "in0000077",
        "Financial Drama": "in0000078",
        "Historical Epic": "in0000079",
        "History": "in0000080",
        "Korean Drama": "in0000209",
        "Legal Drama": "in0000081",
        "Medical Drama": "in0000082",
        "Period Drama": "in0000083",
        "Political Drama": "in0000084",
        "Prison Drama": "in0000085",
        "Psychological Drama": "in0000086",
        "Showbiz Drama": "in0000087",
        "Soap Opera": "in0000088",
        "Teen Drama": "in0000089",
        "Telenovela": "in0000210",
        "Tragedy": "in0000090",
        "Workplace Drama": "in0000091"
    },
    "Family": {
        "Animal Adventure": "in0000092",
        "Family": "in0000093"
    },
    "Fantasy": {
        "Dark Fantasy": "in0000095",
        "Fairy Tale": "in0000097",
        "Fantasy": "in0000098",
        "Fantasy Epic": "in0000096",
        "Supernatural Fantasy": "in0000099",
        "Sword & Sorcery": "in0000100",
        "Teen Fantasy": "in0000101"
    },
    "Game Show": {
        "Beauty Competition": "in0000102",
        "Cooking Competition": "in0000103",
        "Game Show": "in0000105",
        "Quiz Show": "in0000104",
        "Survival Competition": "in0000106",
        "Talent Competition": "in0000107"
    },
    "Horror": {
        "B-Horror": "in0000108",
        "Body Horror": "in0000109",
        "Folk Horror": "in0000110",
        "Found Footage Horror": "in0000111",
        "Horror": "in0000112",
        "Monster Horror": "in0000113",
        "Psychological Horror": "in0000114",
        "Slasher Horror": "in0000115",
        "Splatter Horror": "in0000116",
        "Supernatural Horror": "in0000117",
        "Teen Horror": "in0000118",
        "Vampire Horror": "in0000119",
        "Werewolf Horror": "in0000120",
        "Witch Horror": "in0000121",
        "Zombie Horror": "in0000122"
    },
    "Lifestyle": {
        "Beauty Makeover": "in0000123",
        "Cooking & Food": "in0000124",
        "Home Improvement": "in0000125",
        "Lifestyle": "in0000126",
        "News": "in0000211",
        "Talk Show": "in0000127",
        "Travel": "in0000128"
    },
    "Music": {
        "Concert": "in0000129",
        "Music": "in0000130"
    },
    "Musical": {
        "Classic Musical": "in0000131",
        "Jukebox Musical": "in0000132",
        "Musical": "in0000133",
        "Pop Musical": "in0000134",
        "Rock Musical": "in0000135"
    },
    "Mystery": {
        "Bumbling Detective": "in0000136",
        "Cozy Mystery": "in0000137",
        "Hard-boiled Detective": "in0000138",
        "Mystery": "in0000139",
        "Suspense Mystery": "in0000140",
        "Whodunnit": "in0000141"
    },
    "Reality TV": {
        "Business Reality TV": "in0000142",
        "Crime Reality TV": "in0000143",
        "Dating Reality TV": "in0000144",
        "Docusoap Reality TV": "in0000145",
        "Hidden Camera": "in0000146",
        "Paranormal Reality TV": "in0000147",
        "Reality TV": "in0000148"
    },
    "Romance": {
        "Dark Romance": "in0000149",
        "Feel-Good Romance": "in0000151",
        "Romance": "in0000152",
        "Romantic Comedy": "in0000153",
        "Romantic Epic": "in0000150",
        "Steamy Romance": "in0000154",
        "Teen Romance": "in0000155",
        "Tragic Romance": "in0000156"
    },
    "Sci-Fi": {
        "Alien Invasion": "in0000157",
        "Artificial Intelligence": "in0000158",
        "Cyberpunk": "in0000159",
        "Dystopian Sci-Fi": "in0000160",
        "Kaiju": "in0000161",
        "Sci-Fi": "in0000162",
        "Sci-Fi Epic": "in0000163",
        "Space Sci-Fi": "in0000164",
        "Steampunk": "in0000165",
        "Time Travel": "in0000166"
    },
    "Seasonal": {
        "Holiday": "in0000192",
        "Holiday Animation": "in0000193",
        "Holiday Comedy": "in0000194",
        "Holiday Family": "in0000195",
        "Holiday Romance": "in0000196"
    },
    "Short": {
        "Short": "in0000212"
    },
    "Sport": {
        "Baseball": "in0000167",
        "Basketball": "in0000168",
        "Boxing": "in0000169",
        "Extreme Sport": "in0000170",
        "Football": "in0000171",
        "Motorsport": "in0000172",
        "Soccer": "in0000173",
        "Sport": "in0000174",
        "Water Sport": "in0000175"
    },
    "Thriller": {
        "Conspiracy Thriller": "in0000176",
        "Cyber Thriller": "in0000177",
        "Erotic Thriller": "in0000178",
        "Giallo": "in0000179",
        "Legal Thriller": "in0000180",
        "Political Thriller": "in0000181",
        "Psychological Thriller": "in0000182",
        "Serial Killer": "in0000183",
        "Spy": "in0000184",
        "Survival": "in0000185",
        "Thriller": "in0000186"
    },
    "Western": {
        "Classical Western": "in0000187",
        "Contemporary Western": "in0000188",
        "Spaghetti Western": "in0000190",
        "Western": "in0000191",
        "Western Epic": "in0000189"
    }
}

COMPANY_ID = {
    "20th Century Fox": ["co0000756", "co0176225", "co0201557", "co0017497"],
    "DreamWorks": ["co0067641", "co0040938", "co0252576", "co0003158"],
    "MGM": ["co0007143", "co0026841"],
    "Paramount": ["co0023400"],
    "Sony": ["co0050868", "co0026545", "co0121181"],
    "Universal": ["co0005073", "co0055277", "co0042399"],
    "Walt Disney": ["co0008970", "co0017902", "co0098836", "co0059516", "co0092035", "co0049348"],
    "Warner Bros.": ["co0002663", "co0005035", "co0863266", "co0072876", "co0080422", "co0046718"],
    "HBO": ["co0008693", "co0754095", "co0306346", "co0148466", "co0909975", "co0638197", "co0391378"],
    "Netflix": ["co0144901", "co0805756"],
    "Hulu": ["co0218858", "co0381648"],
    "Amazon Prime Video": ["co0476953", "co1160313", "co0939864", "co0931938"],
    "Apple TV": ["co0931939", "co0546168"],
    "British Broadcasting Corporation (BBC)": ['co0043107'],
    "Tencent Video": ["co0487058"],
    "Youku": ["co0264223"],
    "iQIYI": ["co0493506", "co0691262"],
    "China Central Television (CCTV)": ['co0001524'],
    "Huayi Brothers Media": ["co0099734"],
    "Beijing Enlight Pictures": ["co0208796"],
    "Bona Film Group": ["co0452101"],
}

CACHE_LIFETIME: Final[int] = 86400
IMDB_GRAPHQL_QUERY: Final[str] = dedent("""
    query VerticalListPageItems( $titles: [ID!]! $names: [ID!]! $images: [ID!]! $videos: [ID!]!) {
      titles(ids: $titles) { ...TitleParts meterRanking { currentRank meterType rankChange {changeDirection difference} } ratingsSummary { aggregateRating voteCount} }
      names(ids: $names) { ...NameParts }
      videos(ids: $videos) { ...VideoParts }
      images(ids: $images) { ...ImageParts }
    }
    fragment TitleParts on Title {
      id
      titleText { text }
      titleType { id }
      releaseYear { year }
      akas(first: 50) { edges { node { text country { id text } language { text } } } }
      plot { plotText {plainText}}
      primaryImage { id url width height }
      releaseDate {day month year}
      titleGenres {genres {genre { text }}}
      certificate { rating }
      originalTitleText{ text }
      runtime { seconds }
    }
    fragment NameParts on Name {
      id
      nameText { text }
      primaryImage { id url width height }
    }
    fragment ImageParts on Image {
      id
      height
      width
      url
    }
    fragment VideoParts on Video {
      id
      name { value }
      contentType { displayName { value } id }
      previewURLs { displayName { value } url videoDefinition videoMimeType }
      playbackURLs { displayName { value } url videoDefinition videoMimeType }
      thumbnail { height url width }
    }
""")


class PersistedQueryNotFound(Exception):
    def __init__(self, message: str, code: int = None):
        super().__init__(message)
        self.code = code


class OfficialApiClient:
    BASE_URL = "https://caching.graphql.imdb.com/"

    def __init__(self, proxies: Optional[Dict[str, str]] = None,
                 ua: Optional[str] = None):
        self._req = RequestUtils(accept_type="application/json",
                                 content_type="application/json",
                                 timeout=10,
                                 ua=ua,
                                 proxies=proxies,
                                 session=requests.Session())
        if proxies:
            proxy_url = proxies.get("https") or proxies.get("http")
        else:
            proxy_url = None
        self._client = httpx.AsyncClient(timeout=10, proxy=proxy_url)
        self._async_req = AsyncRequestUtils(accept_type="application/json", content_type="application/json",
                                            client=self._client, ua=ua)
        self.flat_interest_id = {}
        for category, value in INTERESTS_ID.items():
            for name, in_id in value.items():
                self.flat_interest_id[name] = in_id

    @cached(maxsize=1024, ttl=CACHE_LIFETIME)
    async def _async_request(self, params: Dict[str, Any], sha256: str) -> Optional[Dict]:
        params["extensions"] = {"persistedQuery": {"sha256Hash": sha256, "version": 1}}
        data = await self._async_req.post_json(f"{self.BASE_URL}", json=params, raise_exception=True)
        if not data:
            return None
        if "errors" in data:
            error = data.get("errors")[0] if data.get("errors") else {}
            return {'error': error}
        return data.get("data")

    @retry(Exception, logger=logger, delay=1)
    @cached(maxsize=1024, ttl=CACHE_LIFETIME)
    def _query_graphql(self, query: str, variables: Dict[str, Any]) -> Optional[dict]:
        params = {'query': query, 'variables': variables}
        data = self._req.post_json(f"{self.BASE_URL}", json=params, raise_exception=True)
        if not data:
            return {'error': 'Query failed.'}
        if "errors" in data:
            error = data.get("errors")[0] if data.get("errors") else {}
            return {'error': error}
        return data.get("data")

    @retry(Exception, logger=logger, delay=1)
    @cached(maxsize=1024, ttl=CACHE_LIFETIME)
    async def _async_query_graphql(self, query: str, variables: Dict[str, Any]) -> Optional[Dict]:
        params = {'query': query, 'variables': variables}
        data = await self._async_req.post_json(f"{self.BASE_URL}", json=params, raise_exception=True)
        if not data:
            return None
        if "errors" in data:
            error = data.get("errors")[0] if data.get("errors") else {}
            return {'error': error}
        return data.get("data")

    @cached(maxsize=1024, ttl=CACHE_LIFETIME)
    def vertical_list_page_items(self,
                                 titles: Optional[List[str]] = None,
                                 names: Optional[List[str]] = None,
                                 images: Optional[List[str]] = None,
                                 videos: Optional[List[str]] = None,
                                 is_registered: bool = False
                                 ) -> Optional[VerticalList]:
        variables = {'images': images or [],
                     'titles': titles or [],
                     'names': names or [],
                     'videos': videos or [],
                     'isRegistered': is_registered,
                     }
        try:
            data = self._query_graphql(IMDB_GRAPHQL_QUERY, variables)
            if 'error' in data:
                error = data['error']
                if error:
                    logger.error(f"Error querying VerticalListPageItems: {error}")
                return None
            ret = VerticalList.model_validate(data)
        except Exception as e:
            logger.debug(f"An error occurred while querying VerticalListPageItems: {e}")
            return None

        return ret

    @cached(maxsize=1024, ttl=CACHE_LIFETIME)
    async def async_vertical_list_page_items(self,
                                             titles: Optional[List[str]] = None,
                                             names: Optional[List[str]] = None,
                                             images: Optional[List[str]] = None,
                                             videos: Optional[List[str]] = None,
                                             is_registered: bool = False
                                             ) -> Optional[VerticalList]:
        variables = {'images': images or [],
                     'titles': titles or [],
                     'names': names or [],
                     'videos': videos or [],
                     'isRegistered': is_registered,
                     }
        try:
            data = await self._async_query_graphql(IMDB_GRAPHQL_QUERY, variables)
            if 'error' in data:
                error = data['error']
                if error:
                    logger.error(f"Error querying VerticalListPageItems: {error}")
                return None
            ret = VerticalList.model_validate(data)
        except Exception as e:
            logger.debug(f"An error occurred while querying VerticalListPageItems: {e}")
            return None

        return ret

    @retry(Exception, logger=logger, delay=1)
    async def async_advanced_title_search(self,
                                          params: SearchParams,
                                          sha256: str,
                                          last_cursor: Optional[str] = None,
                                          ) -> Optional[AdvancedTitleSearch]:

        variables: Dict[str, Any] = {"first": 50,
                                     "locale": "en-US",
                                     "sortBy": params.sort_by,
                                     "sortOrder": params.sort_order,
                                     }
        operation_name = 'AdvancedTitleSearch'
        if params.title_types:
            title_type_ids = []
            for title_type in params.title_types:
                if title_type in ImdbType._value2member_map_:
                    title_type_ids.append(title_type)
            if len(title_type_ids):
                variables["titleTypeConstraint"] = {"anyTitleTypeIds": title_type_ids}
        if params.genres:
            variables["genreConstraint"] = {"allGenreIds": params.genres, "excludeGenreIds": []}
        if params.countries:
            variables["originCountryConstraint"] = {"allCountries": params.countries}
        if params.languages:
            variables["languageConstraint"] = {"anyPrimaryLanguages": params.languages}
        if params.rating_min or params.rating_max:
            rating_min = params.rating_min if params.rating_min else 1
            rating_min = max(rating_min, 1)
            rating_max = params.rating_max if params.rating_max else 10
            rating_max = min(rating_max, 10)
            variables["userRatingsConstraint"] = {"aggregateRatingRange": {"max": rating_max, "min": rating_min}}
        if params.release_date_start or params.release_date_end:
            release_dict = {}
            if params.release_date_start:
                release_dict["start"] = params.release_date_start
            if params.release_date_end:
                release_dict["end"] = params.release_date_end
            variables["releaseDateConstraint"] = {"releaseDateRange": release_dict}
        if params.award_constraint:
            constraints = []
            for award in params.award_constraint:
                c = self._award_to_constraint(award)
                if c:
                    constraints.append(c)
            variables["awardConstraint"] = {"allEventNominations": constraints}
        if params.ranked:
            constraints = []
            for r in params.ranked:
                c = OfficialApiClient._ranked_list_to_constraint(r)
                if c:
                    constraints.append(c)
            variables["rankedTitleListConstraint"] = {"allRankedTitleLists": constraints,
                                                      "excludeRankedTitleLists": []}
        if params.interests:
            constraints = []
            for interest in params.interests:
                in_id = self.flat_interest_id.get(interest)
                if in_id:
                    constraints.append(in_id)
            variables["interestConstraint"] = {"allInterestIds": constraints, "excludeInterestIds": []}

        if params.company:
            company_ids = COMPANY_ID.get(params.company)
            if company_ids:
                variables["creditedCompanyConstraint"] = {"anyCompanyIds": company_ids, "excludeCompanyIds": []}

        if last_cursor:
            variables["after"] = last_cursor

        params = {"operationName": operation_name,
                  "variables": variables}
        data = await self._async_request(params, sha256)
        if not data:
            return None
        if 'error' in data:
            error = data['error']
            if error:
                if error.get('message') == 'PersistedQueryNotFound':
                    await self._async_request.cache_clear()
                    raise PersistedQueryNotFound(error['message'])
            return None
        try:
            ret = AdvancedTitleSearchResponse.model_validate(data)
        except ValidationError as err:
            logger.error(f"{err}")
            return None
        return ret.advanced_title_search

    async def advanced_title_search_generator(self, params: SearchParams, sha256: str) -> AsyncGenerator[
        TitleEdge, None]:
        last_cursor = None
        while True:
            response = await self.async_advanced_title_search(params, sha256, last_cursor=last_cursor)
            if not response:
                return

            for edge in response.edges:
                yield edge

            last_cursor = response.page_info.end_cursor
            if not last_cursor or not response.page_info.has_next_page:
                break

    @staticmethod
    def _ranked_list_to_constraint(ranked: str) -> Optional[Dict]:
        """
            "TOP_RATED_MOVIES-100": "IMDb Top 100",
            "TOP_RATED_MOVIES-250": "IMDb Top 250",
            "TOP_RATED_MOVIES-1000": "IMDb Top 1000",
            "LOWEST_RATED_MOVIES-100": "IMDb Bottom 100",
            "LOWEST_RATED_MOVIES-250": "IMDb Bottom 250",
            "LOWEST_RATED_MOVIES-1000": "IMDb Bottom 1000"
        """
        pattern = r'^(TOP_RATED_MOVIES|LOWEST_RATED_MOVIES)-(\d+)$'
        match = re.match(pattern, ranked)
        if match:
            ranked_title_list_type = match.group(1)
            rank_range = int(match.group(2))
            constraint = {"rankRange": {"max": rank_range}, "rankedTitleListType": ranked_title_list_type}
            return constraint
        return None

    @staticmethod
    def _award_to_constraint(award: str) -> Optional[Dict]:
        pattern = r'^(ev\d+)(?:-(best\w+))?-(Winning|Nominated)$'
        match = re.match(pattern, award)
        constraint = {}
        if match:
            # 第一部分：evXXXXXXXX
            ev_id = match.group(1)
            # 第二部分：bestXX（可选）
            best = match.group(2)
            # 第三部分：Winning/Nominated
            status = match.group(3)
            constraint["eventId"] = ev_id
            if status == "Winning":
                constraint["winnerFilter"] = "WINNER_ONLY"
            if best:
                constraint["searchAwardCategoryId"] = best
            return constraint
        else:
            return None

    @property
    def interests_id(self) -> Dict[str, str]:
        return self.flat_interest_id