"""patch_v13_test.py <pid> [--revert]

v13: Hook SmartBox::DoPickupEvent to call unparent_children on the
unequipped CPhysicsObj's children before continuing.

This fixes the weapon-switch leak: each unequip currently calls
unset_parent (clearing the weapon's link to player) but never iterates
the weapon's OWN children (visual effects, sub-physobjs). Those keep
parent=weapon and accumulate in weapon->children forever.

Mechanism:
  Replace the `CALL unset_parent` at 0x004522bf with a CALL to a tiny
  12-byte thunk that does:
    PUSH ECX
    CALL unset_parent       (0x00513F70)
    POP ECX
    JMP unparent_children   (0x00513FE0)
  Net effect: both unset_parent AND unparent_children fire, then
  control returns to DoPickupEvent's next instruction (leave_world).
"""
import argparse, ctypes, ctypes.wintypes as wt, struct, sys


PATCH_SITE_VA       = 0x004522BF
ORIG_CALL_BYTES     = bytes([0xE8, 0xAC, 0x1C, 0x0C, 0x00])  # CALL unset_parent
UNSET_PARENT_VA     = 0x00513F70
UNPARENT_CHILDREN_VA = 0x00513FE0

PROCESS_VM_READ = 0x10
PROCESS_VM_WRITE = 0x20
PROCESS_VM_OPERATION = 0x8
PROCESS_QUERY_INFORMATION = 0x400
MEM_COMMIT_RESERVE = 0x1000 | 0x2000
PAGE_EXECUTE_READWRITE = 0x40


k32 = ctypes.windll.kernel32
OpenProcess = k32.OpenProcess
OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]; OpenProcess.restype = wt.HANDLE
CloseHandle = k32.CloseHandle
CloseHandle.argtypes = [wt.HANDLE]; CloseHandle.restype = wt.BOOL
VirtualAllocEx = k32.VirtualAllocEx
VirtualAllocEx.argtypes = [wt.HANDLE, ctypes.c_void_p, ctypes.c_size_t, wt.DWORD, wt.DWORD]
VirtualAllocEx.restype = wt.LPVOID
WriteProcessMemory = k32.WriteProcessMemory
WriteProcessMemory.argtypes = [wt.HANDLE, wt.LPVOID, wt.LPCVOID, ctypes.c_size_t,
                               ctypes.POINTER(ctypes.c_size_t)]
WriteProcessMemory.restype = wt.BOOL
ReadProcessMemory = k32.ReadProcessMemory
ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID, ctypes.c_size_t,
                              ctypes.POINTER(ctypes.c_size_t)]
ReadProcessMemory.restype = wt.BOOL
VirtualProtectEx = k32.VirtualProtectEx
VirtualProtectEx.argtypes = [wt.HANDLE, wt.LPVOID, ctypes.c_size_t, wt.DWORD,
                             ctypes.POINTER(wt.DWORD)]
VirtualProtectEx.restype = wt.BOOL


def read_bytes(h, addr, n):
    buf = (ctypes.c_ubyte * n)()
    sz = ctypes.c_size_t(0)
    if not ReadProcessMemory(h, addr, buf, n, ctypes.byref(sz)):
        raise OSError(f"read 0x{addr:08x} err={ctypes.get_last_error()}")
    return bytes(buf[:sz.value])


def write_bytes(h, addr, data):
    old = wt.DWORD(0)
    if not VirtualProtectEx(h, addr, len(data), PAGE_EXECUTE_READWRITE, ctypes.byref(old)):
        raise OSError(f"VirtualProtectEx 0x{addr:08x} err={ctypes.get_last_error()}")
    sz = ctypes.c_size_t(0)
    ok = WriteProcessMemory(h, addr, data, len(data), ctypes.byref(sz))
    err = ctypes.get_last_error() if not ok else 0
    restored = wt.DWORD(0)
    VirtualProtectEx(h, addr, len(data), old.value, ctypes.byref(restored))
    if not ok:
        raise OSError(f"write 0x{addr:08x} err={err}")


def build_thunk(thunk_va):
    out = bytearray()
    out += bytes([0x51])  # push ecx
    # call unset_parent
    rel = UNSET_PARENT_VA - (thunk_va + len(out) + 5)
    out += bytes([0xE8]) + struct.pack("<i", rel)
    out += bytes([0x59])  # pop ecx
    # jmp unparent_children
    rel = UNPARENT_CHILDREN_VA - (thunk_va + len(out) + 5)
    out += bytes([0xE9]) + struct.pack("<i", rel)
    return bytes(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pid", type=int)
    ap.add_argument("--revert", action="store_true")
    args = ap.parse_args()
    h = OpenProcess(PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION |
                    PROCESS_QUERY_INFORMATION, False, args.pid)
    if not h:
        print(f"OpenProcess({args.pid}) err={ctypes.get_last_error()}"); sys.exit(2)

    cur = read_bytes(h, PATCH_SITE_VA, 5)
    print(f"PID {args.pid}")
    print(f"  patch site @ 0x{PATCH_SITE_VA:08x}: {cur.hex()}")

    if args.revert:
        if cur == ORIG_CALL_BYTES:
            print(f"    already original"); CloseHandle(h); return
        write_bytes(h, PATCH_SITE_VA, ORIG_CALL_BYTES)
        print(f"    reverted; now: {read_bytes(h, PATCH_SITE_VA, 5).hex()}")
        CloseHandle(h); return

    if cur != ORIG_CALL_BYTES:
        if cur[0] == 0xE8:
            print(f"    already has a CALL — likely already patched"); CloseHandle(h); sys.exit(3)
        print(f"    UNEXPECTED — expected {ORIG_CALL_BYTES.hex()}"); CloseHandle(h); sys.exit(4)

    thunk_page = VirtualAllocEx(h, None, 0x40, MEM_COMMIT_RESERVE, PAGE_EXECUTE_READWRITE)
    if not thunk_page:
        print(f"VirtualAllocEx failed err={ctypes.get_last_error()}"); sys.exit(5)
    print(f"  thunk page @ 0x{thunk_page:08x}")
    thunk = build_thunk(thunk_page)
    print(f"  thunk bytes: {thunk.hex()}")

    sz = ctypes.c_size_t(0)
    if not WriteProcessMemory(h, thunk_page, thunk, len(thunk), ctypes.byref(sz)):
        print(f"write thunk failed err={ctypes.get_last_error()}"); sys.exit(6)
    after_thunk = read_bytes(h, thunk_page, len(thunk))
    if after_thunk != thunk:
        print(f"    thunk MISMATCH"); sys.exit(7)

    rel = thunk_page - (PATCH_SITE_VA + 5)
    new_call = bytes([0xE8]) + struct.pack("<i", rel)
    write_bytes(h, PATCH_SITE_VA, new_call)
    after = read_bytes(h, PATCH_SITE_VA, 5)
    print(f"    patch site now: {after.hex()}")
    if after != new_call:
        print(f"    PATCH MISMATCH"); sys.exit(8)
    print(f"  OK")
    CloseHandle(h)


if __name__ == "__main__":
    main()