"""count_leak_classes.py <pid>

Count live instances of all known leak-candidate classes for diagnostic
delta analysis. Emit timestamped, machine-parseable output.
"""
import argparse, ctypes, ctypes.wintypes as wt, struct, sys, time


VTABLES = {
    # GraphicsResource family (texture/surface cache)
    "GraphicsResource":       0x0079bf64,
    "RenderSurface":          0x0079a67c,
    "RenderTexture":          0x0079c198,
    "CSurface":               0x007ca4dc,
    "ImgTex":                 0x007cab04,
    "RenderTextureD3D":       0x00801a18,
    "RenderSurfaceD3D":       0x00801a94,
    # UI
    "UIElement_UIItem":       0x007c0498,
    "NoticeHandler_subvt":    0x007ccb60,
    # CObjCell family
    "CObjCell_primary":       0x007c98e8,
    "CObjCell_subvt":         0x0079385c,
    "CEnvCell_primary":       0x007c9a60,
    # Physics
    "CPhysicsObj":            0x007c78ec,
    # Already-patched (reference)
    "Palette":                0x007caa08,
    "CGfxObj":                0x007ca418,
    "D3DXMesh":               0x007ed3b0,
}


PROCESS_VM_READ = 0x10
PROCESS_QUERY_INFORMATION = 0x400
MEM_COMMIT = 0x1000
MEM_PRIVATE = 0x20000


class MBI(ctypes.Structure):
    _fields_ = [('BaseAddress', ctypes.c_void_p),
                ('AllocationBase', ctypes.c_void_p),
                ('AllocationProtect', wt.DWORD),
                ('PartitionId', wt.WORD),
                ('RegionSize', ctypes.c_size_t),
                ('State', wt.DWORD),
                ('Protect', wt.DWORD),
                ('Type', wt.DWORD)]


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
k.VirtualQueryEx.argtypes = [wt.HANDLE, ctypes.c_void_p, ctypes.POINTER(MBI), ctypes.c_size_t]
k.VirtualQueryEx.restype = ctypes.c_size_t


ap = argparse.ArgumentParser()
ap.add_argument("pid", type=int)
args = ap.parse_args()

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

counts = {name: 0 for name in VTABLES}
vt_to_name = {vt: name for name, vt in VTABLES.items()}

mbi = MBI()
addr = 0
while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
    pr = mbi.Protect & 0xff
    if (mbi.State == MEM_COMMIT and mbi.Type == MEM_PRIVATE
            and pr in (0x04, 0x40)):
        buf = (ctypes.c_ubyte * mbi.RegionSize)()
        sz = ctypes.c_size_t(0)
        if k.ReadProcessMemory(h, mbi.BaseAddress, buf, mbi.RegionSize, ctypes.byref(sz)):
            data = bytes(buf[:sz.value])
            end = (len(data) // 4) * 4
            for off in range(0, end, 4):
                v = struct.unpack_from("<I", data, off)[0]
                if v in vt_to_name:
                    counts[vt_to_name[v]] += 1
    addr = (mbi.BaseAddress or 0) + mbi.RegionSize
    if addr >= 0x80000000:
        break

ts = time.strftime("%H:%M:%S")
print(f"PID={args.pid} @ {ts}")
for name, vt in VTABLES.items():
    print(f"  {name:<22} 0x{vt:08x}  {counts[name]:>6}")