#!/usr/bin/env python3
"""Read-only process/model accounting, for diagnostics (not deployment runtime)."""
import argparse
import collections
import json
from pathlib import Path
import re
import struct

p = argparse.ArgumentParser()
p.add_argument("--pid", type=int)
p.add_argument("--model", type=Path)
a = p.parse_args()
result = {}
if a.model:
    categories = collections.Counter()
    with a.model.open("rb") as f:
        header = f.read(256)
        count = struct.unpack_from("<I", header, 20)[0]
        offset = struct.unpack_from("<Q", header, 32)[0]
        f.seek(offset)
        for _ in range(count):
            entry = f.read(160)
            dtype, role = struct.unpack_from("<I", entry, 8)[0], struct.unpack_from("<I", entry, 16)[0]
            data_bytes = struct.unpack_from("<Q", entry, 56)[0]
            categories[f"dtype{dtype}_role{role}"] += data_bytes
    result["model_mib"] = {k: v / 1048576 for k, v in categories.items()}
    result["model_file_mib"] = a.model.stat().st_size / 1048576
if a.pid:
    categories = collections.defaultdict(collections.Counter)
    counts = collections.Counter()
    category = "unknown"
    for line in Path(f"/proc/{a.pid}/smaps").read_text().splitlines():
        if re.match(r"^[0-9a-f]+-[0-9a-f]+ ", line):
            parts = line.split(maxsplit=5)
            path = parts[5] if len(parts) == 6 else ""
            category = "model" if path.endswith(".l3r") else "drm" if "/dev/dri/" in path else (
                "anonymous" if not path or path.startswith("[") else "libraries_and_other_files")
            counts[category] += 1
        elif re.match(r"^(Size|Rss|Pss|Private_Dirty|Shared_Dirty|Private_Clean|Shared_Clean|Swap):", line):
            k, v, *_ = line.split()
            categories[category][k[:-1]] += int(v) / 1024
    result["smaps_mib"] = dict(categories)
    result["mapping_counts"] = dict(counts)
    result["status"] = {line.split(":", 1)[0]: line.split(":", 1)[1].strip()
                        for line in Path(f"/proc/{a.pid}/status").read_text().splitlines()
                        if line.startswith(("VmRSS:", "VmHWM:", "VmSwap:", "Threads:"))}
print(json.dumps(result, indent=2))