MoviePilot/app/agent/tools/impl/_torrent_search_utils.py

"""种子搜索工具辅助函数"""

import re
from typing import List, Optional

from app.core.context import Context
from app.schemas.types import MediaType, media_type_to_agent
from app.utils.crypto import HashUtils
from app.utils.string import StringUtils
from ._music_utils import simplify_music_info

SEARCH_RESULT_CACHE_FILE = "__search_result__"
TORRENT_RESULT_LIMIT = 50


def build_torrent_ref(context: Optional[Context]) -> str:
    """生成用于下载校验的短引用"""
    if not context or not context.torrent_info:
        return ""
    return HashUtils.sha1(context.torrent_info.enclosure or "")[:7]


def sort_season_options(options: List[str]) -> List[str]:
    """按前端逻辑排序季集选项"""
    if len(options) <= 1:
        return options

    parsed_options = []
    for index, option in enumerate(options):
        match = re.match(r"^S(\d+)(?:-S(\d+))?\s*(?:E(\d+)(?:-E(\d+))?)?$", option or "")
        if not match:
            parsed_options.append({
                "original": option,
                "season_num": 0,
                "episode_num": 0,
                "max_episode_num": 0,
                "is_whole_season": False,
                "index": index,
            })
            continue

        episode_num = int(match.group(3)) if match.group(3) else 0
        max_episode_num = int(match.group(4)) if match.group(4) else episode_num
        parsed_options.append({
            "original": option,
            "season_num": int(match.group(1)),
            "episode_num": episode_num,
            "max_episode_num": max_episode_num,
            "is_whole_season": not match.group(3),
            "index": index,
        })

    whole_seasons = [item for item in parsed_options if item["is_whole_season"]]
    episodes = [item for item in parsed_options if not item["is_whole_season"]]

    whole_seasons.sort(key=lambda item: (-item["season_num"], item["index"]))
    episodes.sort(
        key=lambda item: (
            -item["season_num"],
            -(item["max_episode_num"] or item["episode_num"]),
            -item["episode_num"],
            item["index"],
        )
    )
    return [item["original"] for item in whole_seasons + episodes]


def append_option(options: List[str], value: Optional[str]) -> None:
    """按前端逻辑收集去重后的筛选项"""
    if value and value not in options:
        options.append(value)


def build_filter_options(items: List[Context]) -> dict:
    """从搜索结果中构建筛选项汇总"""
    filter_options = {
        "site": [],
        "season": [],
        "freeState": [],
        "edition": [],
        "resolution": [],
        "videoCode": [],
        "releaseGroup": [],
    }

    for item in items:
        torrent_info = item.torrent_info
        meta_info = item.meta_info
        append_option(filter_options["site"], getattr(torrent_info, "site_name", None))
        append_option(filter_options["season"], getattr(meta_info, "season_episode", None))
        append_option(filter_options["freeState"], getattr(torrent_info, "volume_factor", None))
        append_option(filter_options["edition"], getattr(meta_info, "edition", None))
        append_option(filter_options["resolution"], getattr(meta_info, "resource_pix", None))
        append_option(filter_options["videoCode"], getattr(meta_info, "video_encode", None))
        append_option(filter_options["releaseGroup"], getattr(meta_info, "resource_team", None))

    filter_options["season"] = sort_season_options(filter_options["season"])
    return filter_options


def match_filter(filter_values: Optional[List[str]], value: Optional[str]) -> bool:
    """匹配前端同款多选筛选规则"""
    return not filter_values or bool(value and value in filter_values)


def filter_contexts(items: List[Context],
                    site: Optional[List[str]] = None,
                    season: Optional[List[str]] = None,
                    free_state: Optional[List[str]] = None,
                    video_code: Optional[List[str]] = None,
                    edition: Optional[List[str]] = None,
                    resolution: Optional[List[str]] = None,
                    release_group: Optional[List[str]] = None) -> List[Context]:
    """按前端同款维度筛选结果"""
    filtered_items = []
    for item in items:
        torrent_info = item.torrent_info
        meta_info = item.meta_info
        if (
            match_filter(site, getattr(torrent_info, "site_name", None))
            and match_filter(free_state, getattr(torrent_info, "volume_factor", None))
            and match_filter(season, getattr(meta_info, "season_episode", None))
            and match_filter(release_group, getattr(meta_info, "resource_team", None))
            and match_filter(video_code, getattr(meta_info, "video_encode", None))
            and match_filter(resolution, getattr(meta_info, "resource_pix", None))
            and match_filter(edition, getattr(meta_info, "edition", None))
        ):
            filtered_items.append(item)
    return filtered_items


def simplify_search_result(
    context: Context,
    index: int,
    include_description: bool = False,
) -> dict:
    """
    精简单条搜索结果

    :param context: 搜索结果上下文
    :param index: 搜索结果在原始缓存中的序号
    :param include_description: 是否返回种子简介
    :return: 精简后的搜索结果
    """
    simplified = {}
    torrent_info = context.torrent_info
    meta_info = context.meta_info
    media_info = context.media_info

    if torrent_info:
        simplified["torrent_info"] = {
            "title": torrent_info.title,
            "size": StringUtils.format_size(torrent_info.size),
            "seeders": torrent_info.seeders,
            "peers": torrent_info.peers,
            "site_name": torrent_info.site_name,
            "torrent_url": f"{build_torrent_ref(context)}:{index}",
            "page_url": torrent_info.page_url,
            "volume_factor": torrent_info.volume_factor,
            "freedate_diff": torrent_info.freedate_diff,
            "pubdate": torrent_info.pubdate,
        }
        if include_description:
            simplified["torrent_info"]["description"] = torrent_info.description

    if media_info:
        if getattr(media_info, "type", None) == MediaType.MUSIC:
            simplified["media_info"] = simplify_music_info(media_info)
        else:
            simplified["media_info"] = {
                "title": getattr(media_info, "title", None),
                "en_title": getattr(media_info, "en_title", None),
                "year": getattr(media_info, "year", None),
                "type": media_type_to_agent(getattr(media_info, "type", None)),
                "season": getattr(media_info, "season", None),
                "tmdb_id": getattr(media_info, "tmdb_id", None),
            }

    if meta_info:
        if getattr(meta_info, "type", None) == MediaType.MUSIC:
            simplified["meta_info"] = {
                key: getattr(meta_info, key, None)
                for key in (
                    "title", "artists", "album", "album_artist", "year",
                    "disc_number", "track_number", "total_tracks", "version",
                    "audio_format", "bit_depth", "sample_rate", "bitrate",
                    "duration", "isrc", "media_source", "media_id",
                )
                if getattr(meta_info, key, None) not in (None, "", [])
            }
            simplified["meta_info"]["type"] = "music"
        else:
            simplified["meta_info"] = {
                "name": getattr(meta_info, "name", None),
                "cn_name": getattr(meta_info, "cn_name", None),
                "en_name": getattr(meta_info, "en_name", None),
                "year": getattr(meta_info, "year", None),
                "type": media_type_to_agent(getattr(meta_info, "type", None)),
                "begin_season": getattr(meta_info, "begin_season", None),
                "season_episode": getattr(meta_info, "season_episode", None),
                "resource_team": getattr(meta_info, "resource_team", None),
                "video_encode": getattr(meta_info, "video_encode", None),
                "edition": getattr(meta_info, "edition", None),
                "resource_pix": getattr(meta_info, "resource_pix", None),
            }

    return simplified