#!/usr/bin/env python3
"""Classify and prune a character working lane.

A "lane" is a staged working folder like
`characters/female/lena_nude/hires_claude/` where each step script writes a full
`.blend` snapshot. Left alone those snapshots accumulate at ~80 MB per attempt.
This sorts a lane into three tiers and, with --apply, deletes the scratch tier.

    KEEP       recipe + registered artifacts. Never touched.
    MASTER     the few .blend files a registered artifact was built from,
               listed explicitly in the lane's `.lanekeep`, plus the chain head.
    SCRATCH    per-attempt .blend, .blend1 autosaves, QA render dirs, run logs,
               and byte-identical duplicate binaries. Regenerable from KEEP.
    UNKNOWN    a .blend with no step script beside it, so no recipe can rebuild
               it -- probably hand-authored. Reported, never auto-deleted.

Dry-run by default. Nothing is deleted without --apply.

    python tools/prune_lane.py characters/female/lena_nude/hires_claude
    python tools/prune_lane.py characters/female/lena_nude --recursive
    python tools/prune_lane.py <lane> --apply
"""
from __future__ import annotations

import argparse
import fnmatch
import hashlib
import shutil
import sys
from collections import defaultdict
from pathlib import Path

# --- classification rules -------------------------------------------------

# Extensions that are the recipe or a registered artifact. Never scratch.
KEEP_SUFFIXES = {".py", ".md", ".json", ".txt", ".npz", ".npy", ".jpg", ".jpeg"}

# Directories that hold regenerable QA renders / probe dumps.
SCRATCH_DIR_GLOBS = ("review", "review_*", "review[0-9]*", "dbg_*",
                     "probe_*", "beauty", "beauty_*")

# Files that are always regenerable output.
SCRATCH_FILE_GLOBS = ("*.blend1", "*.blend2", "*.log")

LANEKEEP = ".lanekeep"


def load_lanekeep(lane: Path) -> tuple[set[str], list[str]]:
    """Read the lane's keep manifest. Returns (names, comment lines)."""
    f = lane / LANEKEEP
    if not f.exists():
        return set(), []
    names, notes = set(), []
    for line in f.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        # Strip trailing inline comment: "06_final.blend  # geometry master"
        name = line.split("#", 1)[0].strip()
        if name:
            names.add(name)
    return names, notes


def human(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if abs(n) < 1024 or unit == "GB":
            return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1024.0
    return f"{n:.1f} GB"


def sha256(path: Path, chunk: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for block in iter(lambda: fh.read(chunk), b""):
            h.update(block)
    return h.hexdigest()


def matches(name: str, globs) -> bool:
    return any(fnmatch.fnmatch(name, g) for g in globs)


def has_step_scripts(lane: Path) -> bool:
    """True if this folder holds NN_*.py step scripts -- i.e. its .blend files
    are script output that can be rebuilt, not hand-authored work."""
    return any(lane.glob("[0-9][0-9]*_*.py"))


def classify(lane: Path, keep_names: set[str], repo: Path, exclude=()):
    """Walk the lane and bucket every entry.

    Returns (keep, master, scratch, unknown). `exclude` names subdirectories
    scanned as lanes in their own right, so they are not counted twice.
    """
    keep, master, scratch, unknown = [], [], [], []
    scripted = has_step_scripts(lane)

    # Scratch directories are pruned whole; don't descend into them.
    entries = []
    for path in sorted(lane.rglob("*")):
        parts = path.relative_to(lane).parts
        if parts[0] in exclude:
            continue  # belongs to a nested lane
        if any(matches(p, SCRATCH_DIR_GLOBS) for p in parts[:-1]):
            continue  # inside a dir already bucketed as scratch
        entries.append(path)

    for path in entries:
        rel = path.relative_to(lane).as_posix()
        if path.is_dir():
            if matches(path.name, SCRATCH_DIR_GLOBS) and rel not in keep_names:
                size = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
                scratch.append((rel + "/", size, "QA render dir"))
            continue

        if path.name == LANEKEEP:
            keep.append((rel, path.stat().st_size, "keep manifest"))
            continue

        size = path.stat().st_size

        if rel in keep_names or path.name in keep_names:
            master.append((rel, size, "pinned in .lanekeep"))
        elif path.suffix.lower() in KEEP_SUFFIXES:
            keep.append((rel, size, "recipe/input"))
        elif matches(path.name, SCRATCH_FILE_GLOBS):
            why = "autosave" if path.suffix.startswith(".blend") else "run log"
            (scratch if scripted else unknown).append((rel, size, why))
        elif path.suffix.lower() == ".blend":
            if scripted:
                scratch.append((rel, size, "per-attempt snapshot"))
            else:
                unknown.append((rel, size, "no step script here - cannot rebuild"))
        else:
            keep.append((rel, size, "binary (checked for dupes)"))

    scratch += find_duplicates(lane, keep, repo)
    keep = [k for k in keep if k[0] not in {s[0] for s in scratch}]
    return keep, master, scratch, unknown


def find_duplicates(lane: Path, keep, repo: Path):
    """Flag lane binaries that are byte-identical to a registered original.

    Also flags within-lane duplicate groups, keeping the oldest member.
    """
    dupes = []
    binaries = [lane / rel for rel, _, why in keep if why.startswith("binary")]
    if not binaries:
        return dupes

    # Reference set: the read-only registered originals.
    originals = {}
    orig_dir = repo / "characters" / "originals"
    if orig_dir.is_dir():
        for f in orig_dir.rglob("*"):
            if f.is_file() and f.suffix.lower() in (".glb", ".fbx"):
                originals[sha256(f)] = f.relative_to(repo).as_posix()

    by_hash = defaultdict(list)
    for f in binaries:
        by_hash[sha256(f)].append(f)

    for digest, files in by_hash.items():
        if digest in originals:
            for f in files:
                dupes.append((f.relative_to(lane).as_posix(), f.stat().st_size,
                              f"byte-identical to {originals[digest]}"))
        elif len(files) > 1:
            # The canonical name survives: registry grammar is lowercase
            # [a-z0-9_] only, so anything with " - Copy", " (1)" or spaces is
            # the accidental duplicate regardless of which one is older.
            def rank(p: Path):
                stem = p.stem.lower()
                junk = (" " in p.stem or "copy" in stem or "(1)" in stem)
                return (junk, p.stat().st_mtime)
            files.sort(key=rank)
            for f in files[1:]:
                dupes.append((f.relative_to(lane).as_posix(), f.stat().st_size,
                              f"byte-identical to {files[0].name}"))
    return dupes


def report(lane: Path, keep, master, scratch, unknown, notes, repo: Path) -> int:
    def block(title, rows, show_all):
        total = sum(r[1] for r in rows)
        print(f"\n  {title}  ({len(rows)} entries, {human(total)})")
        shown = rows if show_all else sorted(rows, key=lambda r: -r[1])[:12]
        for rel, size, why in sorted(shown, key=lambda r: r[0]):
            print(f"    {human(size):>9}  {rel:<46} {why}")
        if len(shown) < len(rows):
            print(f"    {'':>9}  ... and {len(rows) - len(shown)} more")
        return total

    try:
        shown_lane = lane.relative_to(repo).as_posix()
    except ValueError:
        shown_lane = str(lane)
    print(f"\n=== {shown_lane} ===")
    for n in notes:
        print(f"  note: {n}")
    block("KEEP     recipe + artifacts", keep, show_all=False)
    block("MASTER   pinned .blend", master, show_all=True)
    if unknown:
        block("UNKNOWN  decide by hand -- never auto-deleted", unknown, show_all=True)
    freed = block("SCRATCH  regenerable", scratch, show_all=False)
    return freed


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("lane", type=Path, help="working lane directory")
    ap.add_argument("--apply", action="store_true",
                    help="actually delete the scratch tier (default: dry run)")
    ap.add_argument("--recursive", action="store_true",
                    help="also treat each immediate subdirectory as its own lane")
    args = ap.parse_args()

    lane = args.lane.resolve()
    if not lane.is_dir():
        print(f"error: not a directory: {lane}", file=sys.stderr)
        return 1

    repo = Path(__file__).resolve().parent.parent
    sublanes = ([d for d in sorted(lane.iterdir()) if d.is_dir()]
                if args.recursive else [])
    # Parent scanned first, with its sublanes excluded so nothing is counted twice.
    plan = [(lane, {d.name for d in sublanes})] + [(d, set()) for d in sublanes]

    total_freed, all_scratch, unknown_total = 0, [], 0
    for ln, exclude in plan:
        keep_names, notes = load_lanekeep(ln)
        if not keep_names and any(ln.glob("*.blend")) and has_step_scripts(ln):
            notes.append(f"no {LANEKEEP} - every .blend here is treated as scratch")
        keep, master, scratch, unknown = classify(ln, keep_names, repo, exclude)
        if not (keep or master or scratch or unknown):
            continue
        total_freed += report(ln, keep, master, scratch, unknown, notes, repo)
        all_scratch += [(ln, rel) for rel, _, _ in scratch]
        unknown_total += sum(u[1] for u in unknown)

    print(f"\n{'=' * 60}")
    print(f"scratch tier: {len(all_scratch)} entries, {human(total_freed)} reclaimable")
    if unknown_total:
        print(f"unknown tier: {human(unknown_total)} held back for a manual call")

    if not args.apply:
        print("DRY RUN - nothing deleted. Re-run with --apply to prune.")
        return 0

    print("APPLYING - deleting scratch tier...")
    for ln, rel in all_scratch:
        target = ln / rel.rstrip("/")
        if target.is_dir():
            shutil.rmtree(target)
        elif target.exists():
            target.unlink()
    print(f"done - freed {human(total_freed)}")
    return 0


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