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

"""更新站点工具"""

import json
from typing import Optional, Type

from pydantic import BaseModel, Field

from app.agent.tools.base import MoviePilotTool
from app.agent.tools.tags import ToolTag
from app.core.event import eventmanager
from app.db.site_oper import SiteOper
from app.log import logger
from app.schemas.types import EventType
from app.utils.string import StringUtils


class UpdateSiteInput(BaseModel):
    """更新站点工具的输入参数模型"""

    site_id: int = Field(
        ...,
        description="The ID of the site to update (can be obtained from query_sites tool)",
    )
    name: Optional[str] = Field(None, description="Site name (optional)")
    url: Optional[str] = Field(
        None, description="Site URL (optional, will be automatically formatted)"
    )
    pri: Optional[int] = Field(
        None,
        description="Site priority (optional, smaller value = higher priority, e.g., pri=1 has higher priority than pri=10)",
    )
    rss: Optional[str] = Field(None, description="RSS feed URL (optional)")
    cookie: Optional[str] = Field(None, description="Site cookie (optional)")
    ua: Optional[str] = Field(None, description="User-Agent string (optional)")
    apikey: Optional[str] = Field(None, description="API key (optional)")
    token: Optional[str] = Field(None, description="API token (optional)")
    proxy: Optional[int] = Field(
        None, description="Whether to use proxy: 0 for no, 1 for yes (optional)"
    )
    filter: Optional[str] = Field(
        None, description="Filter rule as regular expression (optional)"
    )
    note: Optional[str] = Field(None, description="Site notes/remarks (optional)")
    timeout: Optional[int] = Field(
        None, description="Request timeout in seconds (optional, default: 15)"
    )
    limit_interval: Optional[int] = Field(
        None, description="Rate limit interval in seconds (optional)"
    )
    limit_count: Optional[int] = Field(
        None, description="Rate limit count per interval (optional)"
    )
    limit_seconds: Optional[int] = Field(
        None, description="Rate limit seconds between requests (optional)"
    )
    is_active: Optional[bool] = Field(
        None,
        description="Whether site is active: True for enabled, False for disabled (optional)",
    )
    downloader: Optional[str] = Field(
        None, description="Downloader name for this site (optional)"
    )


class UpdateSiteTool(MoviePilotTool):
    name: str = "update_site"
    tags: list[str] = [
        ToolTag.Write,
        ToolTag.Site,
        ToolTag.Admin,
    ]
    description: str = "Update site configuration including URL, priority, authentication credentials (cookie, UA, API key), proxy settings, rate limits, and other site properties. Supports updating multiple site attributes at once. Site priority (pri): smaller values have higher priority (e.g., pri=1 has higher priority than pri=10)."
    args_schema: Type[BaseModel] = UpdateSiteInput
    require_admin: bool = True

    def get_tool_message(self, **kwargs) -> Optional[str]:
        """根据更新参数生成友好的提示消息"""
        site_id = kwargs.get("site_id")
        fields_updated = []

        if kwargs.get("name"):
            fields_updated.append("名称")
        if kwargs.get("url"):
            fields_updated.append("URL")
        if kwargs.get("pri") is not None:
            fields_updated.append("优先级")
        if kwargs.get("cookie"):
            fields_updated.append("Cookie")
        if kwargs.get("ua"):
            fields_updated.append("User-Agent")
        if kwargs.get("proxy") is not None:
            fields_updated.append("代理设置")
        if kwargs.get("is_active") is not None:
            fields_updated.append("启用状态")
        if kwargs.get("downloader"):
            fields_updated.append("下载器")

        if fields_updated:
            return f"更新站点 #{site_id}: {', '.join(fields_updated)}"
        return f"更新站点 #{site_id}"

    async def run(
        self,
        site_id: int,
        name: Optional[str] = None,
        url: Optional[str] = None,
        pri: Optional[int] = None,
        rss: Optional[str] = None,
        cookie: Optional[str] = None,
        ua: Optional[str] = None,
        apikey: Optional[str] = None,
        token: Optional[str] = None,
        proxy: Optional[int] = None,
        filter: Optional[str] = None,
        note: Optional[str] = None,
        timeout: Optional[int] = None,
        limit_interval: Optional[int] = None,
        limit_count: Optional[int] = None,
        limit_seconds: Optional[int] = None,
        is_active: Optional[bool] = None,
        downloader: Optional[str] = None,
        **kwargs,
    ) -> str:
        logger.info(f"执行工具: {self.name}, 参数: site_id={site_id}")

        try:
            site_oper = SiteOper()
            site = await site_oper.async_get(site_id)
            if not site:
                return json.dumps(
                    {"success": False, "message": f"站点不存在: {site_id}"},
                    ensure_ascii=False,
                )

            # 构建更新字典
            site_dict = {}

            # 基本信息
            if name is not None:
                site_dict["name"] = name

            # URL处理（需要校正格式）
            if url is not None:
                _scheme, _netloc = StringUtils.get_url_netloc(url)
                site_dict["url"] = f"{_scheme}://{_netloc}/"

            if pri is not None:
                site_dict["pri"] = pri
            if rss is not None:
                site_dict["rss"] = rss

            # 认证信息
            if cookie is not None:
                site_dict["cookie"] = cookie
            if ua is not None:
                site_dict["ua"] = ua
            if apikey is not None:
                site_dict["apikey"] = apikey
            if token is not None:
                site_dict["token"] = token

            # 配置选项
            if proxy is not None:
                site_dict["proxy"] = proxy
            if filter is not None:
                site_dict["filter"] = filter
            if note is not None:
                site_dict["note"] = note
            if timeout is not None:
                site_dict["timeout"] = timeout

            # 流控设置
            if limit_interval is not None:
                site_dict["limit_interval"] = limit_interval
            if limit_count is not None:
                site_dict["limit_count"] = limit_count
            if limit_seconds is not None:
                site_dict["limit_seconds"] = limit_seconds

            # 状态和下载器
            if is_active is not None:
                site_dict["is_active"] = is_active
            if downloader is not None:
                site_dict["downloader"] = downloader

            # 如果没有要更新的字段
            if not site_dict:
                return json.dumps(
                    {"success": False, "message": "没有提供要更新的字段"},
                    ensure_ascii=False,
                )

            # 更新站点
            await site_oper.async_update(site_id, site_dict)

            # 重新获取更新后的站点数据
            updated_site = await site_oper.async_get(site_id)

            # 发送站点更新事件
            await eventmanager.async_send_event(
                EventType.SiteUpdated,
                {"domain": updated_site.domain if updated_site else site.domain},
            )

            # 构建返回结果
            result = {
                "success": True,
                "message": f"站点 #{site_id} 更新成功",
                "site_id": site_id,
                "updated_fields": list(site_dict.keys()),
            }

            if updated_site:
                result["site"] = {
                    "id": updated_site.id,
                    "name": updated_site.name,
                    "domain": updated_site.domain,
                    "url": updated_site.url,
                    "pri": updated_site.pri,
                    "is_active": updated_site.is_active,
                    "downloader": updated_site.downloader,
                    "proxy": updated_site.proxy,
                    "timeout": updated_site.timeout,
                }

            return json.dumps(result, ensure_ascii=False, indent=2)

        except Exception as e:
            error_message = f"更新站点失败: {str(e)}"
            logger.error(f"更新站点失败: {e}", exc_info=True)
            return json.dumps(
                {"success": False, "message": error_message, "site_id": site_id},
                ensure_ascii=False,
            )