import mmap
import struct
import sys

if len(sys.argv) < 2:
    print("Usage: python3 mmap_patch.py <filename.gguf>")
    sys.exit(1)

filename = sys.argv[1]
print(f"Memory-mapping {filename} (Uses almost 0 RAM)...")

# Open the file in read/write binary mode
with open(filename, "r+b") as f:
    # Memory map the file directly to disk
    mm = mmap.mmap(f.fileno(), 0)
    
    patches_applied = 0
    
    # 1. Fix block_count
    target_key1 = b"qwen35moe.block_count"
    idx1 = mm.find(target_key1)
    if idx1 != -1:
        val_offset1 = idx1 + len(target_key1)
        val_type1 = struct.unpack("<I", mm[val_offset1:val_offset1+4])[0]
        
        if val_type1 == 4: # UINT32
            data_offset1 = val_offset1 + 4
            old_val1 = struct.unpack("<I", mm[data_offset1:data_offset1+4])[0]
            if old_val1 == 41:
                print(f"  [PATCH] block_count: {old_val1} -> 40")
                mm[data_offset1:data_offset1+4] = struct.pack("<I", 40)
                patches_applied += 1
            else:
                print(f"  [OK] block_count is already {old_val1}")
    else:
        print("  [WARN] block_count not found!")
        
    # 2. Fix nextn_predict_layers
    target_key2 = b"qwen35moe.nextn_predict_layers"
    idx2 = mm.find(target_key2)
    if idx2 != -1:
        val_offset2 = idx2 + len(target_key2)
        val_type2 = struct.unpack("<I", mm[val_offset2:val_offset2+4])[0]
        
        if val_type2 == 4: # UINT32
            data_offset2 = val_offset2 + 4
            old_val2 = struct.unpack("<I", mm[data_offset2:data_offset2+4])[0]
            if old_val2 == 1:
                print(f"  [PATCH] nextn_predict_layers: {old_val2} -> 0")
                mm[data_offset2:data_offset2+4] = struct.pack("<I", 0)
                patches_applied += 1
            else:
                print(f"  [OK] nextn_predict_layers is already {old_val2}")
    else:
        print("  [WARN] nextn_predict_layers not found!")

    mm.close()
    
print(f"\nDone! Applied {patches_applied} patches in-place.")