wuling-overseas-dealer-expansion-skill/scripts/common/artifact_manager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Shared run artifact, backup, and cleanup helpers for the Wuling skill."""
from __future__ import annotations

import argparse
import json
import re
import shutil
import zipfile
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional


DEFAULT_BACKUP_KEEP = 10
DEFAULT_RETENTION_DAYS = 30


@dataclass
class CleanupItem:
    action: str
    path: Path
    size: int
    reason: str
    target: Optional[Path] = None


def project_root(start: Optional[Path] = None) -> Path:
    """Return the project root; prefer the current working tree over the skill folder."""
    start = (start or Path.cwd()).resolve()
    candidates = [start, *start.parents]
    for candidate in candidates:
        if (candidate / ".kimi").exists() or (candidate / "摩洛哥客户建联表-按渠道分类.xlsx").exists():
            return candidate
    return Path.cwd().resolve()


def today_key() -> str:
    return datetime.now().strftime("%Y%m%d")


def timestamp() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def safe_slug(value: str, default: str = "run") -> str:
    slug = re.sub(r"[^0-9A-Za-z_.-]+", "_", (value or "").strip()).strip("._-")
    return slug or default


def new_run_id(prefix: str = "run") -> str:
    return f"{safe_slug(prefix)}_{timestamp()}"


def run_dir(run_id: Optional[str] = None, root: Optional[Path] = None) -> Path:
    run_id = safe_slug(run_id or new_run_id())
    return project_root(root) / "runs" / today_key() / run_id


def backups_dir(root: Optional[Path] = None) -> Path:
    return project_root(root) / "backups" / today_key()


def archives_dir(root: Optional[Path] = None) -> Path:
    return project_root(root) / "archives"


def is_bare_filename(path: Path) -> bool:
    return not path.is_absolute() and str(path.parent) in {"", "."}


def resolve_artifact_path(
    raw: str,
    *,
    kind: str,
    default_name: str,
    run_id: Optional[str] = None,
    root: Optional[Path] = None,
) -> Path:
    """Resolve output artifacts into runs/ unless the user supplied a real path."""
    if raw:
        path = Path(raw).expanduser()
        if path.is_absolute():
            return path
        if not is_bare_filename(path):
            return (Path.cwd() / path).resolve()
        filename = path.name
    else:
        filename = default_name
    base = run_dir(run_id or f"{safe_slug(kind)}_{timestamp()}", root)
    return base / filename


def create_backup_once(
    workbook_path: Path,
    *,
    purpose: str,
    run_id: Optional[str] = None,
    root: Optional[Path] = None,
) -> Path:
    """Create one workbook backup for a run/purpose; return existing backup if present."""
    workbook_path = workbook_path.resolve()
    run_id = safe_slug(run_id or f"{safe_slug(purpose)}_{timestamp()}")
    backup_dir = backups_dir(root)
    backup_dir.mkdir(parents=True, exist_ok=True)
    backup_path = backup_dir / f"{workbook_path.stem}_backup_before_{safe_slug(purpose)}_{run_id}{workbook_path.suffix}"
    if not backup_path.exists():
        shutil.copy2(workbook_path, backup_path)
    return backup_path


def write_json(path: Path, data: Dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def file_size(path: Path) -> int:
    return path.stat().st_size if path.exists() and path.is_file() else 0


def iter_legacy_root_artifacts(root: Path) -> Iterable[Path]:
    patterns = [
        "email_preview*",
        "email_retry*",
        "*.sent-log.jsonl",
        "facebook_social_status_updates_*.json",
    ]
    seen = set()
    for pattern in patterns:
        for path in root.glob(pattern):
            key = str(path.resolve()).casefold()
            if key not in seen:
                seen.add(key)
                yield path


def build_cleanup_plan(
    *,
    root: Optional[Path] = None,
    retention_days: int = DEFAULT_RETENTION_DAYS,
    backup_keep: int = DEFAULT_BACKUP_KEEP,
) -> List[CleanupItem]:
    root = project_root(root)
    cutoff = datetime.now() - timedelta(days=retention_days)
    items: List[CleanupItem] = []

    tmp_dir = root / ".tmp"
    if tmp_dir.exists():
        for path in tmp_dir.iterdir():
            if path.is_file():
                items.append(CleanupItem("delete", path, file_size(path), ".tmp temporary file"))

    for path in iter_legacy_root_artifacts(root):
        if not path.is_file():
            continue
        reason = "legacy run artifact in project root"
        if path.name == "email_preview_summary_rows_357_419_20260724_170613.html" or path.name == "email_preview_summary_rows_357_419_20260724_170613.json":
            items.append(CleanupItem("delete", path, file_size(path), "invalid intermediate preview"))
            continue
        if re.match(r"email_retry_summary_rows_357_419_chunk_\d+_20260724\.json$", path.name):
            items.append(CleanupItem("delete", path, file_size(path), "superseded retry manifest"))
            continue
        target = root / "runs" / "legacy_202607" / path.name
        items.append(CleanupItem("move", path, file_size(path), reason, target))

    backups = sorted(
        [p for p in root.glob("*_backup_*.xlsx") if p.is_file()],
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    for idx, path in enumerate(backups):
        target = root / "backups" / "legacy_202607" / path.name
        # Keep recent backups discoverable by moving, not deleting. The newest
        # backup_keep rule applies to future cleanup runs inside backups/.
        items.append(CleanupItem("move", path, file_size(path), f"legacy workbook backup #{idx + 1}", target))

    backup_groups: Dict[str, List[Path]] = {}
    for path in (root / "backups").rglob("*_backup_*.xlsx") if (root / "backups").exists() else []:
        key = re.sub(r"_backup_before_.+$", "", path.name)
        backup_groups.setdefault(key, []).append(path)
    for group in backup_groups.values():
        group.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        for path in group[backup_keep:]:
            modified = datetime.fromtimestamp(path.stat().st_mtime)
            if modified < cutoff:
                target = root / "archives" / "old_backups" / (path.name + ".zip")
                items.append(CleanupItem("archive", path, file_size(path), "old backup beyond retention", target))

    return items


def apply_cleanup_plan(items: Iterable[CleanupItem], *, dry_run: bool = True) -> Dict[str, Any]:
    applied = []
    total_size = 0
    for item in items:
        total_size += item.size
        entry = {
            "action": item.action,
            "path": str(item.path),
            "target": str(item.target) if item.target else "",
            "size": item.size,
            "reason": item.reason,
        }
        if not dry_run:
            try:
                if item.action == "delete":
                    item.path.unlink(missing_ok=True)
                elif item.action == "move" and item.target:
                    item.target.parent.mkdir(parents=True, exist_ok=True)
                    if item.target.exists():
                        item.target = item.target.with_name(f"{item.target.stem}_{timestamp()}{item.target.suffix}")
                    shutil.move(str(item.path), str(item.target))
                    entry["target"] = str(item.target)
                elif item.action == "archive" and item.target:
                    item.target.parent.mkdir(parents=True, exist_ok=True)
                    with zipfile.ZipFile(item.target, "w", compression=zipfile.ZIP_DEFLATED) as archive:
                        archive.write(item.path, arcname=item.path.name)
                    item.path.unlink(missing_ok=True)
                entry["status"] = "ok"
            except Exception as exc:  # keep cleanup best-effort and auditable
                entry["status"] = "failed"
                entry["error"] = str(exc)
        applied.append(entry)
    return {
        "dry_run": dry_run,
        "count": len(applied),
        "total_size_bytes": total_size,
        "total_size_mb": round(total_size / 1024 / 1024, 3),
        "items": applied,
    }


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Preview or apply Wuling skill artifact cleanup.")
    parser.add_argument("--apply", action="store_true", help="Actually move/delete/archive files. Omit for dry-run.")
    parser.add_argument("--retention-days", type=int, default=DEFAULT_RETENTION_DAYS)
    parser.add_argument("--backup-keep", type=int, default=DEFAULT_BACKUP_KEEP)
    args = parser.parse_args(argv)
    plan = build_cleanup_plan(retention_days=args.retention_days, backup_keep=args.backup_keep)
    report = apply_cleanup_plan(plan, dry_run=not args.apply)
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())