#!/usr/bin/env python3
"""Sample npu-smi info and emit nvidia-smi-compatible CSV rows.

Usage: npu_smi_sampler.py <interval_seconds>
Emits a header row first, then one row per chip every interval seconds.
Output columns match nvidia-smi --query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu --format=csv
"""
import sys, re, datetime, time

def sample():
    import subprocess
    try:
        out = subprocess.run(["npu-smi", "info"], capture_output=True, text=True, timeout=10).stdout
    except Exception:
        return []
    ts = datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
    rows = []
    idx = None
    for line in out.splitlines():
        m = re.match(r"\|\s*(\d+)\s+Ascend", line)
        if m:
            idx = m.group(1)
            continue
        if idx is None:
            continue
        hbm = re.search(r"(\d+)\s*/\s*(\d+)", line)
        util = re.search(r"(\d+)\s*%", line)
        if hbm:
            u, t = hbm.group(1), hbm.group(2)
            r = util.group(1) if util else "0"
            rows.append(f"{ts}, {idx}, {u} MiB, {t} MiB, {r} %")
            idx = None
    return rows

def main():
    interval = float(sys.argv[1]) if len(sys.argv) > 1 else 1.0
    print("timestamp, index, memory.used [MiB], memory.total [MiB], utilization.gpu [%]")
    sys.stdout.flush()
    while True:
        for row in sample():
            print(row)
        sys.stdout.flush()
        time.sleep(interval)

if __name__ == "__main__":
    main()