#!/usr/bin/env python3
"""
Enable Beta APIs experiment in a Bedrock level.dat file.

Bedrock level.dat format:
  - 4 bytes: version (little-endian int32)
  - 4 bytes: payload length (little-endian int32)
  - rest: uncompressed little-endian NBT payload

This script patches the 'experiments' compound tag to add:
  - gametest = 1 (byte)
  - experiments_ever_used = 1 (byte)
  - saved_with_toggled_experiments = 1 (byte)

Requires: amulet-nbt (pip install amulet-nbt)
"""

import struct
import sys
import shutil
import amulet_nbt


def patch_level_dat(path: str) -> None:
    # Read the file
    with open(path, "rb") as f:
        header = f.read(8)
        payload = f.read()

    version, length = struct.unpack("<II", header)
    print(f"  Header: version={version}, payload_length={length}")
    print(f"  Actual payload size: {len(payload)}")

    # Parse little-endian NBT (Bedrock format)
    nbt = amulet_nbt.load(payload, little_endian=True)
    root = nbt.compound

    # Find experiments compound
    if "experiments" not in root:
        print("  ERROR: 'experiments' compound tag not found!")
        sys.exit(1)

    experiments = root["experiments"]
    print(f"  Current experiments: {dict((k, v) for k, v in experiments.items())}")

    changed = False

    # Set flags to 1
    for key in ("experiments_ever_used", "saved_with_toggled_experiments"):
        if key in experiments:
            if experiments[key].py_int == 0:
                experiments[key] = amulet_nbt.ByteTag(1)
                changed = True
                print(f"  Set {key} = 1")
        else:
            experiments[key] = amulet_nbt.ByteTag(1)
            changed = True
            print(f"  Added {key} = 1")

    # Add gametest experiment
    if "gametest" not in experiments:
        experiments["gametest"] = amulet_nbt.ByteTag(1)
        changed = True
        print("  Added gametest = 1")
    elif experiments["gametest"].py_int == 0:
        experiments["gametest"] = amulet_nbt.ByteTag(1)
        changed = True
        print("  Set gametest = 1")

    if not changed:
        print("  Already patched — no changes needed.")
        return

    # Serialize back (uncompressed, little-endian for Bedrock)
    new_payload = nbt.save_to(little_endian=True, compressed=False)

    # Build new header with updated length
    new_header = struct.pack("<II", version, len(new_payload))

    # Backup original
    backup_path = path + ".bak"
    shutil.copy2(path, backup_path)
    print(f"  Backup saved to {backup_path}")

    # Write patched file
    with open(path, "wb") as f:
        f.write(new_header)
        f.write(new_payload)

    print(f"  Patched! New payload size: {len(new_payload)}")


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <level.dat> [level.dat ...]")
        sys.exit(1)

    for path in sys.argv[1:]:
        print(f"\nPatching: {path}")
        patch_level_dat(path)

    print("\nDone!")


if __name__ == "__main__":
    main()