feat: add exact subtitle search

app/chain/search.py

@@ -204,6 +204,7 @@ class SearchChain(ChainBase):
             "title": str(params.get("title") or ""),
             "year": str(params.get("year") or ""),
             "season": str(params.get("season") or ""),
+            "episode": str(params.get("episode") or ""),
             "sites": str(params.get("sites") or ""),
             "result_type": str(params.get("result_type") or "torrent"),
         }
@@ -218,6 +219,7 @@ class SearchChain(ChainBase):
             title: Optional[str] = None,
             year: Optional[str] = None,
             season: Optional[int] = None,
+            episode: Optional[int] = None,
             sites: Optional[List[int]] = None,
             result_type: Optional[str] = "torrent",
     ) -> None:
@@ -232,6 +234,7 @@ class SearchChain(ChainBase):
                 "title": title,
                 "year": year,
                 "season": season,
+                "episode": episode,
                 "sites": self._stringify_sites(sites),
                 "result_type": result_type or "torrent",
             }
@@ -248,6 +251,7 @@ class SearchChain(ChainBase):
             title: Optional[str] = None,
             year: Optional[str] = None,
             season: Optional[int] = None,
+            episode: Optional[int] = None,
             sites: Optional[List[int]] = None,
             result_type: Optional[str] = "torrent",
     ) -> None:
@@ -262,6 +266,7 @@ class SearchChain(ChainBase):
                 "title": title,
                 "year": year,
                 "season": season,
+                "episode": episode,
                 "sites": self._stringify_sites(sites),
                 "result_type": result_type or "torrent",
             }
@@ -650,6 +655,101 @@ class SearchChain(ChainBase):
             "total_items": len(subtitles)
         }
 
+    async def async_search_subtitles_by_id(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
+                                           mtype: MediaType = None, season: Optional[int] = None,
+                                           episode: Optional[int] = None, sites: List[int] = None,
+                                           cache_local: bool = False) -> List[SubtitleInfo]:
+        """
+        根据TMDBID/豆瓣ID异步精确搜索字幕，不应用过滤规则。
+        :param tmdbid: TMDB ID
+        :param doubanid: 豆瓣 ID
+        :param mtype: 媒体，电影 or 电视剧
+        :param season: 季数
+        :param episode: 集数
+        :param sites: 站点ID列表
+        :param cache_local: 是否缓存到本地
+        """
+        if cache_local:
+            self.cancel_ai_recommend()
+            await self.async_save_last_search_params(
+                keyword=self._build_search_keyword(tmdbid=tmdbid, doubanid=doubanid),
+                mtype=mtype,
+                area="title",
+                season=season,
+                episode=episode,
+                sites=sites,
+                result_type="subtitle",
+            )
+        mediainfo = await self.async_recognize_media(tmdbid=tmdbid, doubanid=doubanid, mtype=mtype)
+        if not mediainfo:
+            logger.error(f'{tmdbid} 媒体信息识别失败！')
+            return []
+        subtitles = await self.__async_search_subtitles_for_media(
+            mediainfo=mediainfo,
+            tmdbid=tmdbid,
+            doubanid=doubanid,
+            season=season,
+            episode=episode,
+            sites=sites,
+        )
+        if cache_local:
+            await self.async_save_cache(subtitles, self.__subtitle_result_temp_file)
+        return subtitles
+
+    async def async_search_subtitles_by_id_stream(
+            self,
+            tmdbid: Optional[int] = None,
+            doubanid: Optional[str] = None,
+            mtype: MediaType = None,
+            season: Optional[int] = None,
+            episode: Optional[int] = None,
+            sites: List[int] = None,
+            cache_local: bool = False,
+    ) -> AsyncIterator[dict]:
+        """
+        根据TMDBID/豆瓣ID渐进式精确搜索字幕，先返回站点候选，再返回标题和剧集匹配后的结果。
+        """
+        if cache_local:
+            self.cancel_ai_recommend()
+            await self.async_save_last_search_params(
+                keyword=self._build_search_keyword(tmdbid=tmdbid, doubanid=doubanid),
+                mtype=mtype,
+                area="title",
+                season=season,
+                episode=episode,
+                sites=sites,
+                result_type="subtitle",
+            )
+        mediainfo = await self.async_recognize_media(tmdbid=tmdbid, doubanid=doubanid, mtype=mtype)
+        if not mediainfo:
+            logger.error(f'{tmdbid} 媒体信息识别失败！')
+            yield {
+                "type": "error",
+                "success": False,
+                "message": "媒体信息识别失败"
+            }
+            return
+
+        subtitles: List[SubtitleInfo] = []
+        async for event in self.__async_search_subtitles_for_media_stream(
+                mediainfo=mediainfo,
+                tmdbid=tmdbid,
+                doubanid=doubanid,
+                season=season,
+                episode=episode,
+                sites=sites):
+            if event.get("type") == "done":
+                subtitles = event.get("subtitles") or []
+                event = {
+                    key: value
+                    for key, value in event.items()
+                    if key != "subtitles"
+                }
+            yield event
+
+        if cache_local:
+            await self.async_save_cache(subtitles, self.__subtitle_result_temp_file)
+
     async def async_search_by_id(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
                                  mtype: MediaType = None, area: Optional[str] = "title", season: Optional[int] = None,
                                  sites: List[int] = None, cache_local: bool = False) -> List[Context]:
@@ -1341,6 +1441,321 @@ class SearchChain(ChainBase):
             "contexts": contexts
         }
 
+    @staticmethod
+    def __build_subtitle_season_episodes(mediainfo: MediaInfo,
+                                          season: Optional[int] = None,
+                                          episode: Optional[int] = None) -> Optional[Dict[int, List[int]]]:
+        """
+        构造字幕匹配用季集约束，未指定集数时只约束到同一季。
+        """
+        if mediainfo.type != MediaType.TV:
+            return None
+        media_season = season if season is not None else mediainfo.season
+        if media_season is None:
+            return None
+        return {media_season: [episode] if episode is not None else []}
+
+    @staticmethod
+    def __build_subtitle_torrent(subtitle: SubtitleInfo, title: Optional[str] = None) -> TorrentInfo:
+        """
+        将字幕结果转换为轻量资源对象，复用既有标题匹配逻辑。
+        """
+        return TorrentInfo(
+            site=subtitle.site,
+            site_name=subtitle.site_name,
+            site_cookie=subtitle.site_cookie,
+            site_ua=subtitle.site_ua,
+            site_proxy=subtitle.site_proxy,
+            site_order=subtitle.site_order,
+            title=title or subtitle.title or subtitle.file_name,
+            description=subtitle.description,
+            enclosure=subtitle.enclosure,
+            page_url=subtitle.page_url,
+            size=subtitle.size,
+            grabs=subtitle.grabs,
+            pubdate=subtitle.pubdate,
+            date_elapsed=subtitle.date_elapsed,
+        )
+
+    @staticmethod
+    def __build_subtitle_names(subtitle: SubtitleInfo) -> List[str]:
+        """
+        提取字幕标题和下载文件名，作为精确匹配的名称候选。
+        """
+        return list(dict.fromkeys(
+            name.strip()
+            for name in (subtitle.title, subtitle.file_name)
+            if name and name.strip()
+        ))
+
+    @staticmethod
+    def __build_subtitle_meta(title: str,
+                              subtitle: SubtitleInfo,
+                              custom_words: Optional[List[str]] = None) -> MetaInfo:
+        """
+        识别字幕名称。
+        """
+        return MetaInfo(
+            title=title,
+            subtitle=subtitle.description,
+            custom_words=custom_words,
+        )
+
+    @staticmethod
+    def __match_subtitle_episode(meta: MetaInfo,
+                                 season_episodes: Optional[Dict[int, List[int]]],
+                                 episode: Optional[int] = None) -> bool:
+        """
+        判断字幕识别出的季集是否落在目标媒体季集内。
+        """
+        if not season_episodes:
+            return True
+        subtitle_torrent = TorrentInfo(title=meta.org_string)
+        if not TorrentHelper.match_season_episodes(
+                torrent=subtitle_torrent,
+                meta=meta,
+                season_episodes=season_episodes):
+            return False
+        if episode is not None:
+            return bool(meta.episode_list) and episode in meta.episode_list
+        return True
+
+    def __parse_subtitle_result(self,
+                                subtitles: List[SubtitleInfo],
+                                mediainfo: MediaInfo,
+                                keyword: Optional[str] = None,
+                                season_episodes: Optional[Dict[int, List[int]]] = None,
+                                episode: Optional[int] = None,
+                                custom_words: Optional[List[str]] = None) -> List[SubtitleInfo]:
+        """
+        识别并精确匹配字幕搜索结果，不使用任何过滤规则。
+        """
+        if not subtitles:
+            logger.warn(f'{keyword or mediainfo.title} 未搜索到字幕')
+            return []
+
+        match_subtitles = []
+        logger.info(f"开始匹配字幕 标题：{mediainfo.title}，原标题：{mediainfo.original_title}，别名：{mediainfo.names}")
+        for subtitle in subtitles:
+            if global_vars.is_system_stopped:
+                break
+            subtitle_names = self.__build_subtitle_names(subtitle)
+            if not subtitle_names:
+                continue
+
+            for subtitle_name in subtitle_names:
+                subtitle_meta = self.__build_subtitle_meta(
+                    title=subtitle_name,
+                    subtitle=subtitle,
+                    custom_words=custom_words,
+                )
+                if not self.__match_subtitle_episode(
+                        meta=subtitle_meta,
+                        season_episodes=season_episodes,
+                        episode=episode):
+                    continue
+
+                subtitle_torrent = self.__build_subtitle_torrent(
+                    subtitle=subtitle,
+                    title=subtitle_name,
+                )
+                if TorrentHelper.match_torrent(
+                        mediainfo=mediainfo,
+                        torrent_meta=subtitle_meta,
+                        torrent=subtitle_torrent):
+                    match_subtitles.append(subtitle)
+                    break
+
+        logger.info(f"字幕匹配完成，共匹配到 {len(match_subtitles)} 个字幕")
+        return self.__remove_duplicate_subtitles(match_subtitles)
+
+    @staticmethod
+    def __remove_duplicate_subtitles(subtitles: List[SubtitleInfo]) -> List[SubtitleInfo]:
+        """
+        去除重复的字幕结果。
+        """
+        return list({
+            f"{subtitle.site_name}_{subtitle.torrent_id}_{subtitle.subtitle_id}_{subtitle.title}_{subtitle.enclosure}": subtitle
+            for subtitle in subtitles
+        }.values())
+
+    async def __async_search_subtitles_for_media(self,
+                                                 mediainfo: MediaInfo,
+                                                 tmdbid: Optional[int] = None,
+                                                 doubanid: Optional[str] = None,
+                                                 season: Optional[int] = None,
+                                                 episode: Optional[int] = None,
+                                                 sites: List[int] = None,
+                                                 custom_words: List[str] = None) -> List[SubtitleInfo]:
+        """
+        根据媒体信息搜索并精确匹配字幕结果。
+        """
+        if not mediainfo.tmdb_id:
+            meta = MetaInfo(title=mediainfo.title)
+            mediainfo.title = meta.name
+            mediainfo.season = meta.begin_season
+        logger.info(f'开始精确搜索字幕，关键词：{mediainfo.title} ...')
+
+        if not mediainfo.names:
+            mediainfo = await self.async_recognize_media(mtype=mediainfo.type,
+                                                         tmdbid=mediainfo.tmdb_id,
+                                                         doubanid=mediainfo.douban_id)
+            if not mediainfo:
+                logger.error('媒体信息识别失败！')
+                return []
+
+        no_exists = None
+        if season is not None:
+            no_exists = {
+                tmdbid or doubanid: {
+                    season: NotExistMediaInfo(episodes=[episode] if episode is not None else [])
+                }
+            }
+        season_episodes, keywords = self.__prepare_params(
+            mediainfo=mediainfo,
+            no_exists=no_exists,
+        )
+        season_episodes = self.__build_subtitle_season_episodes(
+            mediainfo=mediainfo,
+            season=season,
+            episode=episode,
+        ) or season_episodes
+
+        subtitles: List[SubtitleInfo] = []
+        search_count = 0
+        for search_word in keywords:
+            if search_count > 0:
+                logger.info(f"已搜索 {search_count} 次，强制休眠 1-10 秒 ...")
+                await asyncio.sleep(random.randint(1, 10))
+            subtitles.extend(
+                await self.__async_search_subtitles_all_sites(
+                    keyword=search_word,
+                    sites=sites,
+                ) or []
+            )
+            search_count += 1
+            if not settings.SEARCH_MULTIPLE_NAME and subtitles:
+                logger.info(f"共搜索到 {len(subtitles)} 个字幕，停止搜索")
+                break
+
+        return await run_in_threadpool(
+            self.__parse_subtitle_result,
+            subtitles=subtitles,
+            mediainfo=mediainfo,
+            keyword=mediainfo.title,
+            season_episodes=season_episodes,
+            episode=episode,
+            custom_words=custom_words,
+        )
+
+    async def __async_search_subtitles_for_media_stream(
+            self,
+            mediainfo: MediaInfo,
+            tmdbid: Optional[int] = None,
+            doubanid: Optional[str] = None,
+            season: Optional[int] = None,
+            episode: Optional[int] = None,
+            sites: List[int] = None,
+            custom_words: List[str] = None,
+    ) -> AsyncIterator[Dict[str, Any]]:
+        """
+        根据媒体信息渐进式搜索并精确匹配字幕结果。
+        """
+        if not mediainfo.tmdb_id:
+            meta = MetaInfo(title=mediainfo.title)
+            mediainfo.title = meta.name
+            mediainfo.season = meta.begin_season
+        logger.info(f'开始渐进式精确搜索字幕，关键词：{mediainfo.title} ...')
+
+        if not mediainfo.names:
+            mediainfo = await self.async_recognize_media(mtype=mediainfo.type,
+                                                         tmdbid=mediainfo.tmdb_id,
+                                                         doubanid=mediainfo.douban_id)
+            if not mediainfo:
+                logger.error('媒体信息识别失败！')
+                yield {
+                    "type": "error",
+                    "success": False,
+                    "message": "媒体信息识别失败"
+                }
+                return
+
+        no_exists = None
+        if season is not None:
+            no_exists = {
+                tmdbid or doubanid: {
+                    season: NotExistMediaInfo(episodes=[episode] if episode is not None else [])
+                }
+            }
+        season_episodes, keywords = self.__prepare_params(
+            mediainfo=mediainfo,
+            no_exists=no_exists,
+        )
+        season_episodes = self.__build_subtitle_season_episodes(
+            mediainfo=mediainfo,
+            season=season,
+            episode=episode,
+        ) or season_episodes
+
+        subtitles: List[SubtitleInfo] = []
+        search_count = 0
+        for search_word in keywords:
+            if search_count > 0:
+                logger.info(f"已搜索 {search_count} 次，强制休眠 1-10 秒 ...")
+                await asyncio.sleep(random.randint(1, 10))
+
+            async for event in self.__async_search_subtitles_all_sites_stream(
+                    keyword=search_word,
+                    sites=sites):
+                result = event.pop("items", []) or []
+                subtitles.extend(result)
+                yield {
+                    **event,
+                    "type": "append",
+                    "stage": "searching",
+                    "items": [subtitle.to_dict() for subtitle in result],
+                    "total_items": len(subtitles)
+                }
+
+            search_count += 1
+            if not settings.SEARCH_MULTIPLE_NAME and subtitles:
+                logger.info(f"共搜索到 {len(subtitles)} 个字幕，停止搜索")
+                break
+
+        yield {
+            "type": "progress",
+            "stage": "filtering",
+            "value": 98,
+            "text": f"正在识别匹配 {len(subtitles)} 个候选字幕 ..."
+        }
+
+        match_subtitles = await run_in_threadpool(
+            self.__parse_subtitle_result,
+            subtitles=subtitles,
+            mediainfo=mediainfo,
+            keyword=mediainfo.title,
+            season_episodes=season_episodes,
+            episode=episode,
+            custom_words=custom_words,
+        )
+        final_items = [subtitle.to_dict() for subtitle in match_subtitles]
+        yield {
+            "type": "replace",
+            "stage": "filtered",
+            "value": 100,
+            "text": f"识别匹配完成，共 {len(match_subtitles)} 个字幕",
+            "items": final_items,
+            "total_items": len(match_subtitles)
+        }
+        yield {
+            "type": "done",
+            "stage": "done",
+            "text": f"搜索完成，共 {len(match_subtitles)} 个字幕",
+            "items": final_items,
+            "total_items": len(match_subtitles),
+            "subtitles": match_subtitles
+        }
+
     def __search_all_sites(self, keyword: str,
                            mediainfo: Optional[MediaInfo] = None,
                            sites: List[int] = None,