#!/usr/bin/env python3
"""Measure peak per-process RSS for a build process tree under the L17 limits."""
import os
import subprocess
import sys
import time
from pathlib import Path


def descendants(root):
    pending = [root]
    seen = set()
    while pending:
        pid = pending.pop()
        if pid in seen:
            continue
        seen.add(pid)
        try:
            children = Path(f"/proc/{pid}/task/{pid}/children").read_text().split()
            pending.extend(int(child) for child in children)
        except (FileNotFoundError, PermissionError, ProcessLookupError):
            pass
    return seen


def status(pid):
    try:
        fields = {}
        for line in Path(f"/proc/{pid}/status").read_text().splitlines():
            if line.startswith(("Name:", "VmRSS:", "VmSize:")):
                key, *value = line.split()
                fields[key[:-1]] = value[0] if value else ""
        return (int(fields.get("VmRSS", "0")), int(fields.get("VmSize", "0")),
                fields.get("Name", ""))
    except (FileNotFoundError, PermissionError, ProcessLookupError, ValueError):
        return None


if len(sys.argv) < 2:
    raise SystemExit("usage: MeasureBuild.py command [arguments ...]")
if not os.sched_getaffinity(0).issubset(set(range(12))):
    raise SystemExit("run under taskset -c 0-11")

started = time.monotonic()
process = subprocess.Popen(sys.argv[1:], start_new_session=True)
peak_rss = (0, 0, "")
peak_vms = (0, 0, "")
while process.poll() is None:
    for pid in descendants(process.pid):
        measured = status(pid)
        if measured is None:
            continue
        rss, vms, name = measured
        if rss > peak_rss[0]:
            peak_rss = (rss, pid, name)
        if vms > peak_vms[0]:
            peak_vms = (vms, pid, name)
    time.sleep(0.1)

exit_code = process.wait()
elapsed = time.monotonic() - started
print(f"[BuildTreeResources] exit_code={exit_code} wall_s={elapsed:.3f} "
      f"peak_process_rss_kib={peak_rss[0]} peak_rss_pid={peak_rss[1]} "
      f"peak_rss_name={peak_rss[2]} peak_process_vms_kib={peak_vms[0]} "
      f"peak_vms_pid={peak_vms[1]} peak_vms_name={peak_vms[2]} "
      "monitor_interval_s=0.1", flush=True)
raise SystemExit(exit_code)