#!/usr/bin/env python3
# airmon (ncurses) - passive whole-chip activity monitor for the Apple M1 on Asahi.
#   CPU: per-core utilisation (/proc/stat) + per-cluster/core frequency (cpufreq)
#   GPU: utilisation (rail duty-cycle) + clock (pstate) + rail voltage + region temp
#   MEM: used / cache / swap
#   BAT: charge level + charge/discharge power + time remaining
# Nothing here dispatches GPU work - it reads the SMC (via the gpumon debugfs
# endpoint) and /proc + /sys metadata only. Needs root for the debugfs reads.
#
#   sudo ./airmon.py            (q or Esc to quit; auto-handles resize)
#   GN=40 GAP=0.02 sudo ./airmon.py
import curses, os, time, subprocess, sys

RAIL    = "/sys/kernel/debug/gpu_rail_mv"
GPUMON  = "/sys/kernel/debug/gpumon"
CLIENTS = "/sys/kernel/debug/dri/206400000.gpu/clients"
BATT    = "/sys/class/power_supply/macsmc-battery"
GN  = int(os.environ.get("GN", 16))      # GPU rail samples per refresh
GAP = float(os.environ.get("GAP", 0.06)) # gap between samples (sets the refresh window ~1s)

# rail mV -> GPU pstate MHz (measured rail voltages from the device-tree OPP table)
PSTATES = [(400,0),(596,396),(631,528),(681,720),(771,924),(862,1128),(934,1278)]
VMAX = 933  # 0.93 V = 1278 MHz, the top pstate -> full-scale for the rail meter

# --- driver: the gpumon.ko module creates the debugfs endpoints we read
# (gpu_rail_mv, gpumon). If it isn't loaded, every GPU/rail/temp read silently
# returns 0 and the monitor looks dead, so we check-and-load it on startup.
KO_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "gpumon.ko")

def _loaded(mod):
    return os.path.isdir(f"/sys/module/{mod}")

# gpumon_init() looks up this platform device and takes the SMC handle from its
# parent, so its presence - not merely the macsmc module being loaded - is the
# real precondition for the module loading at all.
HWMON_DEV = "/sys/bus/platform/devices/macsmc-hwmon"

def _endpoints_live():
    """Both debugfs files must be there. Checking only `gpumon` is not enough:
    if gpu_rail_mv is missing the monitor still starts, but every rail sample
    reads 0 and the GPU meters sit flat at 0% - which looks like a dead GPU
    rather than a broken driver."""
    return os.path.exists(GPUMON) and os.path.exists(RAIL)

def _load_module():
    """Try modprobe first, then the .ko beside this script.

    Installed from the smcmon-dkms package the module is built into
    /lib/modules/<kver>/updates and listed in modules.dep, so modprobe finds
    it and there is no .ko next to the script at all. Run from a source
    checkout it is the other way round."""
    r = subprocess.run(["modprobe", "gpumon"], capture_output=True, text=True)
    if r.returncode == 0 or _loaded("gpumon"):
        return None                              # loaded, or something raced us in
    if os.path.exists(KO_PATH):
        k = subprocess.run(["insmod", KO_PATH], capture_output=True, text=True)
        if k.returncode == 0 or _loaded("gpumon"):
            return None
        return (k.stderr or "").strip() or f"insmod exited {k.returncode}"
    return (r.stderr or "").strip() or f"modprobe exited {r.returncode}"

def ensure_driver():
    # debugfs reads AND insmod both need root; airmon is meant to run under sudo.
    # Note /sys/kernel/debug is mode 0700, so as non-root every os.path.exists()
    # below would return False regardless of what is actually there.
    if os.geteuid() != 0:
        sys.exit("airmon: must run as root — try 'sudo ./airmon.py' "
                 "(needs gpumon.ko + root debugfs reads)")
    if _loaded("gpumon") and _endpoints_live():
        return                                   # already up

    if not os.path.ismount("/sys/kernel/debug"):
        sys.exit("airmon: debugfs is not mounted — "
                 "try 'sudo mount -t debugfs none /sys/kernel/debug'")
    if not _loaded("macsmc"):                    # gpumon depends on macsmc
        subprocess.run(["modprobe", "macsmc"], check=False)
    if not os.path.exists(HWMON_DEV):
        sys.exit(f"airmon: {HWMON_DEV} missing — the SMC has not probed, so "
                 "gpumon.ko cannot bind. Is this an Apple Silicon machine "
                 "running the Asahi kernel?")

    if _loaded("gpumon") and not _endpoints_live():
        # Loaded but the endpoints are not both there. Reload rather than give
        # up: this is what a pre-fix gpumon leaves behind, since it used to
        # leak gpu_rail_mv on unload and the re-create then silently failed.
        subprocess.run(["rmmod", "gpumon"], capture_output=True, text=True)

    if not _loaded("gpumon"):
        err = _load_module()
        if err:
            # In a source checkout a kernel upgrade leaves the prebuilt .ko
            # unloadable; rebuild once against the running kernel and retry.
            # Under DKMS this cannot arise - dkms rebuilds on kernel install -
            # so only attempt it when there is a tree to build in.
            if os.path.exists(KO_PATH) and ("format" in err.lower()
                                            or "magic" in err.lower()):
                print(f"airmon: {err}\nairmon: rebuilding gpumon.ko for "
                      f"{os.uname().release}...", file=sys.stderr)
                b = subprocess.run(["make", "-C", os.path.dirname(KO_PATH)],
                                   capture_output=True, text=True)
                if b.returncode != 0:
                    sys.exit(f"airmon: rebuild failed:\n{b.stderr.strip()}")
                err = _load_module()
            if err:
                sys.exit(f"airmon: could not load the gpumon module: {err}\n"
                         "airmon: install smcmon-dkms, or run 'make' in a "
                         "source checkout")

    if not _endpoints_live():
        missing = [p for p in (GPUMON, RAIL) if not os.path.exists(p)]
        sys.exit("airmon: gpumon loaded but missing: " + ", ".join(missing))

def read_text(path):
    try:
        with open(path) as f: return f.read()
    except OSError: return ""

def read_int(path):
    s = read_text(path).strip()
    try: return int(s)
    except ValueError: return 0

def mv2mhz(v):
    if v < 450: return 0
    return min(PSTATES, key=lambda p: abs(p[0]-v))[1]

def read_stat():
    d = {}
    for ln in read_text("/proc/stat").splitlines():
        if ln.startswith("cpu") and ln[3].isdigit():
            f = ln.split(); tot = sum(int(x) for x in f[1:]); idle = int(f[4])+int(f[5])
            d[f[0]] = (tot, idle)
    return d

def cpu_util(prev, cur):
    u = {}
    for k, (t, i) in cur.items():
        pt, pi = prev.get(k, (t, i)); dt = t-pt; di = i-pi
        u[k] = (dt-di)*100//dt if dt > 0 else 0
    return u

def cl_mhz(p):  return read_int(f"/sys/devices/system/cpu/cpufreq/policy{p}/scaling_cur_freq")//1000
def core_mhz(c):return read_int(f"/sys/devices/system/cpu/cpu{c}/cpufreq/scaling_cur_freq")//1000

def gpu_temp():
    for ln in read_text(GPUMON).splitlines():
        if "TEMP Ts1z" in ln:
            try: return float(ln.split()[2])
            except (IndexError, ValueError): return None
    return None

def gpu_ctx():
    lines = read_text(CLIENTS).splitlines()[1:]
    c = {}
    for ln in lines:
        f = ln.split()
        if f: c[f[0]] = c.get(f[0], 0) + 1
    return " ".join(f"{k}({v})" for k, v in c.items())

BSTAT = {"Charging": "chg", "Discharging": "bat", "Full": "full",
         "Not charging": "ac", "Unknown": "?"}

def batt():
    """(charge %, short status, watts, seconds left). macsmc reports power_now in
    µW, signed: negative while discharging. time_to_* is 0 when not applicable."""
    st  = read_text(f"{BATT}/status").strip()
    eta = read_int(f"{BATT}/time_to_full_now" if st == "Charging"
                   else f"{BATT}/time_to_empty_now")
    return (read_int(f"{BATT}/capacity"), BSTAT.get(st, st[:4].lower() or "?"),
            abs(read_int(f"{BATT}/power_now")) / 1e6, eta)

def hm(s):
    return f"{s//3600}h{s%3600//60:02d}" if s > 0 else ""

def mem():
    m = {}
    for ln in read_text("/proc/meminfo").splitlines():
        k, _, v = ln.partition(":"); m[k] = int(v.split()[0]) if v.split() else 0
    tot = m.get("MemTotal", 1); used = tot - m.get("MemAvailable", 0)
    g = 1048576.0
    return (used*100//tot, used/g, tot/g, m.get("Cached", 0)/g,
            (m.get("SwapTotal", 0)-m.get("SwapFree", 0))/1024.0)

# color by load: green / yellow / red
def cload(p): return 1 if p < 50 else (2 if p < 80 else 3)
def cbatt(p): return 3 if p < 20 else (2 if p < 50 else 1)   # inverted: empty is bad
def ctemp(t): return 1 if t < 60 else (2 if t < 80 else 3)

def meter_cells(pct, w):
    """Bar fill: whole solid blocks (█) only — renders in any CP437/console font.
    Returns (filled_str, n_blank_cells)."""
    pct = max(0, min(100, pct)); n = pct*w//100
    return "█"*n, w-n

def draw_meter(scr, y, w, label, pct, val, cpair):
    rows, cols = scr.getmaxyx()
    if y >= rows: return
    filled, nblank = meter_cells(pct, w)
    try:
        scr.addstr(y, 2, f"{label:<4} [")
        scr.addstr(filled, curses.color_pair(cpair))
        scr.addstr(" "*nblank)
        scr.addstr(f"] {val}"[:max(0, cols-1-scr.getyx()[1])])
    except curses.error:
        pass

def line(scr, y, x, s, attr=0):
    rows, cols = scr.getmaxyx()
    if 0 <= y < rows:
        try: scr.addstr(y, x, s[:cols-1-x], attr)
        except curses.error: pass

def main(scr):
    curses.curs_set(0); scr.nodelay(True)
    curses.start_color(); curses.use_default_colors()
    curses.init_pair(1, curses.COLOR_GREEN, -1)
    curses.init_pair(2, curses.COLOR_YELLOW, -1)
    curses.init_pair(3, curses.COLOR_RED, -1)
    curses.init_pair(4, curses.COLOR_CYAN, -1)

    prev = read_stat()
    while True:
        busy = mx = 0
        for _ in range(GN):
            mv = read_int(RAIL)
            # Track the rail whether or not the GPU is gated. Only counting mx
            # above the ungated threshold left it at 0 through any idle window,
            # so the Vdd readout printed "0 mV" while the rail really sat at
            # ~8 mV - indistinguishable from the driver being dead.
            mx = max(mx, mv)
            if mv > 450: busy += 1
            ch = scr.getch()
            if ch in (ord("q"), ord("Q"), 27): return
            time.sleep(GAP)
        cur = read_stat(); util = cpu_util(prev, cur); prev = cur
        gutil = busy*100//GN; gmhz = mv2mhz(mx)
        rpct  = min(100, mx*100//VMAX)
        t     = gpu_temp(); tnum = int(t) if t is not None else 0
        mp, mu, mt, mc, ms = mem()
        ctx = gpu_ctx()
        bat = batt() if os.path.isdir(BATT) else None

        rows, cols = scr.getmaxyx()
        w = max(8, cols - 24)
        sep = "─" * (cols - 4)
        scr.erase()
        line(scr, 0, 2, f"M1 · Asahi · airmon          {time.strftime('%H:%M:%S')}     (passive · no GPU work · q to quit)", curses.color_pair(4))
        line(scr, 1, 2, sep)
        line(scr, 2, 2, f"P-cluster Firestorm  {cl_mhz(4):4d} MHz")
        y = 3
        for c in (4, 5, 6, 7):
            u = util.get(f"cpu{c}", 0)
            draw_meter(scr, y, w, f"c{c}", u, f"{u:3d}% {core_mhz(c):4d} MHz", cload(u)); y += 1
        line(scr, y, 2, f"E-cluster Icestorm   {cl_mhz(0):4d} MHz"); y += 1
        for c in (0, 1, 2, 3):
            u = util.get(f"cpu{c}", 0)
            draw_meter(scr, y, w, f"c{c}", u, f"{u:3d}% {core_mhz(c):4d} MHz", cload(u)); y += 1
        line(scr, y, 2, sep); y += 1
        draw_meter(scr, y, w, "GPU", gutil, f"{gutil:3d}% {gmhz:4d} MHz", cload(gutil)); y += 1
        draw_meter(scr, y, w, "Vdd", rpct, f"{mx:4d} mV", cload(rpct)); y += 1
        draw_meter(scr, y, w, "T°", tnum, f"{tnum:3d}°C" if t is not None else "  ?°C", ctemp(tnum)); y += 1
        line(scr, y, 2, f"{'ACTIVE' if gutil else 'idle':<6}  contexts: {ctx or 'none'}"); y += 1
        line(scr, y, 2, sep); y += 1
        draw_meter(scr, y, w, "MEM", mp, f"{mp:3d}% {mu:.1f} GiB", cload(mp)); y += 1
        line(scr, y, 2, f"      cache {mc:.1f} GiB   swap {ms:.0f} MiB"); y += 1
        if bat:
            bp, bst, bw, beta = bat
            det = f"{bp:3d}% {bst}"
            if bw >= 0.05: det += f" {bw:4.1f} W"
            if beta:       det += f" {hm(beta)}"
            draw_meter(scr, y, w, "BAT", bp, det, cbatt(bp)); y += 1
        scr.noutrefresh(); curses.doupdate()

if __name__ == "__main__":
    ensure_driver()
    try: curses.wrapper(main)
    except KeyboardInterrupt: pass
