#!/usr/bin/env python3
"""
hw-blackbox: throttle/power RECURRENCE recorder.

Confirms the PSU/cable fix held; catches PROCHOT / freq-collapse / RAPL / PSI
recurrence with a pre-event window. It is NOT a hard-hang blackbox -- hard-hang
forensics is pstore/ramoops + kdump/lockup-panic, armed separately. This
daemon does NOT capture silent lockups.
"""
import glob
import json
import os
import shutil
import subprocess
import sys
import time
from collections import deque

RING_DIR = "/run/hw-blackbox"
RING_FILE = os.path.join(RING_DIR, "ring.jsonl")
RING_N = 2000
LOG_DIR = "/var/log/hw-blackbox"
LOG_CAP_BYTES = 200 * 1024 * 1024
CALM_INTERVAL = 10
ANOMALY_INTERVAL = 1
CALM_HOLD = 60
THROTTLE_DIR_GLOB = "/sys/devices/system/cpu/cpu[0-9]*/thermal_throttle"
FREQ_GLOB = "/sys/devices/system/cpu/cpu[0-9]*/cpufreq/scaling_cur_freq"
MSR_0x64F = 0x64F
MSR_0x19C = 0x19C
LOW_FREQ_KHZ = 800_000
PSI_ANOMALY = 20.0

HZ = os.sysconf("SC_CLK_TCK") or 100


def read_int(path):
    try:
        with open(path) as f:
            return int(f.read().strip())
    except (OSError, ValueError):
        return None


def read_freqs():
    freqs = []
    for p in glob.glob(FREQ_GLOB):
        v = read_int(p)
        if v is not None:
            freqs.append(v)
    if not freqs:
        return {"min": None, "max": None, "avg": None}
    return {"min": min(freqs), "max": max(freqs), "avg": sum(freqs) // len(freqs)}


def read_temp_c():
    for name_path in glob.glob("/sys/class/hwmon/hwmon*/name"):
        try:
            with open(name_path) as f:
                name = f.read().strip()
        except OSError:
            continue
        if name != "coretemp":
            continue
        hwdir = os.path.dirname(name_path)
        for label_path in glob.glob(os.path.join(hwdir, "temp*_label")):
            try:
                with open(label_path) as f:
                    label = f.read().strip()
            except OSError:
                continue
            if label.startswith("Package id"):
                v = read_int(label_path.replace("_label", "_input"))
                return None if v is None else v / 1000.0
    return None


def rdmsr(cpu, addr):
    if not shutil.which("rdmsr"):
        return None
    try:
        out = subprocess.run(
            ["rdmsr", "-p", str(cpu), "-d", "0x{:x}".format(addr)],
            capture_output=True, text=True, timeout=2,
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    if out.returncode != 0:
        return None
    try:
        return int(out.stdout.strip())
    except ValueError:
        return None


def decode_m64f(raw):
    if raw is None:
        return {"raw": None, "prochot": None, "thermal": None, "pl1": None}
    return {
        "raw": raw,
        "prochot": bool(raw & (1 << 0)),
        "thermal": bool(raw & (1 << 1)),
        "pl1": bool(raw & (1 << 12)),
    }


def decode_m19c(raw):
    if raw is None:
        return {"raw": None, "prochot": None, "power": None, "current": None}
    return {
        "raw": raw,
        "prochot": bool(raw & (1 << 2)),
        "power": bool(raw & (1 << 10)),
        "current": (raw >> 16) & 0x7F,
    }


def read_throttle_counts():
    pkg_total = 0
    core_total = 0
    found = False
    for d in glob.glob(THROTTLE_DIR_GLOB):
        p = read_int(os.path.join(d, "package_throttle_count"))
        c = read_int(os.path.join(d, "core_throttle_count"))
        if p is not None:
            pkg_total += p
            found = True
        if c is not None:
            core_total += c
            found = True
    if not found:
        return None, None
    return pkg_total, core_total


def read_psi_field(path):
    try:
        with open(path) as f:
            line = f.readline()
    except OSError:
        return None
    for tok in line.split():
        if tok.startswith("avg10="):
            try:
                return float(tok.split("=", 1)[1])
            except ValueError:
                return None
    return None


def read_psi():
    return {
        "cpu": read_psi_field("/proc/pressure/cpu"),
        "io": read_psi_field("/proc/pressure/io"),
        "mem": read_psi_field("/proc/pressure/memory"),
    }


def read_load1():
    try:
        with open("/proc/loadavg") as f:
            return float(f.read().split()[0])
    except (OSError, ValueError, IndexError):
        return None


def read_ac():
    # type=Mains is the wall-power indicator; USB-C ucsi source-caps report
    # per-port PD state, not "on AC", so restrict to Mains supplies only.
    for d in glob.glob("/sys/class/power_supply/*"):
        try:
            with open(d + "/type") as f:
                if f.read().strip() != "Mains":
                    continue
        except OSError:
            continue
        v = read_int(d + "/online")
        if v is not None:
            return v
    return None


def read_bat():
    status = None
    current_now = None
    for d in glob.glob("/sys/class/power_supply/BAT*"):
        try:
            with open(os.path.join(d, "status")) as f:
                status = f.read().strip()
        except OSError:
            pass
        current_now = read_int(os.path.join(d, "current_now"))
        break
    return {"status": status, "current_now": current_now}


def read_proc_stat_fields(pid):
    try:
        with open(f"/proc/{pid}/stat", "rb") as f:
            raw = f.read().decode(errors="replace")
    except OSError:
        return None
    rp = raw.rfind(")")
    if rp == -1:
        return None
    lp = raw.find("(")
    comm = raw[lp + 1:rp]
    rest = raw[rp + 2:].split()
    try:
        utime = int(rest[11])
        stime = int(rest[12])
    except (IndexError, ValueError):
        return None
    return {"comm": comm, "utime": utime, "stime": stime}


def sample_top(prev_cpu):
    now = time.time()
    cur_cpu = {}
    rows = []
    for d in glob.glob("/proc/[0-9]*"):
        pid = os.path.basename(d)
        st = read_proc_stat_fields(pid)
        if st is None:
            continue
        cur_cpu[pid] = (st["utime"] + st["stime"], now)
        prev = prev_cpu.get(pid)
        if prev is None:
            continue
        prev_ticks, prev_ts = prev
        dt = now - prev_ts
        if dt <= 0:
            continue
        dticks = cur_cpu[pid][0] - prev_ticks
        if dticks < 0:
            continue
        pcpu = (dticks / HZ) / dt * 100.0
        rows.append({"pid": int(pid), "comm": st["comm"], "pcpu": round(pcpu, 1)})
    rows.sort(key=lambda r: r["pcpu"], reverse=True)
    return rows[:3], cur_cpu


def take_sample(prev_cpu, prev_throttle):
    t = time.time()
    freqs = read_freqs()
    temp_c = read_temp_c()
    m64f_raw = rdmsr(0, MSR_0x64F)
    m19c_raw = rdmsr(0, MSR_0x19C)
    pkg_c, core_c = read_throttle_counts()
    pkg_d = None if (pkg_c is None or prev_throttle[0] is None) else pkg_c - prev_throttle[0]
    core_d = None if (core_c is None or prev_throttle[1] is None) else core_c - prev_throttle[1]
    psi = read_psi()
    load1 = read_load1()
    ac = read_ac()
    bat = read_bat()
    top, cur_cpu = sample_top(prev_cpu)
    sample = {
        "t": int(t),
        "f": freqs,
        "temp_c": temp_c,
        "m64f": decode_m64f(m64f_raw),
        "m19c": decode_m19c(m19c_raw),
        "thr": {"pkg": pkg_c, "core": core_c, "pkg_d": pkg_d, "core_d": core_d},
        "psi": psi,
        "load1": load1,
        "ac": ac,
        "bat": bat,
        "top": top,
    }
    return sample, cur_cpu, (pkg_c, core_c)


def is_anomaly_rising(sample, ncores):
    thr = sample["thr"]
    if thr["pkg_d"] is not None and thr["pkg_d"] > 0:
        return True
    if thr["core_d"] is not None and thr["core_d"] > 0:
        return True
    m64f = sample["m64f"]
    if m64f["prochot"] is True or m64f["thermal"] is True:
        return True
    fmin = sample["f"]["min"]
    load1 = sample["load1"]
    if fmin is not None and load1 is not None and fmin < LOW_FREQ_KHZ and load1 > (ncores * 0.5):
        return True
    psi_cpu = sample["psi"]["cpu"]
    if psi_cpu is not None and psi_cpu > PSI_ANOMALY:
        return True
    return False


def enforce_log_cap(log_dir, cap_bytes):
    try:
        files = glob.glob(os.path.join(log_dir, "event-*.jsonl"))
    except OSError:
        return
    entries = []
    total = 0
    for f in files:
        try:
            sz = os.path.getsize(f)
        except OSError:
            continue
        entries.append((os.path.getmtime(f), sz, f))
        total += sz
    entries.sort()
    i = 0
    while total > cap_bytes and i < len(entries):
        _, sz, f = entries[i]
        try:
            os.remove(f)
            total -= sz
        except OSError:
            pass
        i += 1


class Ring:
    def __init__(self, maxlen):
        self.buf = deque(maxlen=maxlen)

    def append(self, sample):
        self.buf.append(sample)

    def write_tmpfs(self, path):
        tmp = path + ".tmp"
        try:
            os.makedirs(os.path.dirname(path), exist_ok=True)
            with open(tmp, "w") as f:
                for s in self.buf:
                    f.write(json.dumps(s, separators=(",", ":")) + "\n")
            os.replace(tmp, path)
        except OSError:
            pass

    def dump_lines(self):
        return [json.dumps(s, separators=(",", ":")) for s in self.buf]


def run_daemon():
    ncores = os.cpu_count() or 1
    ring = Ring(RING_N)
    prev_cpu = {}
    prev_throttle = (None, None)
    event_fh = None
    calm_since = None
    interval = CALM_INTERVAL

    while True:
        sample, prev_cpu, prev_throttle = take_sample(prev_cpu, prev_throttle)
        ring.append(sample)
        ring.write_tmpfs(RING_FILE)

        anomaly = is_anomaly_rising(sample, ncores)

        if anomaly:
            calm_since = None
            if event_fh is None:
                os.makedirs(LOG_DIR, exist_ok=True)
                enforce_log_cap(LOG_DIR, LOG_CAP_BYTES)
                event_path = os.path.join(LOG_DIR, f"event-{sample['t']}.jsonl")
                event_fh = open(event_path, "a")
                for line in ring.dump_lines():
                    event_fh.write(line + "\n")
                event_fh.flush()
            else:
                event_fh.write(json.dumps(sample, separators=(",", ":")) + "\n")
                event_fh.flush()
            interval = ANOMALY_INTERVAL
        else:
            if event_fh is not None:
                event_fh.write(json.dumps(sample, separators=(",", ":")) + "\n")
                event_fh.flush()
                if calm_since is None:
                    calm_since = time.time()
                elif time.time() - calm_since >= CALM_HOLD:
                    event_fh.close()
                    event_fh = None
                    calm_since = None
                    interval = CALM_INTERVAL
                else:
                    interval = ANOMALY_INTERVAL
            else:
                interval = CALM_INTERVAL

        if event_fh is not None:
            enforce_log_cap(LOG_DIR, LOG_CAP_BYTES)

        time.sleep(interval)


def self_test():
    ring = Ring(5)
    for i in range(7):
        ring.append({"t": i})
    assert len(ring.buf) == 5, "ring did not cap at maxlen"
    assert ring.buf[0]["t"] == 2, "ring did not drop oldest"

    synthetic = {
        "t": int(time.time()),
        "f": {"min": 3200000, "max": 4500000, "avg": 3900000},
        "temp_c": 62.0,
        "m64f": {"raw": 0, "prochot": False, "thermal": False, "pl1": False},
        "m19c": {"raw": 0, "prochot": False, "power": False, "current": 30},
        "thr": {"pkg": 12, "core": 3, "pkg_d": 0, "core_d": 0},
        "psi": {"cpu": 1.2, "io": 0.0, "mem": 0.0},
        "load1": 2.1,
        "ac": 1,
        "bat": {"status": "Charging", "current_now": None},
        "top": [
            {"pid": 111, "comm": "node", "pcpu": 45.2},
            {"pid": 222, "comm": "claude", "pcpu": 12.0},
            {"pid": 333, "comm": "bash", "pcpu": 1.0},
        ],
    }
    print(json.dumps(synthetic, separators=(",", ":")))
    required_top = {"t", "f", "temp_c", "m64f", "m19c", "thr", "psi", "load1", "ac", "bat", "top"}
    assert required_top.issubset(synthetic.keys()), "sample missing required fields"
    assert set(synthetic["f"].keys()) == {"min", "max", "avg"}
    assert set(synthetic["m64f"].keys()) == {"raw", "prochot", "thermal", "pl1"}
    assert set(synthetic["m19c"].keys()) == {"raw", "prochot", "power", "current"}
    assert set(synthetic["thr"].keys()) == {"pkg", "core", "pkg_d", "core_d"}
    assert set(synthetic["psi"].keys()) == {"cpu", "io", "mem"}
    assert set(synthetic["bat"].keys()) == {"status", "current_now"}
    assert len(synthetic["top"]) == 3

    ncores = 16
    assert is_anomaly_rising(dict(synthetic), ncores) is False, "calm sample wrongly flagged anomaly"

    anomaly_sample = dict(synthetic)
    anomaly_sample["thr"] = {"pkg": 13, "core": 3, "pkg_d": 1, "core_d": 0}
    assert is_anomaly_rising(anomaly_sample, ncores) is True, "injected throttle delta did not trip anomaly"

    prochot_sample = dict(synthetic)
    prochot_sample["m64f"] = {"raw": 1, "prochot": True, "thermal": False, "pl1": False}
    assert is_anomaly_rising(prochot_sample, ncores) is True, "prochot bit did not trip anomaly"

    freq_sample = dict(synthetic)
    freq_sample["f"] = {"min": 700000, "max": 4000000, "avg": 2000000}
    freq_sample["load1"] = ncores * 0.6
    assert is_anomaly_rising(freq_sample, ncores) is True, "low-freq+load did not trip anomaly"

    psi_sample = dict(synthetic)
    psi_sample["psi"] = {"cpu": 25.0, "io": 0.0, "mem": 0.0}
    assert is_anomaly_rising(psi_sample, ncores) is True, "PSI spike did not trip anomaly"

    import tempfile
    tmpdir = tempfile.mkdtemp(prefix="hw-blackbox-selftest-")
    try:
        fake_ring = Ring(RING_N)
        for i in range(10):
            fake_ring.append({"t": i, "note": "pre-event"})
        would_flush_path = os.path.join(tmpdir, f"event-{int(time.time())}.jsonl")
        with open(would_flush_path, "w") as f:
            for line in fake_ring.dump_lines():
                f.write(line + "\n")
        with open(would_flush_path) as f:
            written = f.readlines()
        assert len(written) == 10, "flush did not write full pre-event ring window"
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)

    print("self-test OK", file=sys.stderr)
    return True


if __name__ == "__main__":
    if "--self-test" in sys.argv:
        sys.exit(0 if self_test() else 1)
    run_daemon()