scripts/scratch_cleanup.py

from __future__ import annotations

import argparse
import ctypes
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence


FO_DELETE = 0x0003
FOF_SILENT = 0x0004
FOF_NOCONFIRMATION = 0x0010
FOF_ALLOWUNDO = 0x0040
FOF_NOERRORUI = 0x0400
SKIP_FILE_NAMES: frozenset[str] = frozenset({".gitkeep"})


class SHFILEOPSTRUCTW(ctypes.Structure):
    _fields_ = [
        ("hwnd", ctypes.c_void_p),
        ("wFunc", ctypes.c_uint),
        ("pFrom", ctypes.c_wchar_p),
        ("pTo", ctypes.c_wchar_p),
        ("fFlags", ctypes.c_ushort),
        ("fAnyOperationsAborted", ctypes.c_bool),
        ("hNameMappings", ctypes.c_void_p),
        ("lpszProgressTitle", ctypes.c_wchar_p),
    ]


@dataclass(frozen=True)
class CleanupTarget:
    path: Path
    is_directory: bool
    size_bytes: int
    file_count: int
    directory_count: int


@dataclass(frozen=True)
class CleanupError:
    path: Path
    message: str


@dataclass(frozen=True)
class CleanupResult:
    scratch_root: Path
    dry_run: bool
    targets: tuple[CleanupTarget, ...]
    total_files: int
    total_directories: int
    total_bytes: int
    recycled_paths: tuple[Path, ...]
    errors: tuple[CleanupError, ...]

    @property
    def error_paths(self) -> tuple[Path, ...]:
        return tuple(error.path for error in self.errors)


def _format_bytes(size_bytes: int) -> str:
    return f"{size_bytes / 1_048_576:.2f} MB"


def _resolve_path(path: Path) -> Path:
    return path.expanduser().resolve(strict=False)


def _is_within(parent: Path, child: Path) -> bool:
    parent_text = os.path.normcase(str(_resolve_path(parent)))
    child_text = os.path.normcase(str(_resolve_path(child)))
    try:
        common_path = os.path.commonpath([parent_text, child_text])
    except ValueError:
        return False
    return common_path == parent_text


def _validate_scratch_root(scratch_root: Path) -> Path:
    resolved_root = _resolve_path(scratch_root)
    if resolved_root.name.lower() != "scratch":
        raise ValueError(f"Cleanup target must be a directory named 'scratch': {resolved_root}")
    if not resolved_root.exists():
        raise FileNotFoundError(f"Scratch directory not found: {resolved_root}")
    if not resolved_root.is_dir():
        raise ValueError(f"Scratch path is not a directory: {resolved_root}")
    return resolved_root


def _measure_target(path: Path) -> tuple[int, int, int]:
    if path.is_file():
        return path.stat().st_size, 1, 0

    total_bytes = 0
    total_files = 0
    total_directories = 1
    for child_path in path.rglob("*"):
        if child_path.is_file():
            total_bytes += child_path.stat().st_size
            total_files += 1
            continue
        if child_path.is_dir():
            total_directories += 1
    return total_bytes, total_files, total_directories


def _build_cleanup_target(path: Path) -> CleanupTarget:
    size_bytes, file_count, directory_count = _measure_target(path)
    return CleanupTarget(
        path=path,
        is_directory=path.is_dir(),
        size_bytes=size_bytes,
        file_count=file_count,
        directory_count=directory_count,
    )


def _iter_cleanup_targets(scratch_root: Path) -> tuple[CleanupTarget, ...]:
    targets: list[CleanupTarget] = []
    for path in sorted(scratch_root.iterdir(), key=lambda candidate: candidate.name.lower()):
        if path.name in SKIP_FILE_NAMES:
            continue
        targets.append(_build_cleanup_target(path))
    return tuple(targets)


def _move_to_recycle_bin(path: Path) -> None:
    if os.name != "nt":
        raise OSError("Recycle Bin cleanup is only supported on Windows.")

    path_buffer = f"{_resolve_path(path)}\0\0"
    operation = SHFILEOPSTRUCTW()
    operation.wFunc = FO_DELETE
    operation.pFrom = path_buffer
    operation.fFlags = FOF_ALLOWUNDO | FOF_NOCONFIRMATION | FOF_NOERRORUI | FOF_SILENT

    result = ctypes.windll.shell32.SHFileOperationW(ctypes.byref(operation))
    if result != 0:
        raise OSError(result, f"Recycle Bin move failed for: {path}")
    if operation.fAnyOperationsAborted:
        raise OSError(f"Recycle Bin move aborted for: {path}")


def cleanup_scratch(scratch_root: Path, *, dry_run: bool = True) -> CleanupResult:
    resolved_root = _validate_scratch_root(scratch_root)
    targets = _iter_cleanup_targets(resolved_root)

    recycled_paths: list[Path] = []
    errors: list[CleanupError] = []

    if not dry_run:
        for target in targets:
            if not _is_within(resolved_root, target.path):
                errors.append(
                    CleanupError(
                        path=target.path,
                        message=f"Refused to clean outside scratch root: {target.path}",
                    )
                )
                continue
            try:
                _move_to_recycle_bin(target.path)
                recycled_paths.append(target.path)
            except OSError as exc:
                errors.append(CleanupError(path=target.path, message=str(exc)))

    return CleanupResult(
        scratch_root=resolved_root,
        dry_run=dry_run,
        targets=targets,
        total_files=sum(target.file_count for target in targets),
        total_directories=sum(target.directory_count for target in targets),
        total_bytes=sum(target.size_bytes for target in targets),
        recycled_paths=tuple(recycled_paths),
        errors=tuple(errors),
    )


def _build_argument_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Reset scratch artifacts by moving top-level entries to Recycle Bin.")
    parser.add_argument(
        "--root",
        type=Path,
        default=Path(__file__).resolve().parent.parent / "scratch",
        help="Scratch directory to inspect or clean.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview reclaimable targets without changing anything.",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Move cleanup targets to Recycle Bin.",
    )
    return parser


def _print_target(target: CleanupTarget) -> None:
    kind = "dir " if target.is_directory else "file"
    print(
        "    - "
        f"[{kind}] {target.path.name} | "
        f"{_format_bytes(target.size_bytes)} | "
        f"{target.file_count} file(s) | "
        f"{target.directory_count} dir(s)"
    )


def _print_result(result: CleanupResult) -> None:
    mode_label = "DRY-RUN" if result.dry_run else "APPLY"
    print(f"[{mode_label}] Scratch root: {result.scratch_root}")
    print(f"  Targets: {len(result.targets)}")
    print(f"  Files: {result.total_files}")
    print(f"  Directories: {result.total_directories}")
    print(f"  Reclaimable: {_format_bytes(result.total_bytes)}")
    if result.targets:
        print("  Target list:")
        for target in result.targets:
            _print_target(target)
    else:
        print("  Target list: empty")

    if result.dry_run:
        return

    print(f"  Recycled: {len(result.recycled_paths)}")
    if result.errors:
        print(f"  Errors: {len(result.errors)}")
        for error in result.errors:
            print(f"    - {error.path.name}: {error.message}")
        return
    print("  Errors: 0")


def main(argv: Sequence[str] | None = None) -> int:
    parser = _build_argument_parser()
    args = parser.parse_args(argv)
    if args.apply and args.dry_run:
        parser.error("Choose only one mode: --dry-run or --apply.")

    dry_run = not args.apply
    try:
        result = cleanup_scratch(args.root, dry_run=dry_run)
    except (FileNotFoundError, ValueError, OSError) as exc:
        print(f"[ERROR] {exc}")
        return 1

    _print_result(result)
    if result.errors:
        return 1
    return 0


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