"""
Distill XABA/XALZ assemblies.blob into individual .NET assemblies.
Format ref (reverse-engineered): https://github.com/dotnet/android/blob/main/src/Xamarin.Android.Build.Tasks/Utilities/AssemblyStore.cs
"""
import struct, sys, pathlib, lz4.block

if len(sys.argv) != 3:
    print("Usage: unpack_xaba.py <assemblies.blob> <out_dir>")
    sys.exit(1)

blob_path = pathlib.Path(sys.argv[1])
out_dir = pathlib.Path(sys.argv[2])
out_dir.mkdir(parents=True, exist_ok=True)

data = blob_path.read_bytes()
assert data[:4] == b"XABA", f"not a XABA blob, got {data[:4]!r}"

# Scan for XALZ markers (each = one compressed assembly)
i = 0
extracted = 0
while True:
    idx = data.find(b"XALZ", i)
    if idx == -1:
        break
    # XALZ: 4 magic + 4 descriptor_index + 4 uncompressed_size + N lz4_data
    desc_idx = struct.unpack_from("<I", data, idx + 4)[0]
    uncompressed = struct.unpack_from("<I", data, idx + 8)[0]
    payload_start = idx + 12
    # Find next XALZ to bound payload (or EOF)
    next_idx = data.find(b"XALZ", payload_start)
    payload_end = next_idx if next_idx != -1 else len(data)
    compressed = data[payload_start:payload_end]
    try:
        decompressed = lz4.block.decompress(compressed, uncompressed_size=uncompressed)
    except Exception as e:
        # Try without explicit size as fallback
        try:
            decompressed = lz4.block.decompress(compressed)
        except Exception as e2:
            print(f"[skip] desc={desc_idx} decompress failed: {e2}")
            i = payload_start
            continue
    out_file = out_dir / f"{desc_idx:03d}.dll"
    out_file.write_bytes(decompressed)
    extracted += 1
    print(f"[ok]   desc={desc_idx:03d}  size={len(decompressed)}  -> {out_file.name}")
    i = payload_start

print(f"\nExtracted {extracted} assemblies into {out_dir}")