FastAnime/fastanime/libs/api/anilist/mapper.py

import logging
from datetime import datetime
from typing import List, Optional

from ....core.utils.formatting import renumber_titles, strip_original_episode_prefix
from ..types import (
    AiringSchedule,
    MediaImage,
    MediaItem,
    MediaSearchResult,
    MediaStatus,
    MediaTagItem,
    MediaTitle,
    MediaTrailer,
    PageInfo,
    StreamingEpisode,
    Studio,
    UserListItem,
    UserMediaListStatus,
    UserProfile,
)
from .types import (
    AnilistBaseMediaDataSchema,
    AnilistCurrentlyLoggedInUser,
    AnilistDataSchema,
    AnilistDateObject,
    AnilistImage,
    AnilistMediaList,
    AnilistMediaLists,
    AnilistMediaNextAiringEpisode,
    AnilistMediaTag,
    AnilistMediaTitle,
    AnilistMediaTrailer,
    AnilistPageInfo,
    AnilistStudioNodes,
    AnilistViewerData,
)
from .types import (
    StreamingEpisode as AnilistStreamingEpisode,
)

logger = logging.getLogger(__name__)

user_list_status_map = {
    "CURRENT": UserMediaListStatus.WATCHING,
    "PLANNING": UserMediaListStatus.PLANNING,
    "COMPLETED": UserMediaListStatus.COMPLETED,
    "PAUSED": UserMediaListStatus.PAUSED,
    "REPEATING": UserMediaListStatus.REPEATING,
}
status_map = {
    "FINISHED": MediaStatus.FINISHED,
    "RELEASING": MediaStatus.RELEASING,
    "NOT_YET_RELEASED": MediaStatus.NOT_YET_RELEASED,
    "CANCELLED": MediaStatus.CANCELLED,
    "HIATUS": MediaStatus.HIATUS,
}


def _to_generic_date(date: AnilistDateObject) -> Optional[datetime]:
    if not date:
        return
    year = date["year"]
    month = date["month"]
    day = date["day"]
    if year:
        if not month:
            month = 1
        if not day:
            day = 1
        return datetime(year, month, day)


def _to_generic_media_title(anilist_title: AnilistMediaTitle) -> MediaTitle:
    """Maps an AniList title object to a generic MediaTitle."""
    romaji = anilist_title.get("romaji")
    english = anilist_title.get("english")
    native = anilist_title.get("native")
    return MediaTitle(
        romaji=romaji,
        english=(english or romaji or native or "NO_TITLE"),
        native=native,
    )


def _to_generic_media_image(anilist_image: AnilistImage) -> MediaImage:
    """Maps an AniList image object to a generic MediaImage."""
    return MediaImage(
        medium=anilist_image.get("medium"),
        large=anilist_image["large"],
        extra_large=anilist_image.get("extraLarge"),
    )


def _to_generic_media_trailer(
    anilist_trailer: Optional[AnilistMediaTrailer],
) -> Optional[MediaTrailer]:
    """Maps an AniList trailer object to a generic MediaTrailer."""
    if not anilist_trailer or not anilist_trailer.get("id"):
        return None
    return MediaTrailer(
        id=anilist_trailer["id"],
        site=anilist_trailer.get("site"),
        thumbnail_url=anilist_trailer.get("thumbnail"),
    )


def _to_generic_airing_schedule(
    anilist_schedule: Optional[AnilistMediaNextAiringEpisode],
) -> Optional[AiringSchedule]:
    """Maps an AniList nextAiringEpisode object to a generic AiringSchedule."""
    if anilist_schedule:
        return AiringSchedule(
            airing_at=datetime.fromtimestamp(anilist_schedule["airingAt"])
            if anilist_schedule.get("airingAt")
            else None,
            episode=anilist_schedule.get("episode", 0),
        )


def _to_generic_studios(anilist_studios: AnilistStudioNodes) -> List[Studio]:
    """Maps AniList studio nodes to a list of generic Studio objects."""
    return [
        Studio(
            name=s["name"],
            favourites=s["favourites"],
            is_animation_studio=s["isAnimationStudio"],
        )
        for s in anilist_studios["nodes"]
    ]


def _to_generic_tags(anilist_tags: list[AnilistMediaTag]) -> List[MediaTagItem]:
    """Maps a list of AniList tags to generic MediaTag objects."""
    return [
        MediaTagItem(name=t["name"], rank=t.get("rank"))
        for t in anilist_tags
        if t.get("name")
    ]


# def _to_generic_streaming_episodes(
#     anilist_episodes: list[AnilistStreamingEpisode],
# ) -> List[StreamingEpisode]:
#     """Maps a list of AniList streaming episodes to generic StreamingEpisode objects."""
#     return [
#         StreamingEpisode(title=episode["title"], thumbnail=episode.get("thumbnail"))
#         for episode in anilist_episodes
#         if episode.get("title")
#     ]


# def _to_generic_streaming_episodes(
#     anilist_episodes: list[dict],
# ) -> List[StreamingEpisode]:
#     """Maps a list of AniList streaming episodes to generic StreamingEpisode objects with renumbered episode titles."""

#     # Extract titles
#     titles = [ep["title"] for ep in anilist_episodes if "title" in ep]

#     # Generate mapping: title -> renumbered_ep
#     renumbered_map = renumber_titles(titles)

#     # Apply renumbering
#     return [
#         StreamingEpisode(
#             title=f"{renumbered_map[ep['title']]} - {ep['title']}",
#             thumbnail=ep.get("thumbnail"),
#         )
#         for ep in anilist_episodes
#         if ep.get("title")
#     ]


def _to_generic_streaming_episodes(
    anilist_episodes: list[AnilistStreamingEpisode],
) -> List[StreamingEpisode]:
    """Maps a list of AniList streaming episodes to generic StreamingEpisode objects,
    renumbering them fresh if they contain episode numbers."""

    titles = [ep["title"] for ep in anilist_episodes if "title" in ep and ep["title"]]
    renumber_map = renumber_titles(titles)

    result = []
    for ep in anilist_episodes:
        title = ep.get("title")
        if not title:
            continue

        renumbered_ep = renumber_map.get(title)
        display_title = (
            f"Episode {renumbered_ep} - {strip_original_episode_prefix(title)}"
            if renumbered_ep is not None
            else title
        )

        result.append(
            StreamingEpisode(
                title=display_title,
                thumbnail=ep.get("thumbnail"),
            )
        )

    return result


def _to_generic_user_status(
    anilist_media: AnilistBaseMediaDataSchema,
    anilist_list_entry: Optional[AnilistMediaList],
) -> Optional[UserListItem]:
    """Maps an AniList mediaListEntry to a generic UserListStatus."""
    if anilist_list_entry:
        return UserListItem(
            status=user_list_status_map[anilist_list_entry["status"]],
            progress=anilist_list_entry["progress"],
            score=anilist_list_entry["score"],
            repeat=anilist_list_entry["repeat"],
            notes=anilist_list_entry["notes"],
            start_date=_to_generic_date(anilist_list_entry.get("startDate")),
            completed_at=_to_generic_date(anilist_list_entry.get("completedAt")),
            # TODO: should this be a datetime if so what is the raw values type
            created_at=str(anilist_list_entry["createdAt"]),
        )
    else:
        if not anilist_media["mediaListEntry"]:
            return

        return UserListItem(
            id=anilist_media["mediaListEntry"]["id"],
            status=user_list_status_map[anilist_media["mediaListEntry"]["status"]]
            if anilist_media["mediaListEntry"]["status"]
            else None,
            progress=anilist_media["mediaListEntry"]["progress"],
        )


def _to_generic_media_item(
    data: AnilistBaseMediaDataSchema, media_list: AnilistMediaList | None = None
) -> MediaItem:
    """Maps a single AniList media schema to a generic MediaItem."""
    return MediaItem(
        id=data["id"],
        id_mal=data.get("idMal"),
        type=data.get("type", "ANIME"),
        title=_to_generic_media_title(data["title"]),
        status=status_map[data["status"]],
        format=data.get("format"),
        cover_image=_to_generic_media_image(data["coverImage"]),
        banner_image=data.get("bannerImage"),
        trailer=_to_generic_media_trailer(data["trailer"]),
        description=data.get("description"),
        episodes=data.get("episodes"),
        duration=data.get("duration"),
        genres=data.get("genres", []),
        tags=_to_generic_tags(data.get("tags")),
        studios=_to_generic_studios(data.get("studios")),
        synonymns=data.get("synonyms", []),
        average_score=data.get("averageScore"),
        popularity=data.get("popularity"),
        favourites=data.get("favourites"),
        next_airing=_to_generic_airing_schedule(data.get("nextAiringEpisode")),
        start_date=_to_generic_date(data["startDate"]),
        end_date=_to_generic_date(data["endDate"]),
        streaming_episodes=_to_generic_streaming_episodes(
            data.get("streamingEpisodes", [])
        ),
        user_status=_to_generic_user_status(data, media_list),
    )


def _to_generic_page_info(data: AnilistPageInfo) -> PageInfo:
    """Maps an AniList page info object to a generic PageInfo."""
    return PageInfo(
        total=data.get("total", 0),
        current_page=data.get("currentPage", 1),
        has_next_page=data.get("hasNextPage", False),
        per_page=data.get("perPage", 0),
    )


def to_generic_search_result(
    data: AnilistDataSchema, user_media_list: List[AnilistMediaList] | None = None
) -> Optional[MediaSearchResult]:
    """
    Top-level mapper to convert a raw AniList search/list API response
    into a generic MediaSearchResult object.
    """
    page_data = data["data"]["Page"]

    raw_media_list = page_data["media"]
    if user_media_list:
        media_items: List[MediaItem] = [
            _to_generic_media_item(item, user_media_list_item)
            for item, user_media_list_item in zip(raw_media_list, user_media_list)
        ]
        # TODO: further probe this type
        page_info = _to_generic_page_info(page_data)  # type: ignore
    else:
        media_items: List[MediaItem] = [
            _to_generic_media_item(item) for item in raw_media_list
        ]
        page_info = _to_generic_page_info(page_data["pageInfo"])

    return MediaSearchResult(page_info=page_info, media=media_items)


def to_generic_user_list_result(data: AnilistMediaLists) -> Optional[MediaSearchResult]:
    """
    Mapper for user list queries where media data is nested inside a 'mediaList' key.
    """
    page_data = data["data"]["Page"]

    # Extract media objects from the 'mediaList' array
    media_list_items = page_data["mediaList"]
    raw_media_list = [item["media"] for item in media_list_items]

    # Now that we have a standard list of media, we can reuse the main search result mapper
    return to_generic_search_result(
        {"data": {"Page": {"media": raw_media_list}}},  # pyright:ignore
        media_list_items,
    )


def to_generic_user_profile(data: AnilistViewerData) -> Optional[UserProfile]:
    """Maps a raw AniList viewer response to a generic UserProfile."""

    viewer_data: Optional[AnilistCurrentlyLoggedInUser] = data["data"]["Viewer"]

    return UserProfile(
        id=viewer_data["id"],
        name=viewer_data["name"],
        avatar_url=viewer_data["avatar"]["large"],
        banner_url=viewer_data["bannerImage"],
    )


# TODO: complete this
def to_generic_relations(data: dict) -> Optional[List[MediaItem]]:
    """Maps the 'relations' part of an API response."""
    nodes = data["data"].get("Media", {}).get("relations", {}).get("nodes", [])
    return [_to_generic_media_item(node) for node in nodes if node]


def to_generic_recommendations(data: dict) -> Optional[List[MediaItem]]:
    """Maps the 'recommendations' part of an API response."""
    recs = (
        data.get("data", {})
        .get("Media", {})
        .get("recommendations", {})
        .get("nodes", [])
    )
    return [
        _to_generic_media_item(rec.get("mediaRecommendation"))
        for rec in recs
        if rec.get("mediaRecommendation")
    ]