disk_cache_separate.py

"""
Cache files matching .cacheignore patterns to a backup location,
replacing originals with symlinks for space savings.

Patterns follow .gitignore semantics (handled by the ``pathspec`` library).

Usage::

    uv run python disk_cache_separate.py [OPTIONS]

Environment Variables
--------------------
CACHE_LOC : str
    Override the default cache directory (/media/cbs/backup/cache).
"""
import argparse
import os
import shutil
import subprocess
from pathlib import Path

MIN_CACHE_FREE_MB = 100

import pathspec

from rich.progress import (
    BarColumn,
    Progress,
    TaskID,
    TextColumn,
    TimeRemainingColumn,
)

from log_setup import log_setup

logger = log_setup(__name__)

DEFAULT_CACHE_LOC = Path("/media/cbs/backup/cache")
CACHEIGNORE_NAME = ".cacheignore"


def resolve_cache_loc(cli_value: Path | None) -> Path:
    """Resolve cache location — CLI arg > env var > default."""
    if cli_value is not None:
        return cli_value.resolve()
    env = os.environ.get("CACHE_LOC")
    if env:
        return Path(env).resolve()
    return DEFAULT_CACHE_LOC.resolve()


def find_cacheignore(start_dir: Path) -> Path | None:
    """Search CWD first, then *start_dir* and its parents for ``.cacheignore``."""
    cwd = Path.cwd()
    candidate = cwd / CACHEIGNORE_NAME
    if candidate.is_file():
        return candidate
    for parent in [start_dir, *start_dir.parents]:
        candidate = parent / CACHEIGNORE_NAME
        if candidate.is_file():
            return candidate
    return None


def load_pathspec(ignore_file: Path | None, root_dir: Path) -> pathspec.PathSpec:
    """Load ignore patterns from *ignore_file* or auto-discover near *root_dir*.

    Returns an empty ``PathSpec`` when no file is found.
    """
    if ignore_file is None:
        found = find_cacheignore(root_dir)
        if found is None:
            logger.info("No .cacheignore found for %s — caching nothing", root_dir)
            return pathspec.PathSpec.from_lines("gitwildmatch", [])
        ignore_file = found

    logger.info("Using patterns from %s", ignore_file)
    lines = ignore_file.read_text(encoding="utf-8").splitlines()
    return pathspec.PathSpec.from_lines("gitwildmatch", lines)


def get_home_relative_path(item: Path) -> Path:
    """Return *item* path relative to ``$HOME``."""
    home = Path.home()
    return item.relative_to(home)


def cache_item(
    item: Path,
    cache_loc: Path,
    dry_run: bool,
) -> bool:
    """Move *item* into ``cache_loc / home_relative_path`` and symlink it back.

    Returns ``True`` on success, ``False`` on skip / error.
    """
    home_rel = get_home_relative_path(item)
    dest = cache_loc / home_rel

    if dest.exists():
        logger.debug("Already cached: %s", home_rel)
        return False

    try:
        dest.parent.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        logger.warning("Failed to create directory for %s: %s", home_rel, exc)
        return False

    if dry_run:
        logger.info("[dry-run] mv  %s/%s → %s", "~", home_rel, dest)
        logger.info("[dry-run] ln  %s ← %s/%s", dest, "~", home_rel)
        return True

    try:
        shutil.move(str(item), str(dest))
    except OSError as exc:
        logger.warning("Failed to move %s: %s", item, exc)
        return False

    # Symlink at original location → cache (absolute path is clearest cross-device)
    try:
        item.symlink_to(dest)
    except OSError as exc:
        logger.warning("Moved but failed to symlink %s: %s", item, exc)
        return False

    logger.info("Cached: ~/%s → %s", home_rel, dest)
    return True


def _collect_matches(
    root_dir: Path,
    spec: pathspec.PathSpec,
) -> tuple[list[Path], list[Path]]:
    """Walk *root_dir*, match items against *spec*, return (matched_dirs, matched_files)."""
    root = root_dir.resolve()
    matched_dirs: list[Path] = []
    matched_files: list[Path] = []

    for dirpath_str, dirnames, filenames in os.walk(root):
        dirpath = Path(dirpath_str)

        prune = []
        for name in dirnames:
            rel = dirpath.relative_to(root) / name
            if spec.match_file(str(rel) + "/") or spec.match_file(str(rel)):
                matched_dirs.append(dirpath / name)
                prune.append(name)
        for name in prune:
            dirnames.remove(name)

        for name in filenames:
            rel = dirpath.relative_to(root) / name
            if spec.match_file(str(rel)):
                matched_files.append(dirpath / name)

    return matched_dirs, matched_files


def walk_and_cache(
    root_dir: Path,
    spec: pathspec.PathSpec,
    cache_loc: Path,
    dry_run: bool,
    progress: Progress,
    overall_task: TaskID,
) -> None:
    """Collect and cache matched items under *root_dir* with progress reporting."""
    root = root_dir.resolve()

    logger.info("Scanning %s", root)

    matched_dirs, matched_files = _collect_matches(root, spec)

    total = len(matched_dirs) + len(matched_files)
    if total == 0:
        progress.update(overall_task, description=f"[green]{root.name}: nothing to cache")
        return

    root_task = progress.add_task(f"[cyan]{root.name}", total=total)
    for item in matched_dirs + matched_files:
        progress.update(root_task, description=str(get_home_relative_path(item)))
        cache_item(item, cache_loc, dry_run)
        progress.advance(root_task)
        progress.advance(overall_task)
    progress.update(root_task, description=f"[green]{root.name}: done ({total} items)")


def is_symlink_into_cache(item: Path, cache_loc: Path) -> bool:
    """Return ``True`` iff *item* is a symlink targeting somewhere under *cache_loc*."""
    if not item.is_symlink():
        return False
    try:
        raw = item.readlink()
        target = raw.resolve() if raw.is_absolute() else (item.parent / raw).resolve()
    except OSError:
        return False
    return str(target).startswith(str(cache_loc.resolve()))


def find_cached_ancestor(item: Path, cache_loc: Path) -> Path | None:
    """Walk up from *item*; return nearest ancestor whose symlink targets *cache_loc*."""
    resolved_cache = cache_loc.resolve()
    for candidate in [item, *item.parents]:
        if not candidate.is_symlink():
            continue
        try:
            raw = candidate.readlink()
            target = raw.resolve() if raw.is_absolute() else (candidate.parent / raw).resolve()
        except OSError:
            continue
        if str(target).startswith(str(resolved_cache)):
            return candidate
    return None


def restore_item(item: Path, cache_loc: Path, dry_run: bool) -> bool:
    """Restore *item* (or its cached parent) from cache back to original location.

    Returns ``True`` on success, ``False`` on skip / error.
    """
    cached_at = find_cached_ancestor(item, cache_loc)
    if cached_at is None:
        logger.warning("Not cached: %s", item)
        return False

    home_rel = get_home_relative_path(cached_at)
    if cached_at != item:
        logger.info("%s is inside cached parent ~/%s — restoring parent", item, home_rel)

    try:
        raw = cached_at.readlink()
        cache_entry = raw.resolve() if raw.is_absolute() else (cached_at.parent / raw).resolve()
    except OSError as exc:
        logger.warning("Cannot read symlink %s: %s", cached_at, exc)
        return False

    if not cache_entry.exists():
        logger.warning("Cache entry missing for ~/%s: %s", home_rel, cache_entry)
        return False

    if dry_run:
        logger.info("[dry-run] backup  %s → /tmp/backup/%s", cache_entry, home_rel)
        logger.info("[dry-run] restore  ~/%s ← %s", home_rel, cache_entry)
        return True

    # Backup cache entry to /tmp/backup before touching anything
    backup_path = Path("/tmp/backup") / home_rel
    try:
        if cache_entry.is_dir():
            backup_path.parent.mkdir(parents=True, exist_ok=True)
            subprocess.run(
                ["rsync", "-a", "--ignore-errors", str(cache_entry) + "/", str(backup_path) + "/"],
                capture_output=True,
            )
        else:
            backup_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(str(cache_entry), str(backup_path))
    except OSError as exc:
        logger.warning("Backup failed for ~/%s: %s — aborting restore", home_rel, exc)
        return False
    logger.info("Backed up: %s → %s", cache_entry, backup_path)

    try:
        os.unlink(cached_at)
    except OSError as exc:
        logger.warning("Failed to remove symlink %s: %s", cached_at, exc)
        return False

    try:
        shutil.move(str(cache_entry), str(cached_at))
    except OSError as exc:
        logger.warning("Failed to restore ~/%s: %s", home_rel, exc)
        return False

    logger.info("Restored: ~/%s ← %s (backup at %s)", home_rel, cache_entry, backup_path)
    return True


def _collect_restorable(resolved: Path) -> list[Path]:
    """Walk *resolved* (cache dir) and return paths under $HOME whose symlinks
    point into *resolved*."""
    items: list[Path] = []
    for dirpath_str, dirnames, filenames in os.walk(resolved):
        dirpath = Path(dirpath_str)
        rel = dirpath.relative_to(resolved)

        for name in dirnames:
            home_path = Path.home() / rel / name
            if is_symlink_into_cache(home_path, resolved):
                items.append(home_path)
        for name in filenames:
            home_path = Path.home() / rel / name
            if is_symlink_into_cache(home_path, resolved):
                items.append(home_path)
    return items


def restore_all(cache_loc: Path, dry_run: bool, progress: Progress, overall_task: TaskID) -> None:
    """Restore every cached item whose original symlink still exists."""
    resolved = cache_loc.resolve()
    if not resolved.is_dir():
        logger.warning("Cache directory not found: %s", resolved)
        return

    logger.info("Scanning cache: %s", resolved)

    items = _collect_restorable(resolved)
    total = len(items)
    if total == 0:
        progress.update(overall_task, description="[green]restore: nothing to restore")
        return

    restore_task = progress.add_task("[cyan]restore", total=total)
    for home_path in items:
        progress.update(restore_task, description=str(get_home_relative_path(home_path)))
        restore_item(home_path, resolved, dry_run)
        progress.advance(restore_task)
        progress.advance(overall_task)
    progress.update(restore_task, description=f"[green]restore: done ({total} items)")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Cache files matching .cacheignore patterns to a backup location.",
    )
    parser.add_argument(
        "--cache-loc",
        type=Path,
        default=None,
        help="Cache directory (default: %(default)s)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would be done without making changes",
    )
    parser.add_argument(
        "--root",
        type=Path,
        action="append",
        default=[],
        help="Root directory to scan (can be repeated, default: $HOME)",
    )
    parser.add_argument(
        "--cacheignore",
        type=Path,
        default=None,
        help="Path to .cacheignore pattern file (default: auto-discover from root)",
    )
    parser.add_argument(
        "--restore",
        nargs="*",
        default=None,
        type=Path,
        help=(
            "Restore cached items back to their original locations. "
            "Provide zero or more paths. If no paths given, restore all cached items."
        ),
    )
    return parser


def purge_stale_entries(
    roots: list[Path],
    spec_cache: dict[str, pathspec.PathSpec],
    cache_loc: Path,
    dry_run: bool,
    progress: Progress,
    overall_task: TaskID,
) -> None:
    """Restore symlinked items under *roots* that point into *cache_loc*
    but no longer match the current ignore patterns.

    The cache entry is backed up to /tmp/backup/ before restoration for safety.
    """
    resolved = cache_loc.resolve()
    if not resolved.is_dir():
        return

    restorable = _collect_restorable(resolved)
    if not restorable:
        return

    backup_dir = Path("/tmp/backup")
    stale: list[Path] = []
    for item in restorable:
        for root_dir in roots:
            root = root_dir.resolve()
            if str(item).startswith(str(root)):
                spec = spec_cache.get(str(root))
                if spec is None or not spec.patterns:
                    continue
                rel = item.relative_to(root)
                matches = spec.match_file(str(rel) + "/") or spec.match_file(str(rel))
                if not matches:
                    stale.append(item)
                break  # matched a root, no need to check others

    if not stale:
        return

    logger.info("Found %d stale cached item(s) — restoring", len(stale))
    stale_task = progress.add_task("[yellow]restoring stale", total=len(stale))

    for item in stale:
        home_rel = get_home_relative_path(item)
        progress.update(stale_task, description=str(home_rel))

        try:
            raw = item.readlink()
            cache_target = raw.resolve() if raw.is_absolute() else (item.parent / raw).resolve()
        except OSError as exc:
            logger.warning("Cannot read symlink %s: %s", item, exc)
            progress.advance(stale_task)
            progress.advance(overall_task)
            continue

        if dry_run:
            logger.info("[dry-run] backup  %s → /tmp/backup/%s", cache_target, home_rel)
            logger.info("[dry-run] restore ~/%s ← %s", home_rel, cache_target)
        else:
            backup_path = backup_dir / home_rel
            try:
                if cache_target.is_dir():
                    backup_path.parent.mkdir(parents=True, exist_ok=True)
                    subprocess.run(
                        ["rsync", "-a", "--ignore-errors", str(cache_target) + "/", str(backup_path) + "/"],
                        capture_output=True,
                    )
                else:
                    backup_path.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(cache_target), str(backup_path))
            except OSError as exc:
                logger.warning("Backup failed for ~/%s: %s — skipping restore", home_rel, exc)
                progress.advance(stale_task)
                progress.advance(overall_task)
                continue
            logger.info("Backed up: %s → %s", cache_target, backup_path)

            try:
                os.unlink(item)
            except OSError as exc:
                logger.warning("Failed to remove symlink %s: %s", item, exc)
                progress.advance(stale_task)
                progress.advance(overall_task)
                continue

            try:
                shutil.move(str(cache_target), str(item))
            except OSError as exc:
                logger.warning("Failed to restore ~/%s: %s", home_rel, exc)
                try:
                    item.symlink_to(cache_target)
                except OSError:
                    pass
                progress.advance(stale_task)
                progress.advance(overall_task)
                continue

            logger.info("Restored: ~/%s ← %s (backup at %s)", home_rel, cache_target, backup_path)

        progress.advance(stale_task)
        progress.advance(overall_task)

    progress.update(stale_task, description=f"[green]restored {len(stale)} stale item(s)")


def _count_matches(
    roots: list[Path],
    cacheignore: Path | None,
    spec_cache: dict[str, pathspec.PathSpec],
) -> int:
    """Quick pre-scan: return total matched items across all roots."""
    total = 0
    for root_dir in roots:
        root = root_dir.resolve()
        if not root.is_dir():
            continue
        spec = load_pathspec(cacheignore, root)
        spec_cache[str(root)] = spec
        if not spec.patterns:
            continue
        dirs, files = _collect_matches(root, spec)
        total += len(dirs) + len(files)
    return total


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    cache_loc = resolve_cache_loc(args.cache_loc)

    # ── Restore mode ────────────────────────────────────────────────
    if args.restore is not None:
        total_restore = 0
        if not args.restore:
            resolved = cache_loc.resolve()
            if resolved.is_dir():
                items = _collect_restorable(resolved)
                total_restore = len(items)
        else:
            total_restore = len(args.restore)

        with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeRemainingColumn(),
        ) as progress:
            overall = progress.add_task("[cyan]Overall", total=max(total_restore, 1))
            if not args.restore:
                restore_all(cache_loc, args.dry_run, progress, overall)
            else:
                restore_task = progress.add_task("[cyan]restore", total=len(args.restore))
                for item in args.restore:
                    progress.update(restore_task, description=str(get_home_relative_path(item.resolve())))
                    restore_item(item.resolve(), cache_loc, args.dry_run)
                    progress.advance(restore_task)
                    progress.advance(overall)
                progress.update(restore_task, description="[green]restore: done")
            progress.update(overall, description="[green]All restores complete")
        return

    roots = args.root or [Path.home()]

    logger.info("Cache location: %s", cache_loc)
    if args.dry_run:
        logger.info("DRY RUN — no files will be moved")
    else:
        free_mb = shutil.disk_usage(cache_loc).free // 1_048_576
        if free_mb < MIN_CACHE_FREE_MB:
            logger.error(
                "Only %d MB free on cache device — need at least %d MB, aborting",
                free_mb,
                MIN_CACHE_FREE_MB,
            )
            raise SystemExit(1)

    # Load specs before any operation
    spec_cache: dict[str, pathspec.PathSpec] = {}
    _count_matches(roots, args.cacheignore, spec_cache)

    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        TimeRemainingColumn(),
    ) as progress:
        overall = progress.add_task("[cyan]Overall", total=0)
        progress.update(overall, description="[yellow]Purging stale cache entries…")

        # Phase 1: restore cached items that no longer match current patterns
        purge_stale_entries(roots, spec_cache, cache_loc, args.dry_run, progress, overall)

        # Phase 2: cache items that match current patterns
        grand_total = _count_matches(roots, args.cacheignore, spec_cache)
        progress.update(overall, total=max(grand_total, 1))

        for root_dir in roots:
            root = root_dir.resolve()
            if not root.is_dir():
                logger.warning("Not a directory, skipping: %s", root)
                continue

            spec = spec_cache.get(str(root))
            if spec is None:
                spec = load_pathspec(args.cacheignore, root)

            if not spec.patterns:
                logger.info("No patterns for %s — skipping", root)
                continue

            walk_and_cache(root, spec, cache_loc, args.dry_run, progress, overall)

        progress.update(overall, description="[green]All tasks complete")


if __name__ == "__main__":
    main()