#!/usr/bin/env python3
"""Sample CPU and GPU load while llama-server generates.

Answers one question: during tg, is the CPU pinned (compute-bound on the
4 Ivy Bridge cores) or is everything waiting on memory/PCIe?
"""
import glob
import json
import os
import subprocess
import sys
import threading
import time
import http.client

BIN = os.path.expanduser("~/llama/llama-b11960/llama-server")
MODEL = os.path.expanduser("~/models/Qwen3.6-35B-A3B-MTP-UD-IQ3_XXS.gguf")
PORT = 8299
NCPU = os.cpu_count()

PROMPT = ("You are a coding agent. Explain, step by step and in detail, how to "
          "diagnose why a systemd service keeps restarting on a remote Ubuntu "
          "server, which commands to run, and how to read each output. ") * 8


def read_cpu(pid):
    with open(f"/proc/{pid}/stat") as f:
        parts = f.read().split()
    return (int(parts[13]) + int(parts[14])) / os.sysconf("SC_CLK_TCK")


def gpu_busy():
    for p in glob.glob("/sys/class/drm/card*/device/gpu_busy_percent"):
        try:
            return int(open(p).read())
        except (OSError, ValueError):
            pass
    return None


def vram_mb():
    for p in glob.glob("/sys/class/drm/card*/device/mem_info_vram_used"):
        try:
            return int(open(p).read()) // (1024 * 1024)
        except (OSError, ValueError):
            pass
    return None


def freqs():
    out = []
    for p in sorted(glob.glob("/sys/devices/system/cpu/cpu[0-9]*/cpufreq/scaling_cur_freq")):
        try:
            out.append(int(open(p).read()) // 1000)
        except (OSError, ValueError):
            pass
    return out


def governor():
    try:
        return open("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor").read().strip()
    except OSError:
        return "?"


def post(path, payload, timeout=1800):
    c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=timeout)
    c.request("POST", path, body=json.dumps(payload),
              headers={"Content-Type": "application/json"})
    r = c.getresponse()
    body = r.read()
    c.close()
    if r.status != 200:
        raise RuntimeError(f"{path} {r.status}: {body[:300]!r}")
    return json.loads(body)


def main():
    flags = sys.argv[1:] or ["--spec-type", "draft-mtp", "--spec-draft-n-max", "1"]
    log = open(os.path.expanduser("~/ngram/smoke/busy.log"), "wb")
    proc = subprocess.Popen(
        [BIN, "-m", MODEL, "--port", str(PORT), "-c", "8192", "-ngl", "99",
         "--n-cpu-moe", "999", "-fa", "on", "-ctk", "q8_0", "-ctv", "q8_0",
         "-t", "4", "--no-webui", *flags],
        stdout=log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL)
    try:
        for _ in range(600):
            if proc.poll() is not None:
                raise SystemExit("server died")
            try:
                c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
                c.request("GET", "/health")
                ok = c.getresponse().status == 200
                c.close()
                if ok:
                    break
            except OSError:
                pass
            time.sleep(1)

        post("/completion", {"prompt": PROMPT, "n_predict": 8,
                             "temperature": 0, "cache_prompt": False}, 600)

        print(f"nproc={NCPU} governor={governor()} freqs(MHz)={freqs()}")
        print(f"vram idle={vram_mb()} MB  gpu_busy={gpu_busy()}%")

        samples = []
        result = {}

        def run():
            result["out"] = post("/completion", {
                "prompt": PROMPT, "n_predict": 256, "temperature": 0,
                "cache_prompt": False}, 1800)

        t = threading.Thread(target=run)
        cpu0, t0 = read_cpu(proc.pid), time.time()
        t.start()
        while t.is_alive():
            time.sleep(0.25)
            samples.append((time.time(), read_cpu(proc.pid), gpu_busy(), vram_mb()))
        t.join()
        cpu1, t1 = read_cpu(proc.pid), time.time()

        tm = result["out"]["timings"]
        pred_s = tm["predicted_ms"] / 1000
        pp_s = tm["prompt_ms"] / 1000
        gen = [s for s in samples if s[0] >= t0 + pp_s]

        def stat(rows, idx):
            vals = [r[idx] for r in rows if r[idx] is not None]
            return (sum(vals) / len(vals), min(vals), max(vals)) if vals else (None,) * 3

        print(f"\nprompt {tm['prompt_n']} tok in {pp_s:.1f}s ({tm['prompt_n']/pp_s:.1f} t/s), "
              f"gen {tm['predicted_n']} tok in {pred_s:.1f}s "
              f"({tm['predicted_n']/pred_s:.2f} t/s)")
        print(f"total process cpu: {cpu1-cpu0:.1f}s over {t1-t0:.1f}s wall "
              f"= {(cpu1-cpu0)/(t1-t0)*100:.0f}% of {NCPU*100}%")

        # CPU% inside the generation window, from consecutive samples.
        rows = gen
        if len(rows) >= 2:
            dc = rows[-1][1] - rows[0][1]
            dw = rows[-1][0] - rows[0][0]
            print(f"during tg: cpu {dc/dw*100:.0f}% of {NCPU*100}%")
        for name, idx in (("gpu_busy%", 2), ("vram MB", 3)):
            avg, lo, hi = stat(rows, idx)
            if avg is not None:
                print(f"during tg: {name} avg {avg:.0f} min {lo} max {hi}")
        print(f"freqs after (MHz)={freqs()}")
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=30)
        except subprocess.TimeoutExpired:
            proc.kill()
        log.close()


if __name__ == "__main__":
    main()