"""byte_accounting.py <pid>
Walk all committed memory in a target process and categorize bytes by:
  - VAD region type (private vs mapped vs image)
  - Protection (RW vs RX vs RWX)
  - Size bucket
  - Known-class signature scan (vtable bytes within the region)
Output: per-category totals so we can see where the working set lives."""
import ctypes, ctypes.wintypes as wt, sys, struct

PROCESS_VM_READ = 0x10
PROCESS_QUERY_INFORMATION = 0x400
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, wt.LPCVOID, ctypes.c_void_p, ctypes.c_size_t]
k.VirtualQueryEx.restype = ctypes.c_size_t

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),
    ]

MEM_COMMIT = 0x1000
MEM_PRIVATE = 0x20000
MEM_MAPPED  = 0x40000
MEM_IMAGE   = 0x1000000

# Known class vtables from instr.cpp
KNOWN_VTABLES = {
    0x007C78EC: "CPhysicsObj",
    0x0079A67C: "RenderSurface",
    0x0079C198: "RenderTexture",
    0x00801A94: "RenderSurfaceD3D",
    0x00801A18: "RenderTextureD3D",
    0x007CA4DC: "CSurface(GR)",
    0x007CAB04: "ImgTex(GR)",
    0x007CA418: "CGfxObj",
    0x007ED3B0: "GXTri3Mesh",
    0x007E4F70: "ACCWeenieObject",
    0x007E4ED8: "CWeenieObject",
}

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])

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)

# Totals by region type
total_committed = 0
by_type = {"private_rw": 0, "private_rx": 0, "private_rwx": 0,
           "mapped": 0, "image": 0, "other": 0}
# Bytes attributed to each known class (rough estimate: vtable_count × likely class size)
class_size_estimate = {
    "CPhysicsObj": 376,        # +0x178 from constructor allocation
    "RenderSurface": 288,
    "RenderTexture": 152,
    "RenderSurfaceD3D": 304,
    "RenderTextureD3D": 176,
    "CSurface(GR)": 144,
    "ImgTex(GR)": 136,
    "CGfxObj": 200,            # estimate
    "GXTri3Mesh": 1000,        # estimate; large mesh class
    "ACCWeenieObject": 336,
    "CWeenieObject": 200,      # estimate
}
class_instance_count = {name: 0 for name in KNOWN_VTABLES.values()}

# Histogram of private-RW region sizes
size_buckets = [
    (0, 4096, "<4K"),
    (4096, 65536, "4K-64K"),
    (65536, 262144, "64K-256K"),
    (262144, 1048576, "256K-1M"),
    (1048576, 4194304, "1M-4M"),
    (4194304, 16777216, "4M-16M"),
    (16777216, 67108864, "16M-64M"),
    (67108864, 1<<31, "64M+"),
]
bucket_count = [0] * len(size_buckets)
bucket_bytes = [0] * len(size_buckets)

def classify_size(n):
    for i, (lo, hi, _) in enumerate(size_buckets):
        if lo <= n < hi: return i
    return len(size_buckets) - 1

mbi = MBI()
addr = 0
while k.VirtualQueryEx(h, addr, ctypes.byref(mbi), ctypes.sizeof(mbi)):
    region_base = mbi.BaseAddress or 0
    region_size = mbi.RegionSize
    if mbi.State == MEM_COMMIT:
        total_committed += region_size
        prot = mbi.Protect & 0xFF
        # Classify region type
        if mbi.Type == MEM_IMAGE:
            by_type["image"] += region_size
        elif mbi.Type == MEM_MAPPED:
            by_type["mapped"] += region_size
        elif mbi.Type == MEM_PRIVATE:
            if prot == 0x40:    by_type["private_rwx"] += region_size
            elif prot == 0x20:  by_type["private_rx"] += region_size
            elif prot == 0x04:  by_type["private_rw"] += region_size
            else:               by_type["other"] += region_size
        else:
            by_type["other"] += region_size

        # For private RW/RWX: bucket by size + scan for known vtables
        if mbi.Type == MEM_PRIVATE and prot in (0x04, 0x40):
            bi = classify_size(region_size)
            bucket_count[bi] += 1
            bucket_bytes[bi] += region_size
            # Scan for known vtable bytes (skip huge regions to bound time)
            if region_size <= 64 * 1024 * 1024:
                try:
                    data = rd(h, region_base, region_size)
                    if data:
                        for vt, name in KNOWN_VTABLES.items():
                            vt_bytes = struct.pack('<I', vt)
                            count = 0
                            off = 0
                            while True:
                                off = data.find(vt_bytes, off)
                                if off < 0: break
                                if (off & 3) == 0: count += 1
                                off += 4
                            class_instance_count[name] += count
                except Exception:
                    pass

    next_addr = region_base + region_size
    if next_addr <= addr: break
    addr = next_addr
    if addr >= 0x80000000: break

k.CloseHandle(h)

def mb(n): return f"{n/(1024*1024):,.1f}"

print(f"=== pid {pid} byte accounting ===")
print(f"Total committed: {mb(total_committed)} MB")
print()
print("By region type:")
for label, n in by_type.items():
    pct = (n*100/total_committed) if total_committed else 0
    print(f"  {label:14s} {mb(n):>9} MB  ({pct:5.1f}%)")
print()
print("Private RW/RWX region size distribution:")
print(f"  {'bucket':<12} {'count':>6} {'total MB':>10}")
for i, (lo, hi, label) in enumerate(size_buckets):
    if bucket_count[i] == 0: continue
    print(f"  {label:<12} {bucket_count[i]:>6} {mb(bucket_bytes[i]):>10}")
print()
print("Known-class vtable counts (and estimated bytes):")
print(f"  {'class':<22} {'count':>6} {'est bytes':>12}  {'est MB':>8}")
total_class_bytes = 0
for name in sorted(class_instance_count, key=lambda x: -class_instance_count[x]):
    n = class_instance_count[name]
    if n == 0: continue
    sz = class_size_estimate.get(name, 200)
    bytes_total = n * sz
    total_class_bytes += bytes_total
    print(f"  {name:<22} {n:>6} {bytes_total:>12,} {mb(bytes_total):>8}")
print()
print(f"  Identified-class total: ~{mb(total_class_bytes)} MB")
print(f"  Of private RW:          ~{mb(by_type['private_rw'])} MB")
unidentified = by_type['private_rw'] - total_class_bytes
print(f"  UNIDENTIFIED in priv RW: ~{mb(unidentified)} MB  ← what we don't account for")