#!/usr/bin/python3
"""ahel poolstat — minimal node exporter for the M4 pool nodes (:9101/metrics).

Samples powermetrics (root required) every 30s in a background thread and serves
cached values instantly. Metrics: erebor_node_load1, erebor_mem_used_pct,
erebor_thermal_pressure (0=Nominal..4=Sleeping), erebor_gpu_power_mw, erebor_cpu_power_mw,
erebor_cpu_temp_c, erebor_gpu_temp_c (via macmon — Apple Silicon die temps that
powermetrics does not expose; 0.0 when macmon is absent/unreadable).
"""
import json
import os
import re
import subprocess
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer

PRESSURE = {"Nominal": 0, "Fair": 1, "Serious": 2, "Critical": 3, "Sleeping": 4}
state = {"thermal": 0, "gpu_mw": 0.0, "cpu_mw": 0.0, "cpu_temp": 0.0, "gpu_temp": 0.0}
MACMON = "/opt/homebrew/bin/macmon"


def sample_temps():
    # macmon reads IOReport die temps (no root needed; we already run as root).
    # One 1s sample, first JSON line: {"temp":{"cpu_temp_avg":..,"gpu_temp_avg":..},...}
    try:
        out = subprocess.run([MACMON, "pipe", "-s", "1"],
                             capture_output=True, text=True, timeout=20).stdout
        temp = json.loads(out.splitlines()[0]).get("temp", {})
        cpu = float(temp.get("cpu_temp_avg", 0.0))
        gpu = float(temp.get("gpu_temp_avg", 0.0))
        if cpu > 0:
            state["cpu_temp"] = cpu
        if gpu > 0:
            state["gpu_temp"] = gpu
    except Exception:
        pass


def sampler():
    while True:
        try:
            out = subprocess.run(
                ["/usr/bin/powermetrics", "-s", "cpu_power,gpu_power,thermal",
                 "-n", "1", "-i", "1000"],
                capture_output=True, text=True, timeout=30).stdout
            m = re.search(r"pressure level:\s*(\w+)", out)
            if m:
                state["thermal"] = PRESSURE.get(m.group(1), 0)
            m = re.search(r"GPU Power:\s*(\d+)\s*mW", out)
            if m:
                state["gpu_mw"] = float(m.group(1))
            m = re.search(r"(?:CPU|Combined) Power:\s*(\d+)\s*mW", out)
            if m:
                state["cpu_mw"] = float(m.group(1))
        except Exception:
            pass
        sample_temps()
        time.sleep(30)


def mem_used_pct():
    try:
        out = subprocess.run(["/usr/bin/vm_stat"], capture_output=True,
                             text=True, timeout=5).stdout
        pages = dict(re.findall(r"Pages\s+([\w ]+):\s+(\d+)", out))
        page_size = int(re.search(r"page size of (\d+)", out).group(1))
        free = int(pages.get("free", 0)) + int(pages.get("purgeable", 0))
        total = int(subprocess.run(["/usr/sbin/sysctl", "-n", "hw.memsize"],
                                   capture_output=True, text=True, timeout=5).stdout)
        used = total - free * page_size
        return 100.0 * used / total
    except Exception:
        return 0.0


class H(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path != "/metrics":
            self.send_response(404)
            self.end_headers()
            return
        load1 = os.getloadavg()[0]
        body = (
            "# TYPE erebor_node_load1 gauge\n"
            f"erebor_node_load1 {load1:.2f}\n"
            "# TYPE erebor_mem_used_pct gauge\n"
            f"erebor_mem_used_pct {mem_used_pct():.1f}\n"
            "# TYPE erebor_thermal_pressure gauge\n"
            f"erebor_thermal_pressure {state['thermal']}\n"
            "# TYPE erebor_gpu_power_mw gauge\n"
            f"erebor_gpu_power_mw {state['gpu_mw']:.0f}\n"
            "# TYPE erebor_cpu_power_mw gauge\n"
            f"erebor_cpu_power_mw {state['cpu_mw']:.0f}\n"
            "# TYPE erebor_cpu_temp_c gauge\n"
            f"erebor_cpu_temp_c {state['cpu_temp']:.1f}\n"
            "# TYPE erebor_gpu_temp_c gauge\n"
            f"erebor_gpu_temp_c {state['gpu_temp']:.1f}\n"
        ).encode()
        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    threading.Thread(target=sampler, daemon=True).start()
    HTTPServer(("0.0.0.0", int(os.environ.get("POOLSTAT_PORT", "9101"))), H).serve_forever()