HeXi/hexi/plugins/nonebot_plugin_video_analysis/cleanup.py

"""temp 目录清理 — 手动 / 自动可选

hexi/data/temp 是下载媒体中转区(见 utils.get_temp_root),发送成功后保留
供用户取用。清理方式由配置 `video_analysis_temp_cleanup_mode` 决定:

- auto   : 每日定时清理,删除超过配置天数未修改的文件(默认)
- manual : 关闭定时任务,仅通过命令手动清理

命令(manual 模式下可用):
- /清理temp [天数]     清理 temp 下超过 N 天(默认取配置)未修改的文件
- /temp统计            查看 temp 目录占用情况

媒体按「作者目录/作品目录」分层落盘,文件删完后这些空目录会一并收掉
(只删同为过期、且确实为空的目录 —— 正在落盘的目录 mtime 很新,不会误删)。
占用中的文件自动跳过。
"""

import asyncio
import time
from collections.abc import Callable
from pathlib import Path
from typing import Literal

from nonebot import logger, require, get_plugin_config, on_command
from nonebot.adapters.onebot.v11 import Bot, MessageEvent
from pydantic import BaseModel

from hexi.core.custom_utils import check_manage
from hexi.core.task_slots import TASKS, Task

require("nonebot_plugin_apscheduler")
from nonebot_plugin_apscheduler import scheduler
from hexi.core.scheduler import add_job

from nonebot_plugin_alconna import UniMessage

from .utils import get_temp_root


class CleanupConfig(BaseModel):
    # 清理模式:auto=每日定时自动清理, manual=仅命令手动清理
    video_analysis_temp_cleanup_mode: Literal["auto", "manual"] = "auto"
    # 超过该天数未修改的 temp 文件会被清理(默认 7 天)
    video_analysis_temp_max_age_days: int = 7
    # 自动模式下每日清理时间(小时,0-23,默认凌晨 4 点)
    video_analysis_temp_cleanup_hour: int = 4


cleanup_config = get_plugin_config(CleanupConfig)


def _max_age_seconds(days: int | None = None) -> float:
    """由天数配置换算为秒;非法值回落默认 7 天"""
    if days is None:
        days = cleanup_config.video_analysis_temp_max_age_days
    try:
        days = max(0, int(days))
    except (TypeError, ValueError):
        days = 7
    return days * 24 * 60 * 60


def _file_is_stale(path: Path, max_age: float) -> bool:
    """文件最后修改时间距今是否超过 max_age 秒"""
    try:
        return time.time() - path.stat().st_mtime > max_age
    except OSError:
        # 文件已被删除或不可访问 → 视为可清理(下一轮 unlink 会跳过)
        return True


def _walk_files(root: Path) -> list[Path]:
    """收集 root 下所有文件,按修改时间倒序(最新的在前)"""
    files = [p for p in root.rglob("*") if p.is_file()]
    files.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return files


def _prune_empty_dirs(root: Path, max_age: float) -> int:
    """自底向上收掉空目录(作者的层与作品的层都算),返回删除的目录数。

    只删「本身就是空」且 mtime 已过期的目录:正在落盘的目录刚建出来、
    mtime 很新,不会被误删;刚删完文件的目录 mtime 会被刷新,留到下一轮。
    root 自身不在 rglob 结果里,不会被删。
    """
    removed = 0
    # 目录另有 60s 下限:清理temp 0 时不能把"刚建出来、还没写第一个文件"的
    # 目录(下载落盘点先 mkdir 再 open)删掉
    dir_age = max(max_age, 60)
    dirs = [p for p in root.rglob("*") if p.is_dir()]
    # 深的先处理:子目录删掉后父目录才可能变空,同一轮里能被顺带收掉
    for path in sorted(dirs, key=lambda p: len(p.parts), reverse=True):
        try:
            if any(path.iterdir()):
                continue
            if not _file_is_stale(path, dir_age):
                continue
            path.rmdir()
            removed += 1
            logger.info(f"temp 清理: 删除空目录 {path}")
        except OSError:
            # 被占用 / 刚被别的进程删掉 → 留待下轮
            continue
    return removed


def clean_temp_files(
    sub: str = "",
    days: int | None = None,
    on_progress: Callable[..., None] | None = None,
) -> tuple[int, int]:
    """清理 temp[/sub] 下超过期限的文件与随之空掉的目录。

    Args:
        on_progress: 进度回调 `(done, total, phase, message, warning)`（关键字可省）——
            扫描阶段 total 未知，之后按"处理过的文件数"上报。跑在哪个线程调用都行
            （任务槽那边会切回事件循环）。**返回 False 表示调用方请求取消**，
            本函数会在处理完当前文件后尽早收尾（清理天然可中断，不硬杀线程）。

    Returns:
        (removed, total) — 删除的文件数、统计到的文件数（被取消时是"取消前"的进度）
    """

    def report(**kwargs) -> bool:
        if on_progress is None:
            return True  # 没有槽（纯函数调用）时不存在取消
        return bool(on_progress(**kwargs))

    root = get_temp_root(sub)
    if not root.is_dir():
        return 0, 0

    max_age = _max_age_seconds(days)
    report(phase="扫描中", unit="个文件", message="正在统计文件…")
    files = _walk_files(root)
    report(
        done=0, total=len(files), phase="清理中",
        unit="个文件", message=f"共 {len(files)} 个文件",
    )

    removed = 0
    canceled = False
    for index, path in enumerate(files, 1):
        if _file_is_stale(path, max_age):
            try:
                path.unlink(missing_ok=True)
                removed += 1
                logger.info(f"temp 清理: 删除 {path}")
            except OSError as e:
                # 文件被占用(如发送中)等场景,留待下轮 —— 也算"部分失败",报给任务槽
                logger.warning(f"temp 清理: 跳过 {path} ({e})")
                report(warning=f"跳过 {path.name}（被占用？）：{e}")
        if not report(done=index, total=len(files), message=f"已删除 {removed} 个"):
            logger.info("temp 清理: 收到取消请求，提前结束")
            canceled = True
            break

    if not canceled:
        report(phase="回收空目录", done=len(files), total=len(files))
        dirs = _prune_empty_dirs(root, max_age)
        if dirs:
            logger.info(f"temp 清理: 同时收掉 {dirs} 个空目录")
    return removed, len(files)


def temp_stats(sub: str = "") -> dict:
    """统计 temp[/sub] 目录:文件数、总大小(字节)、目录数(含作者/作品层)"""
    root = get_temp_root(sub)
    if not root.is_dir():
        return {"files": 0, "bytes": 0, "dirs": 0}
    files = _walk_files(root)
    total_bytes = sum(p.stat().st_size for p in files if p.exists())
    dirs = sum(1 for p in root.rglob("*") if p.is_dir())
    return {"files": len(files), "bytes": total_bytes, "dirs": dirs}


# ── 手动清理命令(manual 模式,auto 模式下也可用) ──────────────
# 清理是删除操作,要求管理及以上(check_manage:群主/群管理/超管)
clean_cmd = on_command("清理temp", aliases={"清理临时文件", "清temp"}, priority=10, block=True)
# 统计只读,对所有人开放
stats_cmd = on_command("temp统计", aliases={"temp状态"}, priority=10, block=True)


@clean_cmd.handle()
async def _handle_clean(bot: Bot, event: MessageEvent):
    # 鉴权:仅群主/群管理/超管可清理
    if not await check_manage(bot, event):
        await UniMessage.text("只有管理以上才能清理 temp 哦~").send(at_sender=True)
        return
    # 解析可选天数参数:/清理temp 3 → 清理 3 天前的文件
    text = str(event.message).strip()
    tokens = text.replace("/", " ").split()
    days = None
    if len(tokens) >= 2:
        try:
            days = int(tokens[1])
        except ValueError:
            await UniMessage.text("天数参数不合法,示例:/清理temp 3").send()
            return
    task, reused = submit_clean(
        days,
        title="清理 temp（命令）",
        trigger="command",
        actor=event.get_user_id(),
        group_id=getattr(event, "group_id", "") or "",
    )
    if reused:
        await UniMessage.text(
            "已经有一个清理任务在跑了,等它跑完再来吧~（进度可在网页端看）"
        ).send()
        return
    await TASKS.wait(task)  # 命令要"跑完再回复",与网页端的即提交即返回不同

    if task.status == "canceled":
        await UniMessage.text("清理被取消了（网页端点过取消）。").send()
        return
    if task.status != "done":
        await UniMessage.text(f"清理失败了:{task.error}").send()
        return
    removed, total = task.result or (0, 0)
    if total == 0:
        await UniMessage.text("temp 目录目前是空的,没什么可清理的~").send()
    else:
        await UniMessage.text(f"temp 清理完成:删除了 {removed} 个文件(共 {total} 个文件)。").send()


@stats_cmd.handle()
async def _handle_stats(event: MessageEvent):
    st = await run_stats()
    if st["files"] == 0:
        if st.get("dirs"):
            await UniMessage.text(
                f"temp 目录下没有文件了,还剩 {st['dirs']} 个空目录"
                f"(下次清理/清理temp 会一并收掉)。"
            ).send()
        else:
            await UniMessage.text("temp 目录目前是空的。").send()
    else:
        size_mb = st["bytes"] / 1024 / 1024
        await UniMessage.text(
            f"temp 目录:共 {st['files']} 个文件,占用 {size_mb:.1f} MB,"
            f"{st.get('dirs', 0)} 个目录。"
        ).send()


#: 任务槽名:命令 / Web / 每日定时三条路共用同一个槽,同槽只跑一个(防重复触发)
CLEANUP_SLOT = "video_analysis:temp_cleanup"


def submit_clean(
    days: int | None,
    *,
    title: str,
    trigger: str = "",
    actor: str = "",
    group_id: str = "",
) -> tuple[Task, bool]:
    """把清理提交到任务槽,立即返回 (任务, 是否复用了正在跑的任务)。

    任务体跑在线程池里(见 hexi.core.task_slots),进度实时推给网页端;
    同槽已有任务时不会重复清理 —— 服务端这一层才是"防连点"的真正保证。
    trigger / actor / group_id 是审计信息(谁点的:web 用户名 / command 的 QQ 号)。
    """
    return TASKS.submit(
        slot=CLEANUP_SLOT,
        title=title,
        plugin_id="video_analysis",
        trigger=trigger,
        actor=actor,
        group_id=group_id,
        func=lambda report: clean_temp_files(days=days, on_progress=report),
    )


async def run_stats() -> dict:
    """异步统计(命令 / Web 缓存面板共用)"""
    return await asyncio.to_thread(temp_stats)


# ── 自动模式:每日定时清理 ──────────────────────────────────────


def _register_auto_job() -> None:
    """按当前配置注册/撤销定时清理任务(幂等,可反复调用)。

    - auto 模式:注册每日 cron 清理任务(replace_existing 保证不重复)。
    - manual 模式:移除已注册的自动清理任务,仅保留手动命令。
    """
    from apscheduler.jobstores.base import JobLookupError  # noqa: PLC0415

    if cleanup_config.video_analysis_temp_cleanup_mode == "auto":
        add_job(
            _scheduled_cleanup,
            "cron",
            hour=cleanup_config.video_analysis_temp_cleanup_hour,
            minute=0,
            id="video_analysis_temp_cleanup",
            misfire_grace_time=3600,
            replace_existing=True,
        )
        logger.info(
            f"temp 自动清理已启用:每天 {cleanup_config.video_analysis_temp_cleanup_hour}:00"
        )
    else:
        try:
            scheduler.remove_job("video_analysis_temp_cleanup")
            logger.info("temp 自动清理已关闭(manual 模式)")
        except JobLookupError:
            pass


def reload_cleanup_config() -> None:
    """Web 保存清理配置后调用:按新配置重设自动清理任务。"""
    _register_auto_job()


async def _scheduled_cleanup() -> None:
    """每日定时清理 temp 目录(auto 模式);同样走任务槽(与手动清理互斥)"""
    logger.info("temp 自动清理: 开始")
    task, reused = submit_clean(None, title="清理 temp（每日定时）", trigger="schedule")
    if reused:
        logger.info("temp 自动清理: 已有清理任务在跑，跳过本次")
        return
    await TASKS.wait(task)
    if task.status != "done":
        logger.error(f"temp 自动清理未完成（{task.status}）: {task.error}")
        return
    removed, total = task.result or (0, 0)
    logger.info(f"temp 自动清理完成:删除 {removed} 个,当前共 {total} 个文件")


# 插件加载时按模式注册(auto)或提示(manual)
_register_auto_job()