import struct

MODULE = '/tmp/amdgpu_installed.ko'
OUTPUT = '/tmp/amdgpu_v3.ko'

old_lea = bytes.fromhex('49f7ffff')  # -(2231) for MAX=2230
new_lea = bytes.fromhex('3bf6ffff')  # -(2501) for MAX=2500
old_cmp = bytes.fromhex('a6f8ffff')  # -(1882) for range 350-2230
new_cmp = bytes.fromhex('98f7ffff')  # -(2152) for range 350-2500

with open(MODULE, 'rb') as f:
    data = bytearray(f.read())

print(f'Module size: {len(data)} bytes')
total_patches = 0

# Search ALL occurrences of old_lea and patch nearby old_cmp
pos = 0
while True:
    idx = data.find(old_lea, pos)
    if idx == -1:
        break
    
    ctx = bytes(data[max(0,idx-8):idx+16]).hex()
    print(f'\nFound old_lea at 0x{idx:x}: ...{ctx}...')
    
    # Look for matching cmp within 20 bytes
    found_cmp = False
    for j in range(idx+4, min(idx+30, len(data)-4)):
        if data[j:j+4] == old_cmp:
            print(f'  Found matching cmp at 0x{j:x}')
            data[idx:idx+4] = new_lea
            data[j:j+4] = new_cmp
            print(f'  Patched lea: {old_lea.hex()} -> {new_lea.hex()}')
            print(f'  Patched cmp: {old_cmp.hex()} -> {new_cmp.hex()}')
            total_patches += 2
            found_cmp = True
            break
    
    if not found_cmp:
        # Maybe cmp is after more bytes, or uses a different range
        # Check wider range and also look for the original cmp value
        for j in range(max(0, idx-30), min(idx+60, len(data)-4)):
            if j == idx: continue
            val = data[j:j+4]
            if val == old_cmp:
                print(f'  Found cmp at 0x{j:x} (wider search)')
                data[idx:idx+4] = new_lea
                data[j:j+4] = new_cmp
                total_patches += 2
                found_cmp = True
                break
        
        if not found_cmp:
            print(f'  WARNING: No matching cmp found near 0x{idx:x}')
            # Patch just the lea anyway — it might still help
            data[idx:idx+4] = new_lea
            print(f'  Patched lea only: {old_lea.hex()} -> {new_lea.hex()}')
            total_patches += 1
    
    pos = idx + 1

# Double check: any remaining old patterns?
remaining_lea = bytes(data).count(old_lea)
remaining_cmp = bytes(data).count(old_cmp)
print(f'\nRemaining old_lea (49f7ffff): {remaining_lea}')
print(f'Remaining old_cmp (a6f8ffff): {remaining_cmp}')

with open(OUTPUT, 'wb') as f:
    f.write(data)

print(f'\nTotal patches applied: {total_patches}')
print(f'Output: {OUTPUT}')