From a1e178c80584771b6019bbc98d5c227f95ca4a27 Mon Sep 17 00:00:00 2001
From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com>
Date: Wed, 4 Dec 2024 20:21:57 +0800
Subject: [PATCH] feat(event): add ResourceSelection event for update resource
 contexts

---
 app/chain/download.py | 18 +++++++++++++++++-
 app/schemas/event.py  | 28 +++++++++++++++++++++++++++-
 app/schemas/types.py  |  2 ++
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/app/chain/download.py b/app/chain/download.py
index 0569d9e2..faf726f1 100644
--- a/app/chain/download.py
+++ b/app/chain/download.py
@@ -20,7 +20,8 @@ from app.helper.message import MessageHelper
 from app.helper.torrent import TorrentHelper
 from app.log import logger
 from app.schemas import ExistMediaInfo, NotExistMediaInfo, DownloadingTorrent, Notification
-from app.schemas.types import MediaType, TorrentStatus, EventType, MessageChannel, NotificationType
+from app.schemas.event import ResourceSelectionEventData
+from app.schemas.types import MediaType, TorrentStatus, EventType, MessageChannel, NotificationType, ChainEventType
 from app.utils.http import RequestUtils
 from app.utils.string import StringUtils
 
@@ -458,6 +459,21 @@ class DownloadChain(ChainBase):
                 return 9999
             return no_exist[season].total_episode
 
+        # 发送资源选择事件，允许外部修改上下文数据
+        logger.debug(f"Initial contexts: {len(contexts)} items, Downloader: {downloader}")
+        event_data = ResourceSelectionEventData(
+            contexts=contexts,
+            downloader=downloader
+        )
+        event = eventmanager.send_event(ChainEventType.ResourceSelection, event_data)
+        # 如果事件修改了上下文数据，使用更新后的数据
+        if event and event.event_data:
+            event_data: ResourceSelectionEventData = event.event_data
+            if event_data.updated and event_data.updated_contexts is not None:
+                logger.debug(f"Contexts updated by event: "
+                             f"{len(event_data.updated_contexts)} items (source: {event_data.source})")
+                contexts = event_data.updated_contexts
+
         # 分组排序
         contexts = TorrentHelper().sort_group_torrents(contexts)
 
diff --git a/app/schemas/event.py b/app/schemas/event.py
index 58d77005..facb6017 100644
--- a/app/schemas/event.py
+++ b/app/schemas/event.py
@@ -1,8 +1,10 @@
 from pathlib import Path
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, List
 
 from pydantic import BaseModel, Field, root_validator
 
+from app.core.context import Context
+
 
 class BaseEventData(BaseModel):
     """
@@ -143,3 +145,27 @@ class TransferRenameEventData(ChainEventData):
     updated: bool = Field(False, description="是否已更新")
     updated_str: Optional[str] = Field(None, description="更新后的字符串")
     source: Optional[str] = Field("未知拦截源", description="拦截源")
+
+
+class ResourceSelectionEventData(BaseModel):
+    """
+    ResourceSelection 事件的数据模型
+
+    Attributes:
+        # 输入参数
+        contexts (List[Context]): 当前待选择的资源上下文列表
+        source (str): 事件源，指示事件的触发来源
+
+        # 输出参数
+        updated (bool): 是否已更新，默认值为 False
+        updated_contexts (Optional[List[Context]]): 已更新的资源上下文列表，默认值为 None
+        source (str): 更新源，默认值为 "未知更新源"
+    """
+    # 输入参数
+    contexts: Any = Field(None, description="待选择的资源上下文列表")
+    downloader: Optional[str] = Field(None, description="下载器")
+
+    # 输出参数
+    updated: bool = Field(False, description="是否已更新")
+    updated_contexts: Optional[List[Context]] = Field(None, description="已更新的资源上下文列表")
+    source: Optional[str] = Field("未知拦截源", description="拦截源")
diff --git a/app/schemas/types.py b/app/schemas/types.py
index 3cb1d72c..5402f3bd 100644
--- a/app/schemas/types.py
+++ b/app/schemas/types.py
@@ -70,6 +70,8 @@ class ChainEventType(Enum):
     CommandRegister = "command.register"
     # 整理重命名
     TransferRename = "transfer.rename"
+    # 资源选择
+    ResourceSelection = "resource.selection"
 
 
 # 系统配置Key字典