#!/usr/bin/env python3
"""Per-thread CPU of a process over an interval, grouped by thread name (Linux /proc).
Shows where a smoke run's CPU goes: UI/render thread, audio thread, OpenAL mixer,
compositor, GC... (smoke_run.sh step `threads:<s>`).

Usage: scripts/thread_cpu.py PID SECONDS [--top N]
"""
import argparse
import os
import time
from collections import defaultdict


def sample(pid):
    """{tid: (name, cpu_ticks)} for every live thread of pid."""
    out = {}
    base = f"/proc/{pid}/task"
    for tid in os.listdir(base):
        try:
            with open(f"{base}/{tid}/stat") as f:
                stat = f.read()
        except OSError:
            continue  # thread ended
        name = stat[stat.index("(") + 1: stat.rindex(")")]
        fields = stat[stat.rindex(")") + 2:].split()
        out[tid] = (name, int(fields[11]) + int(fields[12]))  # utime + stime
    return out


def per_name(before, after, seconds, hz):
    """[(name, percent of one core, thread count)] busiest first; threads that started
    during the interval count from zero."""
    total = defaultdict(float)
    count = defaultdict(int)
    for tid, (name, ticks) in after.items():
        start = before.get(tid, (name, 0))[1]
        total[name] += (ticks - start) * 100.0 / hz / seconds
        count[name] += 1
    return sorted(((n, total[n], count[n]) for n in total), key=lambda r: -r[1])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pid", type=int)
    ap.add_argument("seconds", type=float)
    ap.add_argument("--top", type=int, default=10)
    args = ap.parse_args()
    hz = os.sysconf("SC_CLK_TCK")
    a = sample(args.pid)
    time.sleep(args.seconds)
    b = sample(args.pid)
    rows = per_name(a, b, args.seconds, hz)
    print(f"threads: {sum(r[1] for r in rows):.1f}% total over {args.seconds:g}s")
    for name, pct, n in rows[: args.top]:
        if pct >= 0.05:
            print(f"  {pct:5.1f}%  {name}" + (f" ×{n}" if n > 1 else ""))


if __name__ == "__main__":
    main()