UmbraAudio/scripts/nas_lifecycle.py

#!/usr/bin/env python3
"""NAS sync + local prune lifecycle for UmbraAudio.

Primary store: local SSD. Secondary: NAS mirror. On each run:
  1. ensure the NAS is mounted
  2. rsync local -> NAS (mirror everything)
  3. prune locally only what is verified present on the NAS,
     respecting age and free-space targets

Fetch fallback for pruned files lives in umbraaudio.storage.
"""

from __future__ import annotations

import argparse
import json
import logging
import shutil
import subprocess
import sys
import time
from pathlib import Path

DEFAULT_LOCAL_ROOT = Path("/Volumes/1TB_APL_Storage/UmbraAudio")
DEFAULT_NAS_ROOT = Path("/Volumes/Umbra/UmbraAudio")
MOUNT_SCRIPT = Path.home() / ".umbra/mount_nas.sh"
CONFIG_PATH = Path.home() / ".umbra/nas_config.json"

RSYNC_EXCLUDES = [".git/", "node_modules/", ".idea/", "__pycache__/", ".venv/", ".python/"]

# Never pruned in normal mode: live code, config, and active TTS models
PROTECTED_DIRS = {"umbraaudio", "scripts", "tests", "specs", "Qwen3-TTS",
                  ".umbra", ".git", ".venv", ".python", ".pytest_cache"}

DEFAULT_POLICY = {
    "min_age_days": 30,       # only prune files untouched this long
    "free_target_gb": 150,    # prune until at least this much is free (normal mode)
    "emergency_free_gb": 40,  # below this, protected dirs may be pruned too
}

logger = logging.getLogger("nas_lifecycle")
LOCK_FILE = Path("/tmp/umbra_nas_lifecycle.lock")


class Locked(RuntimeError):
    pass


def load_policy() -> dict:
    policy = dict(DEFAULT_POLICY)
    try:
        cfg = json.loads(CONFIG_PATH.read_text())
        policy.update(cfg.get("prune_policy", {}))
    except (OSError, json.JSONDecodeError):
        pass
    return policy


def nas_mounted(nas_root: Path = DEFAULT_NAS_ROOT) -> bool:
    return nas_root.is_dir() and any(nas_root.iterdir())


def ensure_mounted() -> bool:
    if nas_mounted():
        return True
    logger.info("NAS not mounted; running %s", MOUNT_SCRIPT)
    try:
        subprocess.run(["bash", str(MOUNT_SCRIPT)], check=True,
                       capture_output=True, timeout=120)
    except (subprocess.SubprocessError, OSError) as e:
        logger.error("mount failed: %s", e)
        return False
    return nas_mounted()


def sync(local_root: Path = DEFAULT_LOCAL_ROOT,
         nas_root: Path = DEFAULT_NAS_ROOT) -> bool:
    cmd = ["rsync", "-a", "--delete"]
    for exc in RSYNC_EXCLUDES:
        cmd += ["--exclude", exc]
    cmd += [f"{local_root}/", f"{nas_root}/"]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        logger.error("rsync failed (%d): %s", result.returncode, result.stderr[-500:])
        return False
    return True


def verified_on_nas(local_file: Path, local_root: Path,
                    nas_root: Path) -> bool:
    mirror = nas_root / local_file.relative_to(local_root)
    try:
        ls, ln = local_file.stat(), mirror.stat()
    except OSError:
        return False
    return ls.st_size == ln.st_size and int(ls.st_mtime) == int(ln.st_mtime)


def free_bytes(path: Path) -> int:
    return shutil.disk_usage(path).free


def prune_plan(local_root: Path, nas_root: Path, policy: dict,
               emergency: bool = False) -> list[Path]:
    """Files safe to delete now, largest-first among eligible."""
    min_age_s = policy["min_age_days"] * 86400
    cutoff = time.time() - min_age_s
    candidates: list[tuple[int, Path]] = []
    for p in local_root.rglob("*"):
        if not p.is_file() or p.is_symlink():
            continue
        rel_parents = set(p.relative_to(local_root).parts[:-1])
        if not emergency and rel_parents & PROTECTED_DIRS:
            continue
        st = p.stat()
        if st.st_mtime > cutoff:
            continue
        if not verified_on_nas(p, local_root, nas_root):
            continue
        candidates.append((st.st_size, p))
    candidates.sort(reverse=True)
    return [p for _, p in candidates]


def prune(local_root: Path = DEFAULT_LOCAL_ROOT,
          nas_root: Path = DEFAULT_NAS_ROOT,
          policy: dict | None = None) -> dict:
    policy = policy or load_policy()
    target = policy["free_target_gb"] * 1024**3
    emergency_at = policy["emergency_free_gb"] * 1024**3
    freed = 0
    deleted = 0
    emergency = free_bytes(local_root) < emergency_at
    for p in prune_plan(local_root, nas_root, policy, emergency=emergency):
        if free_bytes(local_root) >= target and freed > 0:
            break
        size = p.stat().st_size
        try:
            p.unlink()
        except OSError as e:
            logger.warning("could not unlink %s: %s", p, e)
            continue
        freed += size
        deleted += 1
        parent = p.parent
        if parent != local_root and not any(parent.iterdir()):
            parent.rmdir()
    stats = {"deleted": deleted, "freed_gb": round(freed / 1024**3, 2),
             "emergency": emergency,
             "free_gb": round(free_bytes(local_root) / 1024**3, 1)}
    logger.info("prune: %s", stats)
    return stats


def status(local_root: Path = DEFAULT_LOCAL_ROOT,
           nas_root: Path = DEFAULT_NAS_ROOT) -> dict:
    lf = lb = 0
    for p in local_root.rglob("*"):
        if p.is_file() and not p.is_symlink():
            try:
                lf += 1
                lb += p.stat().st_size
            except OSError:
                pass
    usage = shutil.disk_usage(nas_root) if nas_mounted(nas_root) else None
    return {
        "local": {"files": lf, "gb": round(lb / 1024**3, 2),
                  "free_gb": round(free_bytes(local_root) / 1024**3, 1)},
        "nas": {"mounted": nas_mounted(nas_root),
                "free_gb": round(usage.free / 1024**3, 1) if usage else None},
    }


def _acquire_lock() -> int:
    import os
    fd = os.open(LOCK_FILE, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
    os.write(fd, str(os.getpid()).encode())
    return fd


def run_cycle(local_root: Path = DEFAULT_LOCAL_ROOT,
              nas_root: Path = DEFAULT_NAS_ROOT) -> bool:
    if not ensure_mounted():
        return False
    busy = subprocess.run(["pgrep", "-f",
                           f"rsync .*{nas_root}"], capture_output=True)
    if busy.returncode == 0:
        logger.info("another rsync to the NAS is running; deferring cycle")
        return True
    try:
        _acquire_lock()
    except FileExistsError:
        logger.info("cycle already in progress; skipping")
        return True
    try:
        if not sync():
            return False
        prune()
    finally:
        LOCK_FILE.unlink(missing_ok=True)
    return True


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=["sync", "prune", "status", "cycle"])
    args = parser.parse_args(argv)
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    if args.command == "status":
        print(json.dumps(status(), indent=2))
        return 0
    if args.command == "sync":
        ok = ensure_mounted() and sync()
    elif args.command == "prune":
        print(json.dumps(prune(), indent=2))
        return 0
    else:
        ok = run_cycle()
    return 0 if ok else 1


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