MoviePilot/app/utils/media.py

from typing import Any, Optional, Tuple

from app.core.config import settings
from app.schemas.types import (
    MUSIC_ENTITY_ALBUM,
    MUSIC_ENTITY_ARTIST,
    MUSIC_ENTITY_RECORDING,
    MUSIC_ENTITY_TYPES,
    MUSIC_SUBSCRIBABLE_TYPES,
)


MEDIA_SOURCE_ALIASES = {
    "tmdb": "themoviedb",
    "themoviedb": "themoviedb",
    "douban": "douban",
    "bangumi": "bangumi",
    "anilist": "anilist",
    "musicbrainz": "musicbrainz",
    "theaudiodb": "theaudiodb",
    "audio_db": "theaudiodb",
    "doubanmusic": "doubanmusic",
    "douban_music": "doubanmusic",
}

MEDIA_SOURCE_PREFIXES = {
    "themoviedb": "tmdb",
    "douban": "douban",
    "bangumi": "bangumi",
    "anilist": "anilist",
    "musicbrainz": "musicbrainz",
    "theaudiodb": "theaudiodb",
    "doubanmusic": "doubanmusic",
}

MEDIA_SOURCE_ID_FIELDS = {
    "themoviedb": ("tmdb_id", "tmdbid"),
    "douban": ("douban_id", "doubanid"),
    "bangumi": ("bangumi_id", "bangumiid"),
    "anilist": ("anilist_id", "anilistid"),
    "musicbrainz": ("media_id",),
    "theaudiodb": ("media_id", "theaudiodb_id"),
    "doubanmusic": ("media_id", "douban_id", "doubanid"),
}

MUSIC_MEDIA_SOURCE_ORDER = ("musicbrainz", "theaudiodb", "doubanmusic")
MUSIC_MEDIA_SOURCES = frozenset(MUSIC_MEDIA_SOURCE_ORDER)


def normalize_music_type(
        value: Optional[object],
        *,
        allow_artist: bool = True,
) -> Optional[str]:
    """规范化音乐实体类型，非法值返回 None。"""
    normalized = str(value or "").strip().lower()
    allowed = MUSIC_ENTITY_TYPES if allow_artist else MUSIC_SUBSCRIBABLE_TYPES
    return normalized if normalized in allowed else None


def is_music_media_source(source: Optional[str]) -> bool:
    """判断单个请求级来源是否为内置音乐元数据源。"""
    return normalize_media_source(source) in MUSIC_MEDIA_SOURCES


def normalize_media_source(source: Optional[str]) -> Optional[str]:
    """规范化媒体数据源名称，兼容外部使用的 ``tmdb`` 前缀。"""
    if not source:
        return None
    normalized = str(source).strip().casefold()
    return MEDIA_SOURCE_ALIASES.get(normalized, normalized or None)


def is_media_source_selected(source: Optional[str], source_key: str) -> bool:
    """
    判断请求级搜索数据源列表（逗号分隔，可为多个）中是否包含指定数据源。

    :param source: 请求级搜索数据源，逗号分隔多个来源，空表示不作限制
    :param source_key: 当前模块对应的数据源标识
    :return: 是否包含
    """
    if not source:
        return True
    normalized_key = normalize_media_source(source_key) or source_key
    return normalized_key in [
        normalize_media_source(item) for item in str(source).split(",")
    ]


def is_media_source_enabled(source: Optional[str], source_key: str) -> bool:
    """
    判断媒体搜索时数据源是否启用：请求级 source（逗号分隔多数据源）优先，
    未指定时回退到全局 SEARCH_SOURCE 配置，两者均未配置时全部启用。

    :param source: 请求级搜索数据源，逗号分隔多个来源
    :param source_key: 当前模块对应的数据源标识
    :return: 是否启用
    """
    if source:
        return is_media_source_selected(source, source_key)
    if settings.SEARCH_SOURCE:
        return is_media_source_selected(settings.SEARCH_SOURCE, source_key)
    return True


def parse_media_key(media_key: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
    """解析带来源前缀的媒体键，返回规范化数据源与原生 ID。"""
    if not media_key or ":" not in str(media_key):
        return None, None
    prefix, media_id = str(media_key).split(":", 1)
    source = normalize_media_source(prefix)
    media_id = media_id.strip()
    if not source or not media_id:
        return None, None
    return source, media_id


def resolve_media_identity(
        media: Any = None,
        source: Optional[str] = None,
        media_id: Optional[Any] = None,
        tmdbid: Optional[Any] = None,
        doubanid: Optional[Any] = None,
        bangumiid: Optional[Any] = None,
        anilistid: Optional[Any] = None,
) -> Tuple[Optional[str], Optional[str]]:
    """
    从统一媒体对象、通用身份或兼容 ID 中解析主媒体身份。

    显式 ``source/media_id`` 优先；未指定来源时按 TMDB、豆瓣、Bangumi、
    AniList 的兼容顺序选择首个有效 ID。
    """
    normalized_source = normalize_media_source(source)
    if normalized_source and media_id is not None and str(media_id).strip():
        return normalized_source, str(media_id).strip()

    values = {
        "themoviedb": tmdbid,
        "douban": doubanid,
        "bangumi": bangumiid,
        "anilist": anilistid,
    }
    if media is not None:
        normalized_source = normalized_source or normalize_media_source(
            getattr(media, "source", None) or getattr(media, "media_source", None)
        )
        object_media_id = getattr(media, "media_id", None)
        if normalized_source and object_media_id is not None and str(object_media_id).strip():
            return normalized_source, str(object_media_id).strip()
        for media_source, fields in MEDIA_SOURCE_ID_FIELDS.items():
            for field in fields:
                value = getattr(media, field, None)
                if value is not None and str(value).strip():
                    values[media_source] = value
                    break

        legacy_source, legacy_media_id = parse_media_key(
            getattr(media, "mediaid", None)
        )
        if not normalized_source and legacy_source and legacy_media_id:
            return legacy_source, legacy_media_id

    if normalized_source:
        value = values.get(normalized_source)
        return (
            normalized_source,
            str(value).strip() if value is not None and str(value).strip() else None,
        )

    for media_source in MEDIA_SOURCE_ID_FIELDS:
        value = values.get(media_source)
        if value is not None and str(value).strip():
            return media_source, str(value).strip()
    return None, None


def build_media_key(source: Optional[str], media_id: Optional[Any]) -> str:
    """构造 API 使用的带来源前缀媒体键。"""
    normalized_source = normalize_media_source(source)
    if not normalized_source or media_id is None or not str(media_id).strip():
        return ""
    prefix = MEDIA_SOURCE_PREFIXES.get(normalized_source, normalized_source)
    return f"{prefix}:{str(media_id).strip()}"