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 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

require("nonebot_plugin_apscheduler")
from nonebot_plugin_apscheduler import scheduler

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) -> tuple[int, int]:
    """清理 temp[/sub] 下超过期限的文件与随之空掉的目录。

    Returns:
        (removed, total) — 删除的文件数、统计到的文件总数
    """
    root = get_temp_root(sub)
    if not root.is_dir():
        return 0, 0

    max_age = _max_age_seconds(days)
    files = _walk_files(root)
    removed = 0
    for path in files:
        if not _file_is_stale(path, max_age):
            continue
        try:
            path.unlink(missing_ok=True)
            removed += 1
            logger.info(f"temp 清理: 删除 {path}")
        except OSError as e:
            # 文件被占用(如发送中)等场景,留待下轮
            logger.warning(f"temp 清理: 跳过 {path} ({e})")

    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
    removed, total = await _run_clean(days)
    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()


async def _run_clean(days: int | None = None) -> tuple[int, int]:
    """执行清理并记录日志(供命令与定时任务共用)"""
    removed, total = await asyncio.to_thread(clean_temp_files, "", days)
    return removed, total


async def _run_stats() -> dict:
    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":
        scheduler.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 自动清理: 开始")
    try:
        removed, total = await asyncio.to_thread(clean_temp_files)
        logger.info(f"temp 自动清理完成:删除 {removed} 个,当前共 {total} 个文件")
    except Exception as e:
        logger.exception(f"temp 自动清理失败: {e}")


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