query.py

"""
Module responsible for creating requests to and parsing responses from JustWatch.

Functions are prepared in pairs - prepare request and parse response for specific
operation. "Request" functions do no verification of input data; "parse" functions check
if returned JSON/`dict` contain `error` key. In such case a [`JustWatchApiError`]
[simplejustwatchapi.exceptions.JustWatchApiError] is raised.

All "parse" functions convert JSON returned by API into request-specific Python
[`NamedTuple`][typing.NamedTuple].
"""

from typing import Any

from simplejustwatchapi.exceptions import JustWatchApiError, JustWatchError
from simplejustwatchapi.graphql import (
    GRAPHQL_DETAILS_QUERY,
    GRAPHQL_EPISODES_QUERY,
    GRAPHQL_POPULAR_QUERY,
    GRAPHQL_PROVIDERS_QUERY,
    GRAPHQL_SEARCH_QUERY,
    GRAPHQL_SEASONS_QUERY,
    graphql_offers_for_countries_query,
)
from simplejustwatchapi.tuples import (
    Episode,
    Interactions,
    MediaEntry,
    Offer,
    OfferPackage,
    Scoring,
    StreamingCharts,
)

_DETAILS_URL = "https://justwatch.com"
_IMAGES_URL = "https://images.justwatch.com"


def prepare_search_request(
    title: str,
    country: str,
    language: str,
    count: int,
    best_only: bool,
    offset: int,
    providers: list[str] | str | None,
    min_release_year: int | None,
    max_release_year: int | None,
    object_types: list[str] | str | None,
) -> dict[str, Any]:
    """
    Prepare search request for JustWatch GraphQL API.

    Creates a `GetSearchTitles` GraphQL query.

    Country code should be two uppercase letters, however it will be auto-converted to
    uppercase. Language code is not verified.

    Meant to be used together with [`parse_search_response`]
    [simplejustwatchapi.query.parse_search_response].

    Args:
        title (str): Title to search.
        country (str): Country to search for offers.
        language (str): Language of responses.
        count (int): How many responses should be returned.
        best_only (bool): Return only best offers if `True`,
            return all offers if `False`.
        offset (int): Search results offset.
        providers (list[str] | str | None): 3-letter service identifier(s),
            or `None` for all providers.
        min_release_year (int | None): Minimum release year of returned titles.
        max_release_year (int | None): Maximum release year of returned titles.
        object_types (list[str] | str | None): Types of objects to filter for, it seems
            that only "SHOW" and "MOVIE" make sense.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    return {
        "operationName": "GetSearchTitles",
        "variables": {
            "first": count,
            "searchTitlesFilter": {
                "searchQuery": title,
                "packages": providers,
                **_filter_variables(min_release_year, max_release_year, object_types),
            },
            **_common_variables(best_only),
            **_locale_variables(country, language),
            "offset": offset or None,
        },
        "query": GRAPHQL_SEARCH_QUERY,
    }


def parse_search_response(json: dict[str, Any]) -> list[MediaEntry]:
    """
    Parse response from search query from JustWatch GraphQL API.

    Parses response for `GetSearchTitles` query.

    If API didn't return any data, then an empty list is returned.

    Meant to be used together with [`prepare_search_request`]
    [simplejustwatchapi.query.prepare_search_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.

    Returns:
        (list[MediaEntry]): Parsed received JSON as a list of `MediaEntry` NamedTuples.

    Raises:
        exceptions.JustWatchApiError: JSON response from API has internal errors.

    """
    _raise_for_errors_in_response(json)
    return [
        _parse_entry(edge["node"]) for edge in json["data"]["popularTitles"]["edges"]
    ]


def prepare_popular_request(
    country: str,
    language: str,
    count: int,
    best_only: bool,
    offset: int,
    providers: list[str] | str | None,
    min_release_year: int | None,
    max_release_year: int | None,
    object_types: list[str] | str | None,
) -> dict[str, Any]:
    """
    Prepare "get popular" request for JustWatch GraphQL API.

    Creates a `GetPopularTitles` GraphQL query.

    Country code should be two uppercase letters, however it will be auto-converted to
    uppercase. Language code is not verified.

    Meant to be used together with [`parse_popular_response`]
    [simplejustwatchapi.query.parse_popular_response].

    Args:
        country (str): Country to search for offers.
        language (str): Language of responses.
        count (int): How many responses should be returned.
        best_only (bool): Return only best offers if `True`,
            return all offers if `False`.
        offset (int): Search results offset.
        providers (list[str] | str | None): 3-letter service identifier(s),
            or `None` for all providers.
        min_release_year (int | None): Minimum release year of returned titles.
        max_release_year (int | None): Maximum release year of returned titles.
        object_types (list[str] | str | None): Types of objects to filter for, it seems
            that only "SHOW" and "MOVIE" make sense.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    return {
        "operationName": "GetPopularTitles",
        "variables": {
            "first": count,
            "popularTitlesFilter": {
                "packages": providers,
                **_filter_variables(min_release_year, max_release_year, object_types),
            },
            **_common_variables(best_only),
            **_locale_variables(country, language),
            "offset": offset or None,
        },
        "query": GRAPHQL_POPULAR_QUERY,
    }


def parse_popular_response(json: dict[str, Any]) -> list[MediaEntry]:
    """
    Parse response from the "get popular" query from JustWatch GraphQL API.

    Parses response for `GetPopularTitles` query.

    If API didn't return any data, then an empty list is returned.

    Meant to be used together with [`prepare_popular_request`]
    [simplejustwatchapi.query.prepare_popular_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.

    Returns:
        (list[MediaEntry]): Parsed received JSON as a list of `MediaEntry` NamedTuples.

    Raises:
        exceptions.JustWatchApiError: JSON response from API has internal errors.

    """
    _raise_for_errors_in_response(json)
    return [
        _parse_entry(edge["node"]) for edge in json["data"]["popularTitles"]["edges"]
    ]


def prepare_details_request(
    node_id: str, country: str, language: str, best_only: bool
) -> dict[str, Any]:
    """
    Prepare a details request for specified node ID to JustWatch GraphQL API.

    Creates a `GetTitleNode` GraphQL query.

    Country code should be two uppercase letters, however it will be auto-converted to
    uppercase. Language code is not verified.

    Meant to be used together with [`parse_details_response`]
    [simplejustwatchapi.query.parse_details_response].

    Args:
        node_id (str): Node ID of entry to get details for.
        country (str): Country to search for offers.
        language (str): Language of responses.
        best_only (bool): Return only best offers if `True`,
            return all offers if `False`.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    return {
        "operationName": "GetTitleNode",
        "variables": {
            "nodeId": node_id,
            **_common_variables(best_only),
            **_locale_variables(country, language),
        },
        "query": GRAPHQL_DETAILS_QUERY,
    }


def parse_details_response(json: dict[str, Any]) -> MediaEntry:
    """
    Parse response from details query from JustWatch GraphQL API.

    Parses response for `GetTitleNode` query.

    Meant to be used together with [`prepare_details_request`]
    [simplejustwatchapi.query.prepare_details_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.

    Returns:
        (MediaEntry): Parsed received JSON as a `MediaEntry` NamedTuple.

    Raises:
        exceptions.JustWatchApiError: JSON response from API has internal errors.

    """
    _raise_for_errors_in_response(json)
    return _parse_entry(json["data"]["node"])


def prepare_seasons_request(
    show_id: str, country: str, language: str, best_only: bool
) -> dict[str, Any]:
    """
    Prepare a seasons details request for specified show ID to JustWatch GraphQL API.

    Creates a `GetTitleNode` GraphQL query.

    Country code should be two uppercase letters, however it will be auto-converted to
    uppercase. Language code is not verified.

    Meant to be used together with [`parse_seasons_response`]
    [simplejustwatchapi.query.parse_seasons_response].

    Args:
        show_id (str): Show ID of entry to get details for.
        country (str): Country to search for offers.
        language (str): Language of responses.
        best_only (bool): Return only best offers if `True`,
            return all offers if `False`.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    return {
        "operationName": "GetTitleNode",
        "variables": {
            "nodeId": show_id,
            **_common_variables(best_only),
            **_locale_variables(country, language),
        },
        "query": GRAPHQL_SEASONS_QUERY,
    }


def parse_seasons_response(json: dict[str, Any]) -> list[MediaEntry]:
    """
    Parse response from seasons details query from JustWatch GraphQL API.

    Parses response for `GetTitleNode` query.

    Meant to be used together with [`prepare_seasons_request`]
    [simplejustwatchapi.query.prepare_seasons_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.

    Returns:
        (list[MediaEntry]): Parsed received JSON as a `MediaEntry` NamedTuple list.

    Raises:
        exceptions.JustWatchApiError: JSON response from API has internal errors.

    """
    _raise_for_errors_in_response(json)
    return [_parse_entry(season) for season in json["data"]["node"].get("seasons", [])]


def prepare_episodes_request(
    episode_id: str, country: str, language: str, best_only: bool
) -> dict[str, Any]:
    """
    Prepare a episodes details request for specified node ID to JustWatch GraphQL API.

    Creates a `GetTitleNode` GraphQL query.

    Country code should be two uppercase letters, however it will be auto-converted to
    uppercase. Language code is not verified.

    Meant to be used together with [`parse_episodes_response`]
    [simplejustwatchapi.query.parse_episodes_response].

    Args:
        episode_id (str): Episode ID of entry to get details for.
        country (str): Country to search for offers.
        language (str): Language of responses.
        best_only (bool): Return only best offers if `True`,
            return all offers if `False`.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    return {
        "operationName": "GetTitleNode",
        "variables": {
            "nodeId": episode_id,
            **_common_variables(best_only),
            **_locale_variables(country, language),
        },
        "query": GRAPHQL_EPISODES_QUERY,
    }


def parse_episodes_response(json: dict[str, Any]) -> list[Episode]:
    """
    Parse response from episodes details query from JustWatch GraphQL API.

    Parses response for `GetTitleNode` query.

    Meant to be used together with [`prepare_episodes_request`]
    [simplejustwatchapi.query.prepare_episodes_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.

    Returns:
        (list[Episode]): Parsed received JSON as a `Episode` NamedTuple list.

    """
    _raise_for_errors_in_response(json)
    return [
        _parse_episode(episode) for episode in json["data"]["node"].get("episodes", [])
    ]


def prepare_offers_for_countries_request(
    node_id: str,
    countries: set[str],
    language: str,
    best_only: bool,
) -> dict[str, Any]:
    """
    Prepare an offers request for the node ID and for all given countries.

    Creates a `GetTitleOffers` GraphQL query.

    Country codes should be two uppercase letters, however they will be auto-converted
    to uppercase. `countries` argument must not be empty. Language code is not verified.

    Meant to be used together with [`parse_offers_for_countries_response`]
    [simplejustwatchapi.query.parse_offers_for_countries_response].

    Args:
        node_id (str): Node ID of entry to get details for.
        countries (set[str]): Set of country codes to search for offers.
        language (str): Language of responses.
        best_only (bool): Return only best offers if `True`,
            return all offers if `False`.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    if not countries:
        # This should never happen, justwatch.py should take care of this.
        # If it will happen API will respond with an error due to unused variables.
        error_msg = "No country codes, should not happen!"
        raise JustWatchError(error_msg)
    countries = {country.upper() for country in countries}
    return {
        "operationName": "GetTitleOffers",
        "variables": {
            "nodeId": node_id,
            "language": language,
            **_common_variables(best_only),
        },
        "query": graphql_offers_for_countries_query(countries),
    }


def parse_offers_for_countries_response(
    json: dict[str, Any], countries: set[str]
) -> dict[str, list[Offer]]:
    """
    Parse response from offers query from JustWatch GraphQL API.

    Parses response for `GetTitleOffers` query.

    Response if searched for country codes passed as `countries` argument.
    Countries in JSON response which are not present in `countries` set will be ignored.
    If response doesn't have offers for a country, then that country still will be
    present in returned `dict`, just with an empty list as value.

    Meant to be used together with [`prepare_offers_for_countries_request`]
    [simplejustwatchapi.query.prepare_offers_for_countries_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.
        countries (set[str]): Set of country codes to look for in API response.

    Returns:
        (dict[str, list[Offer]]): A `dict`, where keys are matching `countries` argument
            and values are offers for a given country parsed from JSON response.

    Raises:
        exceptions.JustWatchApiError: JSON response from API has internal errors.

    """
    _raise_for_errors_in_response(json)
    return {
        country: list(map(_parse_offer, json["data"]["node"].get(country.upper(), [])))
        for country in countries
    }


def prepare_providers_request(country: str) -> dict[str, Any]:
    """
    Prepare "get all providers" request for JustWatch GraphQL API.

    Creates a `GetProviders` GraphQL query.

    Country code should be two uppercase letters, however it will be auto-converted to
    uppercase.

    Meant to be used together with [`parse_providers_response`]
    [simplejustwatchapi.query.parse_providers_response].

    Args:
        country (str): 2-letter country code for which providers should be looked up.

    Returns:
        (dict[str, Any]): JSON with GraphQL POST body.

    """
    return {
        "operationName": "GetProviders",
        "variables": {
            "country": country.upper(),
            "formatOfferIcon": "PNG",
        },
        "query": GRAPHQL_PROVIDERS_QUERY,
    }


def parse_providers_response(json: dict[str, Any]) -> list[OfferPackage]:
    """
    Parse response from "get all providers" query from JustWatch GraphQL API.

    Parses response for `GetProviders` query.

    If API didn't return any data, then an empty list is returned.

    Meant to be used together with [`prepare_providers_request`]
    [simplejustwatchapi.query.prepare_providers_request].

    Args:
        json (dict[str, Any]): JSON returned by JustWatch GraphQL API.

    Returns:
        (list[OfferPackage]): Parsed received JSON as a list of `OfferPackage`.

    """
    _raise_for_errors_in_response(json)
    return [_parse_package(package) for package in json["data"]["packages"]]


def _common_variables(best_only: bool) -> dict[str, Any]:
    """Return dict with variables common for queries."""
    return {
        "formatPoster": "JPG",
        "formatOfferIcon": "PNG",
        "profile": "S718",
        "backdropProfile": "S1920",
        "filter": {"bestOnly": best_only},
    }


def _locale_variables(country: str, language: str) -> dict[str, str]:
    """Return dict with variables related to locale."""
    return {"country": country.upper(), "language": language}


def _filter_variables(
    min_release_year: int | None,
    max_release_year: int | None,
    object_types: list[str] | str | None,
) -> dict[str, Any]:
    """Return dict with variables related to looking up lists of titles."""
    return {
        "includeTitlesWithoutUrl": True,
        "objectTypes": object_types,
        "releaseYear": {"min": min_release_year, "max": max_release_year},
    }


def _raise_for_errors_in_response(json: dict[str, Any]) -> None:
    """Raise JustWatchApiError if given JSON contains `errors` key."""
    if "errors" in json:
        raise JustWatchApiError(json["errors"])


def _parse_entry(json: Any) -> MediaEntry:
    entry_id = json.get("id")
    object_id = json.get("objectId")
    object_type = json.get("objectType")
    total_episode_count = json.get("totalEpisodeCount")
    content = json["content"]
    season_number = content.get("seasonNumber")
    episode_number = content.get("episodeNumber")
    title = content.get("title")
    url_field = content.get("fullPath")
    # Missing URL should be "None", not an empty string,
    # but I want to keep the MediaEntry structure unchanged for now.
    url = _DETAILS_URL + url_field if url_field else ""
    year = content.get("originalReleaseYear")
    date = content.get("originalReleaseDate")
    runtime_minutes = content.get("runtime")
    short_description = content.get("shortDescription")
    genres = [node.get("shortName") for node in content.get("genres", []) if node]
    external_ids = content.get("externalIds")
    imdb_id = external_ids.get("imdbId") if external_ids else None
    tmdb_id = external_ids.get("tmdbId") if external_ids else None
    poster_url_field = content.get("posterUrl")
    poster = _IMAGES_URL + poster_url_field if poster_url_field else None
    backdrops = [
        _IMAGES_URL + backdrop.get("backdropUrl")
        for backdrop in content.get("backdrops", [])
        if backdrop
    ]
    age_certification = content.get("ageCertification")
    scoring = _parse_scores(content.get("scoring"))
    interactions = _parse_interactions(content.get("interactions"))
    streaming_charts = _parse_streaming_charts(json)
    offers = [_parse_offer(offer) for offer in json.get("offers", []) if offer]
    total_season_count = json.get("totalSeasonCount")
    return MediaEntry(
        entry_id,
        object_id,
        object_type,
        title,
        url,
        year,
        date,
        runtime_minutes,
        short_description,
        genres,
        imdb_id,
        tmdb_id,
        poster,
        backdrops,
        age_certification,
        scoring,
        interactions,
        streaming_charts,
        offers,
        total_season_count,
        total_episode_count,
        season_number,
        episode_number,
    )


def _parse_episode(json: Any) -> Episode:
    episode_id = json.get("id")
    object_id = json.get("objectId")
    object_type = json.get("objectType")
    content = json["content"]
    title = content.get("title")
    year = content.get("originalReleaseYear")
    date = content.get("originalReleaseDate")
    runtime_minutes = content.get("runtime")
    short_description = content.get("shortDescription")
    episode_number = content.get("episodeNumber")
    season_number = content.get("seasonNumber")
    offers = [_parse_offer(offer) for offer in json.get("offers", []) if offer]
    return Episode(
        episode_id,
        object_id,
        object_type,
        title,
        year,
        date,
        runtime_minutes,
        short_description,
        episode_number,
        season_number,
        offers,
    )


def _parse_scores(json: Any) -> Scoring | None:
    if not json:
        return None
    imdb_score = json.get("imdbScore")
    imdb_votes = json.get("imdbVotes")
    tmdb_popularity = json.get("tmdbPopularity")
    tmdb_score = json.get("tmdbScore")
    tomatometer = json.get("tomatoMeter")
    certified_fresh = json.get("certifiedFresh")
    jw_rating = json.get("jwRating")
    return Scoring(
        imdb_score,
        int(imdb_votes) if imdb_votes is not None else None,
        tmdb_popularity,
        tmdb_score,
        int(tomatometer) if tomatometer is not None else None,
        certified_fresh,
        jw_rating,
    )


def _parse_interactions(json: Any) -> Interactions | None:
    if not json:
        return None
    likes = json.get("likelistAdditions")
    dislikes = json.get("dislikelistAdditions")
    return Interactions(likes, dislikes)


def _parse_streaming_charts(json: Any) -> StreamingCharts | None:
    if (
        not (chart_info := json.get("streamingCharts", {}).get("edges"))
        or not (chart_info := chart_info[0].get("streamingChartInfo"))
        # Getting final info is awkward, I think this in general can return a list when
        # searching for ranks for multiple entries. In this case, to unify searching
        # and displaying details, it's always getting single element in a list.
    ):
        return None
    rank = chart_info.get("rank")
    trend = chart_info.get("trend")
    trend_difference = chart_info.get("trendDifference")
    top_rank = chart_info.get("topRank")
    days_in_top_3 = chart_info.get("daysInTop3")
    days_in_top_10 = chart_info.get("daysInTop10")
    days_in_top_100 = chart_info.get("daysInTop100")
    days_in_top_1000 = chart_info.get("daysInTop1000")
    updated = chart_info.get("updatedAt")
    return StreamingCharts(
        rank,
        trend,
        trend_difference,
        top_rank,
        days_in_top_3,
        days_in_top_10,
        days_in_top_100,
        days_in_top_1000,
        updated,
    )


def _parse_offer(json: Any) -> Offer:
    offer_id = json.get("id")
    monetization_type = json.get("monetizationType")
    presentation_type = json.get("presentationType")
    price_string = json.get("retailPrice")
    price_value = json.get("retailPriceValue")
    price_currency = json.get("currency")
    last_change_retail_price_value = json.get("lastChangeRetailPriceValue")
    offer_type = json.get("type")
    package = _parse_package(json["package"])
    url = json.get("standardWebURL")
    element_count = json.get("elementCount")
    available_to = json.get("availableTo")
    deeplink_roku = json.get("deeplinkRoku")
    subtitle_languages = json.get("subtitleLanguages")
    video_technology = json.get("videoTechnology")
    audio_technology = json.get("audioTechnology")
    audio_languages = json.get("audioLanguages")
    return Offer(
        offer_id,
        monetization_type,
        presentation_type,
        price_string,
        price_value,
        price_currency,
        last_change_retail_price_value,
        offer_type,
        package,
        url,
        element_count,
        available_to,
        deeplink_roku,
        subtitle_languages,
        video_technology,
        audio_technology,
        audio_languages,
    )


def _parse_package(json: Any) -> OfferPackage:
    platform_id = json.get("id")
    package_id = json.get("packageId")
    name = json.get("clearName")
    technical_name = json.get("technicalName")
    short_name = json.get("shortName")
    monetization_types = json.get("monetizationTypes", [])
    icon = _IMAGES_URL + icon if (icon := json.get("icon")) else ""
    return OfferPackage(
        platform_id,
        package_id,
        name,
        technical_name,
        short_name,
        monetization_types,
        icon,
    )