"""체크포인트 archive 파일 정리.

`step_runner.save_checkpoint` / `invalidate_downstream` / `clear_checkpoint`가
기존 manifest.json을 `manifest_<YYYYMMDD>_<HHMMSS>.json`으로 **복사** 보관한다.
수 회 force/재실행 후 수천 개 누적 가능 (현재 entity_t2i 한 step에만 5000+).

**정리 정책**:
  - **이름에 label/uuid 접미사가 붙은 파일은 전부 유지** (사용자 스냅샷 보호).
    예: `manifest_20260420_154145_pre_p0_zoom_verify.json`
  - 자동 archive (`manifest_YYYYMMDD_HHMMSS.json`)만 step당 최근 N개 남기고 삭제.
  - 활성 `manifest.json`은 건드리지 않음.

사용:
    python backend/scripts/cleanup_checkpoint_archives.py            # dry-run
    python backend/scripts/cleanup_checkpoint_archives.py --apply    # 실제 삭제
    python backend/scripts/cleanup_checkpoint_archives.py --keep 5   # step당 5개 유지
"""
from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from typing import List, Tuple

_BACKEND = Path(__file__).resolve().parent.parent
PROJECT_ROOT = _BACKEND.parent
PROJECTS_ROOT = PROJECT_ROOT / "projects"

# `manifest_20260420_154238.json` (label 없는 auto archive)
AUTO_PATTERN = re.compile(r"^manifest_\d{8}_\d{6}\.json$")


def scan(keep: int = 3) -> List[Tuple[Path, int]]:
    """삭제 대상 (path, size) 리스트 반환."""
    to_remove: List[Tuple[Path, int]] = []
    if not PROJECTS_ROOT.exists():
        return to_remove
    # projects/<pid>/checkpoints/episodes/<eid>/<step>/
    for step_dir in PROJECTS_ROOT.glob("*/checkpoints/episodes/*/*"):
        if not step_dir.is_dir():
            continue
        auto_files = sorted(
            (f for f in step_dir.iterdir() if f.is_file() and AUTO_PATTERN.match(f.name)),
            key=lambda f: f.stat().st_mtime,
            reverse=True,
        )
        for f in auto_files[keep:]:
            to_remove.append((f, f.stat().st_size))
    return to_remove


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--apply", action="store_true", help="실제 삭제 실행 (생략 시 dry-run)")
    ap.add_argument("--keep", type=int, default=3, help="step당 auto archive 유지 개수 (기본 3)")
    args = ap.parse_args()

    to_remove = scan(keep=args.keep)
    total_bytes = sum(sz for _, sz in to_remove)
    print(f"Candidates: {len(to_remove)} auto-archive files ({total_bytes / (1024*1024):.1f} MB)")
    print(f"Policy: keep latest {args.keep} auto archives per step. Labeled/uuid archives untouched.")

    if not to_remove:
        return 0

    # step별 분포
    by_step = {}
    for p, _sz in to_remove:
        by_step[p.parent.name] = by_step.get(p.parent.name, 0) + 1
    print("\nDistribution (top 10):")
    for step, count in sorted(by_step.items(), key=lambda x: -x[1])[:10]:
        print(f"  {step}: {count}")

    if args.apply:
        removed = 0
        for p, _sz in to_remove:
            try:
                p.unlink()
                removed += 1
            except OSError as exc:
                print(f"  [fail] {p}: {exc}")
        print(f"\nRemoved {removed}/{len(to_remove)} files.")
    else:
        print("\n(dry-run) pass --apply to delete")
    return 0


if __name__ == "__main__":
    sys.exit(main())