"""count_gr_subclasses_live.py <pid>

Count live instances of each GraphicsResource subclass in a running
process by scanning RW heap for vtable pointers. Used to measure
whether v5 PurgeResource patch actually drains the leaked instances
over time.
"""
import argparse, ctypes, ctypes.wintypes as wt, struct, sys, time


VTABLES = {
    "RenderSurface":          0x0079a67c,
    "RenderTexture":          0x0079c198,
    "CSurface":               0x007ca4dc,
    "ImgTex":                 0x007cab04,
    "RenderVertexBufferD3D":  0x007e6520,
    "RenderTextureD3D":       0x00801a18,
    "RenderSurfaceD3D":       0x00801a94,
    "RenderIndexStreamD3D":   0x00801b64,
}


PROCESS_VM_READ           = 0x0010
PROCESS_QUERY_INFORMATION = 0x0400
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


def main():
    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_set = {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)):  # RW or RWX
            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_set:
                        counts[vt_set[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}")
    print(f"{'class':<25} {'vtable':<12} {'count':>6}")
    for name, vt in VTABLES.items():
        marker = " <- LEAKING (v5 patches this)" if name in ("RenderSurface", "RenderTexture") else ""
        print(f"{name:<25} 0x{vt:08x}   {counts[name]:>6}{marker}")
    total = sum(counts.values())
    leakers = counts["RenderSurface"] + counts["RenderTexture"]
    print(f"{'total':<25} {'':<12} {total:>6}")
    print(f"{'leakers (RS+RT)':<25} {'':<12} {leakers:>6}")


if __name__ == "__main__":
    main()