#!/usr/bin/env python3
"""Sample CM4 CPU / RAM / temperature while Layer 1 runs, and write HEADROOM.md.

The brief flags CM4 headroom as the #1 make-or-break risk: openWakeWord + Silero
VAD + the idle loop + audio streaming must fit simultaneously. Run this ON THE
ROBOT alongside the app for ~30 minutes:

    python scripts/headroom_sampler.py --minutes 30

It samples every second, prints a live summary, and writes the results table into
HEADROOM.md (between the RESULTS markers). Temperature comes from `vcgencmd
measure_temp` when present, else /sys/class/thermal, else "n/a".
"""

from __future__ import annotations

import argparse
import subprocess
import time
from pathlib import Path

try:
    import psutil
except Exception:  # noqa: BLE001
    psutil = None

RESULTS_START = "<!-- RESULTS:START -->"
RESULTS_END = "<!-- RESULTS:END -->"
HEADROOM_MD = Path(__file__).resolve().parent.parent / "HEADROOM.md"


def read_temp_c() -> float | None:
    try:
        out = subprocess.check_output(["vcgencmd", "measure_temp"], text=True, timeout=2)
        return float(out.strip().split("=")[1].split("'")[0])
    except Exception:  # noqa: BLE001
        pass
    for zone in Path("/sys/class/thermal").glob("thermal_zone*/temp"):
        try:
            return int(zone.read_text().strip()) / 1000.0
        except Exception:  # noqa: BLE001
            continue
    return None


def summarize(vals: list[float]) -> tuple[float, float, float]:
    if not vals:
        return (0.0, 0.0, 0.0)
    return (min(vals), sum(vals) / len(vals), max(vals))


def write_results(cpu, ram, temp, minutes, samples) -> None:
    if not HEADROOM_MD.exists():
        print(f"(no {HEADROOM_MD.name} to update; printing only)")
        return
    lo_c, avg_c, hi_c = summarize(cpu)
    lo_r, avg_r, hi_r = summarize(ram)
    temps = [t for t in temp if t is not None]
    lo_t, avg_t, hi_t = summarize(temps)
    tcell = "n/a" if not temps else f"{lo_t:.0f} / {avg_t:.0f} / {hi_t:.0f}"
    table = (
        f"{RESULTS_START}\n"
        f"_Measured over {minutes:.0f} min ({samples} samples), app running._\n\n"
        f"| metric | min | avg | max |\n"
        f"|--------|-----|-----|-----|\n"
        f"| CPU (%, all cores) | {lo_c:.0f} | {avg_c:.0f} | {hi_c:.0f} |\n"
        f"| RAM used (%) | {lo_r:.0f} | {avg_r:.0f} | {hi_r:.0f} |\n"
        f"| SoC temp (°C) | {tcell} |\n"
        f"{RESULTS_END}"
    )
    text = HEADROOM_MD.read_text()
    if RESULTS_START in text and RESULTS_END in text:
        pre = text.split(RESULTS_START)[0]
        post = text.split(RESULTS_END)[1]
        HEADROOM_MD.write_text(pre + table + post)
        print(f"Wrote results into {HEADROOM_MD.name}")
    else:
        print("(RESULTS markers not found; printing only)")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--minutes", type=float, default=30.0)
    ap.add_argument("--interval", type=float, default=1.0)
    args = ap.parse_args()
    if psutil is None:
        raise SystemExit("psutil not installed (pip install psutil)")

    cpu: list[float] = []
    ram: list[float] = []
    temp: list[float | None] = []
    end = time.monotonic() + args.minutes * 60
    psutil.cpu_percent()  # prime
    print(f"Sampling for {args.minutes:.0f} min (Ctrl-C to stop early)...")
    try:
        while time.monotonic() < end:
            c = psutil.cpu_percent(interval=args.interval)
            r = psutil.virtual_memory().percent
            t = read_temp_c()
            cpu.append(c)
            ram.append(r)
            temp.append(t)
            print(f"\rCPU {c:5.1f}%  RAM {r:5.1f}%  temp {('%.1f' % t) if t else 'n/a':>5}°C  "
                  f"n={len(cpu)}", end="", flush=True)
    except KeyboardInterrupt:
        print("\nstopped early")
    print()
    write_results(cpu, ram, temp, args.minutes, len(cpu))


if __name__ == "__main__":
    main()