#!/usr/bin/env python3
"""basic-benchmark: background power/thermal sampler.

RAPL energy counters are root-only, so run with sudo:

    sudo python3 powerlog.py --tag before --interval 1

Writes results/power-<tag>.csv at --interval seconds until SIGINT/SIGTERM.
Columns: epoch,iso,pkg_w,core_w,cpu_temp_c,cpu_mhz,gpu_w,gpu_temp_c,gpu_busy,fan1,fan2
"""
import argparse
import csv
import glob
import os
import signal
import sys
import time
from datetime import datetime


def rd(path):
    try:
        with open(path) as fh:
            return fh.read().strip()
    except OSError:
        return None


def rdi(path):
    try:
        return int(rd(path))
    except (TypeError, ValueError):
        return None


def rapl(name):
    for d in sorted(glob.glob("/sys/class/powercap/intel-rapl:*/")):
        if rd(os.path.join(d, "name")) == name:
            energy = os.path.join(d, "energy_uj")
            if os.path.exists(energy):
                return energy, rdi(os.path.join(d, "max_energy_range_uj"))
    return None, None


def hwmon(name):
    for h in sorted(glob.glob("/sys/class/hwmon/hwmon*")):
        if rd(os.path.join(h, "name")) == name:
            return h
    return None


def first(paths):
    for p in paths:
        if p and os.path.exists(p):
            return p
    return None


def amdgpu_hwmons():
    out = []
    for h in sorted(glob.glob("/sys/class/hwmon/hwmon*")):
        if rd(os.path.join(h, "name")) == "amdgpu":
            out.append(h)
    return out


def pick_gpu_hwmon():
    """Prefer the discrete GPU.

    On hybrid AMD systems both the iGPU and the dGPU expose a hwmon named
    "amdgpu"; the dGPU is the one that reports power1_average. Fall back to the
    first amdgpu hwmon (and to GPU_DEVPATH, if set) so this stays stable.
    """
    hs = amdgpu_hwmons()
    want = os.environ.get("GPU_DEVPATH")
    if want:
        for h in hs:
            if os.path.realpath(os.path.join(h, "device")) == os.path.realpath(want):
                return h
    for h in hs:
        if os.path.exists(os.path.join(h, "power1_average")):
            return h
    return hs[0] if hs else None


def card_for_device(devpath):
    """Match a hwmon device to its /sys/class/drm/cardN node."""
    if not devpath:
        return None
    real = os.path.realpath(devpath)
    for c in sorted(glob.glob("/sys/class/drm/card[0-9]*")):
        name = os.path.basename(c)
        if not (name.startswith("card") and name[4:].isdigit()):
            continue
        dev = os.path.join(c, "device")
        if os.path.exists(dev) and os.path.realpath(dev) == real:
            return c
    return None


def cpu_temp():
    for name in ("k10temp", "coretemp", "zenpower"):
        h = hwmon(name)
        if not h:
            continue
        hit = first(sorted(glob.glob(os.path.join(h, "temp*_input"))))
        if hit:
            return hit
    return None


def avg_mhz():
    freq = [rdi(f) for f in glob.glob("/sys/devices/system/cpu/cpu*/cpufreq/scaling_cur_freq")]
    freq = [f for f in freq if f]
    return int(sum(freq) / len(freq) / 1000) if freq else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", required=True)
    ap.add_argument("--interval", type=float, default=1.0)
    ap.add_argument("--out", default=None)
    ap.add_argument("--count", type=int, default=0, help="0 = run until signalled")
    args = ap.parse_args()

    here = os.path.dirname(os.path.abspath(__file__))
    out = args.out or os.path.join(here, "results", f"power-{args.tag}.csv")
    os.makedirs(os.path.dirname(out), exist_ok=True)

    pkg_e, pkg_max = rapl("package-0")
    core_e, core_max = rapl("core")
    cpu_t = cpu_temp()
    gpu_h = pick_gpu_hwmon()
    gpu_p = os.path.join(gpu_h, "power1_average") if gpu_h else None
    gpu_t = os.path.join(gpu_h, "temp1_input") if gpu_h else None
    fan_h = hwmon("nct6798") or hwmon("nct6775")
    fan1 = os.path.join(fan_h, "fan1_input") if fan_h else None
    fan2 = os.path.join(fan_h, "fan2_input") if fan_h else None
    gpu_busy = None
    gpu_card = card_for_device(os.path.join(gpu_h, "device")) if gpu_h else None
    if gpu_card:
        gpu_busy = first([os.path.join(gpu_card, "device", "gpu_busy_percent")])

    if not pkg_e and not core_e:
        sys.stderr.write("powerlog: no RAPL counters found; run as root\n")

    state = {"stop": False}

    def on_signal(*_):
        state["stop"] = True

    signal.signal(signal.SIGINT, on_signal)
    signal.signal(signal.SIGTERM, on_signal)

    prev = {"pkg": None, "core": None}
    last = time.time()
    rows = 0

    with open(out, "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["epoch", "iso", "pkg_w", "core_w", "cpu_temp_c", "cpu_mhz",
                    "gpu_w", "gpu_temp_c", "gpu_busy", "fan1", "fan2"])
        fh.flush()

        def watts(key, path, mx):
            cur = rdi(path) if path else None
            if cur is None:
                return None
            old = prev[key]
            prev[key] = cur
            if old is None:
                return None
            delta = cur - old
            if delta < 0 and mx:
                delta += mx
            return round(delta / 1e6 / max(now - last, 1e-3), 2)

        while not state["stop"]:
            now = time.time()
            pkg_w = watts("pkg", pkg_e, pkg_max)
            core_w = watts("core", core_e, core_max)
            gpu_w = None
            gw = rdi(gpu_p) if gpu_p else None
            if gw is not None:
                gpu_w = round(gw / 1e6, 2)
            temp = rdi(cpu_t) if cpu_t else None
            gt = rdi(gpu_t) if gpu_t else None
            w.writerow([
                int(now),
                datetime.fromtimestamp(now).isoformat(timespec="seconds"),
                pkg_w, core_w,
                round(temp / 1000, 1) if temp is not None else None,
                avg_mhz(),
                gpu_w,
                round(gt / 1000, 1) if gt is not None else None,
                rdi(gpu_busy) if gpu_busy else None,
                rdi(fan1) if fan1 else None,
                rdi(fan2) if fan2 else None,
            ])
            last = now
            rows += 1
            fh.flush()
            if args.count and rows >= args.count:
                break
            time.sleep(max(0.05, args.interval - (time.time() - now)))

    sys.stderr.write(f"powerlog: wrote {rows} samples to {out}\n")


if __name__ == "__main__":
    main()