"""
Re-stamp every Rec Room CDN asset bundle for the Unity 2021.3.41f1 client.

    python ConvertBundles2021.py [SRC] [DST]        default D:\AssetBundles -> D:\ConvertedAssetBundles

What it does to each .assetbundle:
  - sets the engine version to 2021.3.41f1 on the UnityFS header and on every serialized
    file inside (header format 22 and UnityFS v8 already match what the client builds)
  - keeps the embedded type trees, so the player deserialises by field name
  - re-saves LZ4, then re-opens the result and checks every object still reads
Nothing else is touched. The folder layout is mirrored, so DST/<a>/<b>/(Item)/<bodytype>/x.assetbundle.

Output next to the bundles:
  DST\hashes.json   { "<file name>.assetbundle": "<base64 sha256 of the converted file>" }
                    - the UnityAssetHash / UnityAsset2Hash the catalog must carry for it
  DST\convert.log   one line per file; "FAIL" lines are the ones to look at
Already-converted files (present in DST with the right size) are skipped, so a rerun resumes.
"""
import os, sys, json, hashlib, base64, time
from concurrent.futures import ProcessPoolExecutor, as_completed

TARGET = "2021.3.41f1"
SRC = sys.argv[1] if len(sys.argv) > 1 else r"D:\AssetBundles"
DST = sys.argv[2] if len(sys.argv) > 2 else r"D:\ConvertedAssetBundles"


def convert_one(rel):
    import UnityPy
    src = os.path.join(SRC, rel)
    dst = os.path.join(DST, rel)
    try:
        env = UnityPy.load(src)
        bundles = [f for f in env.files.values() if type(f).__name__ == "BundleFile"]
        if not bundles:
            return rel, None, "not a UnityFS bundle"
        bf = bundles[0]
        before = bf.version_engine
        bf.version = 8
        bf.version_engine = TARGET
        n_sf = 0
        for f in bf.files.values():
            if type(f).__name__ == "SerializedFile":
                f.unity_version = TARGET
                n_sf += 1
        data = bf.save(packer="lz4")

        # Verify: re-open and read every object through its type tree.
        env2 = UnityPy.load(data)
        bad = 0
        for o in env2.objects:
            try:
                o.read_typetree()
            except Exception:
                bad += 1
        if bad:
            return rel, None, f"{bad} object(s) unreadable after convert"

        os.makedirs(os.path.dirname(dst), exist_ok=True)
        with open(dst, "wb") as fh:
            fh.write(data)
        h = base64.b64encode(hashlib.sha256(data).digest()).decode()
        return rel, h, f"ok {before} -> {TARGET} ({n_sf} file(s), {len(data)} B)"
    except Exception as ex:
        return rel, None, f"EXC {type(ex).__name__}: {ex}"


def main():
    todo = []
    for dp, _, fn in os.walk(SRC):
        for f in fn:
            if f.lower().endswith(".assetbundle"):
                rel = os.path.relpath(os.path.join(dp, f), SRC)
                out = os.path.join(DST, rel)
                if os.path.exists(out) and os.path.getsize(out) > 0:
                    continue
                todo.append(rel)
    os.makedirs(DST, exist_ok=True)
    hashes_path = os.path.join(DST, "hashes.json")
    hashes = json.load(open(hashes_path)) if os.path.exists(hashes_path) else {}
    print(f"{len(todo)} bundle(s) to convert ({len(hashes)} already hashed)")

    ok = fail = 0
    t0 = time.time()
    with open(os.path.join(DST, "convert.log"), "a", encoding="utf-8") as log, \
         ProcessPoolExecutor(max_workers=max(2, os.cpu_count() - 1)) as pool:
        futs = [pool.submit(convert_one, rel) for rel in todo]
        for i, fut in enumerate(as_completed(futs), 1):
            rel, h, msg = fut.result()
            if h:
                ok += 1
                hashes[os.path.basename(rel)] = h
            else:
                fail += 1
                msg = "FAIL " + msg
            log.write(f"{rel}\t{msg}\n")
            if i % 500 == 0 or i == len(futs):
                json.dump(hashes, open(hashes_path, "w"), indent=0)
                print(f"  {i}/{len(futs)}  ok={ok} fail={fail}  {time.time()-t0:.0f}s")
    json.dump(hashes, open(hashes_path, "w"), indent=0)
    print(f"done: {ok} converted, {fail} failed, hashes for {len(hashes)} file(s) in {hashes_path}")


if __name__ == "__main__":
    main()