#!/usr/bin/env python3
"""
prune-identical.py - delete everything from an extracted tree that we already have,
leaving only genuine differences and additions.

    python3 prune-identical.py <extractedDir> [--apply]

Dry-run by default; nothing is deleted without --apply.

Why two baselines
-----------------
A .mw4 is not a zip of the source tree. Our own packer rewrites ~7,400 files on
the way in (.data .video .instance .contents .audio .damage .torso .subsystems
.engine .lights) and generates a further ~17,700 qualified records that have no
source file at all ('foo.data{gamemodel}', 'foo.contents[joint_hip]{armature}',
'bar.tga{hint}'). Comparing only against Content/ would therefore "keep" about
24,000 files we in fact already have.

So each candidate is tested, in order:

  1. same package + same entry path in OUR extracted packages, identical bytes
  2. same entry path anywhere in OUR extracted packages, identical bytes
  3. unqualified name + same path in our Content/ source tree, identical bytes

Any hit means we have that exact asset -> delete. Rule 3 is what recognises
raw pass-through assets (.tga, .wav, .erf, scripts); rules 1-2 are what
recognise the packer-generated forms.

Afterwards _merged/ is rebuilt from the survivors and empty directories removed.
"""
import sys, os, hashlib, collections, shutil

OURS_EXTRACTED = "/home/rich/Repositories/FS_Ours_extracted"
OUR_SOURCE = "/home/rich/Repositories/firestorm/Gameleap/mw4/Content"

PRECEDENCE = {"props": 50, "core": 40, "textures": 30, "maps": 20, "missions": 10}
MERGE_EXCLUDE = ("variants/", "pilots/")


def digest(path, cache={}):
    key = (path, os.path.getsize(path))
    if key not in cache:
        with open(path, "rb") as fh:
            cache[key] = hashlib.md5(fh.read()).hexdigest()
    return cache[key]


def index_paths(root, skip_top=()):
    """relative-lowercased-path -> absolute path. No hashing (lazy)."""
    out = {}
    for dp, dirs, fs in os.walk(root):
        if dp == root:
            dirs[:] = [d for d in dirs if d.lower() not in skip_top]
        for f in fs:
            p = os.path.join(dp, f)
            out[os.path.relpath(p, root).replace("\\", "/").lower()] = p
    return out


def package_of(rel):
    """Split '<pkgdir>/<entryPath>' for an extracted file.

    Package directories are: core, props, textures, maps/<x>, Missions/<x>,
    Variants/<x>, Pilots/Tesla/options
    """
    parts = rel.split("/")
    head = parts[0].lower()
    if head in ("core", "props", "textures"):
        return parts[0], "/".join(parts[1:])
    if head in ("maps", "missions", "variants") and len(parts) > 2:
        return "/".join(parts[:2]), "/".join(parts[2:])
    if head == "pilots" and len(parts) > 3:
        return "/".join(parts[:3]), "/".join(parts[3:])
    return parts[0], "/".join(parts[1:])


def merged_precedence(pkg):
    return PRECEDENCE.get(pkg.split("/")[0].lower(), 0)


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    root = sys.argv[1]
    apply_ = "--apply" in sys.argv
    merged_root = os.path.join(root, "_merged")

    print("indexing our extracted packages ...", flush=True)
    ours_pkg = index_paths(OURS_EXTRACTED, skip_top={"_merged"})
    ours_by_entry = {}
    for rel, p in ours_pkg.items():
        if rel.startswith("_"):
            continue
        _pkg, entry = package_of(rel)
        ours_by_entry.setdefault(entry, []).append(p)
    print(f"  {len(ours_pkg)} files, {len(ours_by_entry)} distinct entry paths")

    print("indexing our Content source tree ...", flush=True)
    src = index_paths(OUR_SOURCE)
    print(f"  {len(src)} files")

    stats = collections.Counter()
    doomed, kept = [], []

    for dp, _dirs, fs in os.walk(root):
        if dp == merged_root or dp.startswith(merged_root + os.sep):
            continue
        for f in fs:
            p = os.path.join(dp, f)
            rel = os.path.relpath(p, root).replace("\\", "/")
            if rel.startswith("_"):
                continue
            pkg, entry = package_of(rel)
            h = digest(p)
            reason = None

            cand = ours_pkg.get(rel.lower())
            if cand and digest(cand) == h:
                reason = "same-package"
            if reason is None:
                for c in ours_by_entry.get(entry.lower(), ()):
                    if digest(c) == h:
                        reason = "other-package"
                        break
            if reason is None and "{" not in f and "[" not in f:
                c = src.get(entry.lower())
                if c and digest(c) == h:
                    reason = "our-source"

            if reason:
                stats["deleted:" + reason] += 1
                doomed.append((rel, reason))
            else:
                stats["kept"] += 1
                kept.append((rel, pkg, entry, h))

    total = len(doomed) + len(kept)
    print(f"\nexamined {total} files")
    for k, v in sorted(stats.items()):
        print(f"  {v:7d}  {k}")
    print(f"\n  -> would delete {len(doomed)}, keep {len(kept)}")

    with open(os.path.join(root, "_survivors.tsv"), "w", encoding="utf-8") as fh:
        fh.write("path\tpackage\tentry\tmd5\n")
        for rel, pkg, entry, h in sorted(kept):
            fh.write(f"{rel}\t{pkg}\t{entry}\t{h}\n")

    if not apply_:
        print(f"\nDRY RUN - nothing changed. Survivor list written to "
              f"{os.path.join(root, '_survivors.tsv')}. Re-run with --apply.")
        return

    for rel, _reason in doomed:
        os.remove(os.path.join(root, rel))

    with open(os.path.join(root, "_pruned.tsv"), "w", encoding="utf-8") as fh:
        fh.write("deletedPath\tmatchedVia\n")
        for rel, reason in sorted(doomed):
            fh.write(f"{rel}\t{reason}\n")

    # rebuild _merged from survivors
    shutil.rmtree(merged_root, ignore_errors=True)
    best = {}
    for rel, pkg, entry, h in kept:
        if pkg.lower().startswith(MERGE_EXCLUDE):
            continue
        pr = merged_precedence(pkg)
        cur = best.get(entry.lower())
        if cur is None or pr > cur[0]:
            best[entry.lower()] = (pr, entry, os.path.join(root, rel))
    for _pr, entry, srcfile in best.values():
        dest = os.path.join(merged_root, entry)
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        os.link(srcfile, dest)

    removed = 0
    for dp, dirs, fs in os.walk(root, topdown=False):
        if dp == root or not os.path.isdir(dp):
            continue
        if not os.listdir(dp):
            os.rmdir(dp)
            removed += 1

    print(f"\ndeleted {len(doomed)} files, removed {removed} empty directories")
    print(f"_merged rebuilt with {len(best)} paths")
    print(f"log: {os.path.join(root, '_pruned.tsv')}")


if __name__ == "__main__":
    main()