"""probe_gxtri3mesh.py <pid>
Walk committed memory, find every instance of GXTri3Mesh<u16> vtable
0x007ED3B0. For each instance, read:
  +0x000  vfptr (= 0x007ED3B0)
  +0x26c  refcount (per Release: if param_1[0x26c] == 1, free)
  +0x270  aggregate-IUnknown pointer (if non-NULL, the object is delegating
          AddRef/Release elsewhere — used as alternate refcount path)
Goal: characterize refcount distribution. A leak would manifest as many
instances with high refcount (engine holds many lingering refs).
"""
import ctypes, ctypes.wintypes as wt, sys, struct

PROCESS_VM_READ = 0x10
PROCESS_QUERY_INFORMATION = 0x400
MEM_COMMIT = 0x1000
MEM_PRIVATE = 0x20000
PAGE_READWRITE = 0x4
PAGE_EXECUTE_READWRITE = 0x40

k = ctypes.windll.kernel32
k.OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]; k.OpenProcess.restype = wt.HANDLE
k.ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID, ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)]
k.ReadProcessMemory.restype = wt.BOOL

class MBI(ctypes.Structure):
    _fields_ = [("BaseAddress", ctypes.c_void_p),
                ("AllocationBase", ctypes.c_void_p),
                ("AllocationProtect", wt.DWORD),
                ("RegionSize", ctypes.c_size_t),
                ("State", wt.DWORD),
                ("Protect", wt.DWORD),
                ("Type", wt.DWORD)]
k.VirtualQueryEx.argtypes = [wt.HANDLE, ctypes.c_void_p, ctypes.POINTER(MBI), ctypes.c_size_t]
k.VirtualQueryEx.restype = ctypes.c_size_t

VTABLE = 0x007ED3B0

def rd(h, va, n):
    buf = (ctypes.c_ubyte * n)(); sz = ctypes.c_size_t(0)
    if not k.ReadProcessMemory(h, va, buf, n, ctypes.byref(sz)): return None
    return bytes(buf[:sz.value])

def rd_u32(h, va):
    b = rd(h, va, 4); return struct.unpack('<I', b)[0] if b else None

pid = int(sys.argv[1])
h = k.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_INFORMATION, False, pid)
if not h: print(f"OpenProcess err={ctypes.get_last_error()}"); sys.exit(2)

# Pass 1: locate every vtable match.
matches = []
addr = 0
mbi = MBI()
while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
    base = mbi.BaseAddress or 0
    size = mbi.RegionSize
    if (mbi.State == MEM_COMMIT and mbi.Type == MEM_PRIVATE and
        (mbi.Protect & 0xFF) in (PAGE_READWRITE, PAGE_EXECUTE_READWRITE)):
        off = 0
        while off < size:
            chunk = min(4*1024*1024, size - off)
            b = rd(h, base + off, chunk)
            if b:
                n_words = len(b) // 4
                arr = struct.unpack(f"<{n_words}I", b[:n_words*4])
                for i, v in enumerate(arr):
                    if v == VTABLE:
                        matches.append(base + off + i*4)
            off += chunk
    addr = base + size
    if addr >= 0x80000000: break

print(f"pid {pid}: total GXTri3Mesh vtable matches = {len(matches)}")

# Pass 2: for the first N (so we don't blow time), read refcount + aggregate.
SAMPLE = min(2000, len(matches))
buckets = {'rc=0':0, 'rc=1':0, 'rc=2':0, 'rc=3':0, 'rc=4-9':0, 'rc=10-99':0, 'rc>=100':0,
           'aggregate':0, 'bad':0}
max_rc = 0
sum_rc = 0
counted = 0
samples = []
for i, va in enumerate(matches[:SAMPLE]):
    agg = rd_u32(h, va + 0x270)
    rc  = rd_u32(h, va + 0x26c)
    if agg is None or rc is None:
        buckets['bad'] += 1; continue
    if agg != 0:
        buckets['aggregate'] += 1; continue
    # Plausible refcounts are 0-millions; reject only insane values
    if rc > 0x40000000:  # garbage / pointer-shaped value at refcount offset
        buckets['bad'] += 1; continue
    counted += 1
    sum_rc += rc
    if rc > max_rc: max_rc = rc
    if rc == 0: buckets['rc=0'] += 1
    elif rc == 1: buckets['rc=1'] += 1
    elif rc == 2: buckets['rc=2'] += 1
    elif rc == 3: buckets['rc=3'] += 1
    elif rc < 10: buckets['rc=4-9'] += 1
    elif rc < 100: buckets['rc=10-99'] += 1
    else: buckets['rc>=100'] += 1
    if rc >= 5 and len(samples) < 8:
        samples.append((va, rc))

print(f"sampled {SAMPLE} (of {len(matches)}) vtable matches; standalone refcount={counted}, aggregates={buckets['aggregate']}, garbage={buckets['bad']}")
for kbl, v in buckets.items():
    print(f"  {kbl:10s}: {v}")
if counted:
    print(f"  avg refcount (standalone): {sum_rc/counted:.2f}  max: {max_rc}")
for va, rc in samples:
    print(f"  high-refcount sample va=0x{va:08x} rc={rc}")