#!/usr/bin/env python3
from __future__ import annotations

import argparse
import shutil
import sys
import time
from dataclasses import dataclass
from pathlib import Path


DEFAULT_RETENTION_DAYS = 30


@dataclass(frozen=True)
class PruneCandidate:
    path: Path
    age_days: float
    size_bytes: int


def collect_candidates(
    runtime_root: Path,
    *,
    retention_days: int = DEFAULT_RETENTION_DAYS,
    now: float | None = None,
) -> list[PruneCandidate]:
    reference_time = time.time() if now is None else now
    cutoff_seconds = retention_days * 24 * 60 * 60
    if not runtime_root.exists():
        return []
    candidates: list[PruneCandidate] = []
    for path in sorted(runtime_root.iterdir()):
        if not path.is_dir() or not path.name.startswith("v1_run_"):
            continue
        age_seconds = max(0.0, reference_time - path.stat().st_mtime)
        if age_seconds <= cutoff_seconds:
            continue
        candidates.append(
            PruneCandidate(
                path=path,
                age_days=age_seconds / (24 * 60 * 60),
                size_bytes=_dir_size(path),
            )
        )
    return candidates


def prune_candidates(candidates: list[PruneCandidate], *, execute: bool) -> int:
    for candidate in candidates:
        if execute:
            shutil.rmtree(candidate.path)
        print(
            f"{'DELETE' if execute else 'DRY-RUN'}\t"
            f"{candidate.age_days:.1f}d\t"
            f"{_format_size(candidate.size_bytes)}\t"
            f"{candidate.path}"
        )
    return len(candidates)


def main() -> int:
    args = _parse_args()
    candidates = collect_candidates(
        args.runtime_root,
        retention_days=args.retention_days,
    )
    count = prune_candidates(candidates, execute=args.execute)
    action = "deleted" if args.execute else "would_delete"
    print(f"{action}={count} retention_days={args.retention_days} runtime_root={args.runtime_root}")
    return 0


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Prune local runtime/v1 run directories. DB and OSS records are not touched."
    )
    parser.add_argument("--runtime-root", default=Path("runtime/v1"), type=Path)
    parser.add_argument("--retention-days", default=DEFAULT_RETENTION_DAYS, type=int)
    parser.add_argument(
        "--execute",
        action="store_true",
        help="Actually delete matched run directories. Default is dry-run.",
    )
    return parser.parse_args()


def _dir_size(path: Path) -> int:
    return sum(item.stat().st_size for item in path.rglob("*") if item.is_file())


def _format_size(size_bytes: int) -> str:
    value = float(size_bytes)
    for unit in ["B", "KB", "MB", "GB"]:
        if value < 1024 or unit == "GB":
            return f"{value:.1f}{unit}"
        value /= 1024
    return f"{value:.1f}GB"


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