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 目录占用情况

目录结构始终保留,占用中的文件自动跳过。
"""

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 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})")
    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}
    files = _walk_files(root)
    total_bytes = sum(p.stat().st_size for p in files if p.exists())
    return {"files": len(files), "bytes": total_bytes}


# ── 手动清理命令(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:
        await UniMessage.text("temp 目录目前是空的。").send()
    else:
        size_mb = st["bytes"] / 1024 / 1024
        await UniMessage.text(
            f"temp 目录:共 {st['files']} 个文件,占用 {size_mb:.1f} MB。"
        ).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()