#!/usr/bin/env python3
"""The CPU energy a browser's process tree has spent so far, for the end-to-end scenario
`binding` (WAI browser-sink §6): its figure beside the user agent's powerEfficient.

    python3 -I energy.py <run-id>

Prints one JSON line: the processes counted, and the sum over them of the kernel's
per-process energy counter (`proc_pid_rusage`, RUSAGE_INFO_V6: `ri_energy_nj`), plus the
energy services were billed for work done for them (`ri_billed_energy`), less what they
did for others (`ri_serviced_energy`), all in nanojoules. The tree is every process whose
command line carries `moq-web-run=<run-id>` (the browser the test launched) and every
descendant of one. The counter covers CPU energy only: GPU and media-engine energy are not
in it, so a figure built from it covers part of a decode (`Partial`, fabrics `cpu`). A
counter is a reading, not an estimate (energy-measurement §2: `OnChipCounter`).

Exits 2 where the counter does not exist (a host other than macOS); the scenario then
reports the step as skipped. Standard library only.
"""
import ctypes
import json
import subprocess
import sys

RUSAGE_INFO_V6 = 6
SIZE = 392 + 9 * 8  # struct rusage_info_v6
OFF_ENERGY = 336     # ri_energy_nj
OFF_BILLED = 264     # ri_billed_energy
OFF_SERVICED = 272   # ri_serviced_energy


def tree(run_id):
    ps = subprocess.run(["ps", "-A", "-o", "pid=,ppid=,command="], capture_output=True, text=True, check=True).stdout
    rows = []
    for line in ps.splitlines():
        parts = line.strip().split(None, 2)
        if len(parts) >= 2:
            rows.append((int(parts[0]), int(parts[1]), parts[2] if len(parts) > 2 else ""))
    roots = {p for p, _, c in rows if "moq-web-run=" + run_id in c}
    found = set(roots)
    grew = True
    while grew:
        grew = False
        for p, pp, _ in rows:
            if pp in found and p not in found:
                found.add(p)
                grew = True
    return sorted(found)


def main():
    if sys.platform != "darwin" or len(sys.argv) != 2:
        print(json.dumps({"error": "no per-process energy counter on this host"}))
        return 2
    libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib", use_errno=True)
    buf = ctypes.create_string_buffer(SIZE)
    total, counted = 0, 0
    for pid in tree(sys.argv[1]):
        if libc.proc_pid_rusage(ctypes.c_int(pid), ctypes.c_int(RUSAGE_INFO_V6), buf) != 0:
            continue  # gone, or not ours
        raw = buf.raw
        get = lambda off: int.from_bytes(raw[off:off + 8], "little")
        total += get(OFF_ENERGY) + get(OFF_BILLED) - get(OFF_SERVICED)
        counted += 1
    print(json.dumps({"processes": counted, "nj": total}))
    return 0 if counted else 2


if __name__ == "__main__":
    sys.exit(main())