#!/usr/bin/env python3
"""
bench-vs — the SAME prompt through llama.cpp (router :11343) and vLLM (:11348), one table.

Measures per target: decode t/s, prefill t/s, TTFT, VRAM used/free (amdgpu sysfs, no root needed),
optionally decode at a long prompt (--ctx) and aggregate throughput under N parallel requests
(--concurrency). Method = ../llamacpp/bench.py (streamed chat completion, usage.completion_tokens
over wall time after the first token), applied identically to both engines; llama.cpp's own
`timings` figure is shown in parentheses so rows stay comparable with bench-results.md.

Thinking is OFF by default on both engines (chat_template_kwargs.enable_thinking=false) and the
first reasoning_content delta counts as the first token anyway — a reasoning preamble of different
length per engine would otherwise poison TTFT. --think re-enables it.

GUARDS (the card drives the display; two resident engines = GTT spill = frozen PC):
  - a llamacpp target is refused while vLLM answers /health (stop it: ./vllmctl.py down)
  - a vllm target is refused while the router reports a loaded model or llama.service is active

  ./bench-vs.py -t vllm:gemma-4-E4B-it-w4a16                  # vLLM up (vllmctl.py up …), router stopped
  ./bench-vs.py -t llamacpp:gemma-4-E4B-it-UD-Q8_K_XL         # vLLM down, router started
  ./bench-vs.py -t llamacpp:Qwen3.5-9B-UD-Q6_K_XL --ctx 8000 --concurrency 4
Output: bench-vs-results.md (latest) + bench-vs-history.md (append), next to this script. Stdlib only.
"""
import argparse, json, os, re, subprocess, sys, threading, time, urllib.request, urllib.error
from datetime import datetime

HERE = os.path.dirname(os.path.abspath(__file__))
DRM = "/sys/class/drm/renderD128/device"
URLS = {"llamacpp": "http://127.0.0.1:11343/v1", "vllm": "http://127.0.0.1:11348/v1"}
TASK = "Write a Python function that merges two sorted lists, with a short docstring and one example."


def sysfs(name):
    try:
        return int(open(os.path.join(DRM, name)).read().strip())
    except (OSError, ValueError):
        return None


def gb(b):
    return f"{b / 1e9:.1f} GB" if b is not None else "-"


def http(url, key, payload=None, timeout=600):
    data = json.dumps(payload).encode() if payload is not None else None
    hdrs = {"Authorization": f"Bearer {key}"}
    if data:
        hdrs["Content-Type"] = "application/json"
    return urllib.request.urlopen(urllib.request.Request(url, data=data, headers=hdrs), timeout=timeout)


def parse_ini(path):
    sections, sec = {}, None
    try:
        lines = open(path, encoding="utf-8").read().splitlines()
    except OSError:
        return {}
    for ln in lines:
        s = ln.strip()
        if not s or s.startswith("#"):
            continue
        if s.startswith("[") and s.endswith("]"):
            sec = s[1:-1]; sections[sec] = {}
        elif sec and "=" in s:
            k, v = s.split("=", 1)
            sections[sec][k.strip()] = v.split(" #", 1)[0].split("\t#", 1)[0].strip()
    return sections


def run(base, key, model, prompt, n, temp, timeout, think):
    payload = {"model": model, "messages": [{"role": "user", "content": prompt}],
               "max_tokens": n, "temperature": temp, "stream": True,
               "stream_options": {"include_usage": True},
               "timings_per_token": True}                      # llama.cpp extension, ignored by vLLM
    if not think:
        payload["chat_template_kwargs"] = {"enable_thinking": False}   # both engines accept it; templates without it ignore it
    t0 = time.perf_counter(); ttft = None; ntok = 0; usage = None; timings = None
    for raw in http(base.rstrip("/") + "/chat/completions", key, payload, timeout):
        line = raw.decode("utf-8", "ignore").strip()
        if not line.startswith("data:"):
            continue
        body = line[5:].strip()
        if body == "[DONE]":
            break
        try:
            c = json.loads(body)
        except ValueError:
            continue
        ch = c.get("choices") or [{}]
        d = ch[0].get("delta", {}) if ch else {}
        if d.get("content") or d.get("reasoning_content") or d.get("reasoning"):
            if ttft is None:
                ttft = time.perf_counter() - t0
            ntok += 1
        if c.get("usage"):
            usage = c["usage"]
        if c.get("timings"):
            timings = c["timings"]
    total = time.perf_counter() - t0
    comp = (usage or {}).get("completion_tokens") or ntok
    ptok = (usage or {}).get("prompt_tokens")
    gen_s = total - (ttft or total)
    tg = comp / gen_s if gen_s > 0 else 0.0
    pp = (ptok / ttft) if (ptok and ttft) else None
    srv = timings.get("predicted_per_second") if timings else None    # llama.cpp's own number, for the ledger
    return {"tg": tg, "pp": pp, "ttft": ttft or 0.0, "ptok": ptok or 0, "ctok": comp, "wall": total, "srv": srv,
            "t0": t0, "t1": t0 + total}


def with_retries(fn, retries, wait, label=""):
    last = None
    for i in range(retries + 1):
        try:
            return fn()
        except urllib.error.HTTPError as e:
            body = e.read().decode("utf-8", "ignore")[:300]
            if e.code < 500 or "failed to load" in body.lower():
                raise
            last = e
        except (urllib.error.URLError, ConnectionError, TimeoutError) as e:
            last = e
        if i < retries:
            print(f"    … {label}: not ready ({last}); retry in {wait}s", file=sys.stderr)
            time.sleep(wait)
    raise last


def vllm_up(url):
    try:
        return http(re.sub(r"/v1/?$", "", url) + "/health", "x", timeout=3).status == 200
    except Exception:
        return False


def _status(m):
    """Router /models status: newer builds send {'value': 'loaded'|'sleeping'|'unloaded'|…}, older a string."""
    st = m.get("status", "loaded")
    return st.get("value", "loaded") if isinstance(st, dict) else st


def router_loaded(url):
    try:
        d = json.load(http(re.sub(r"/v1/?$", "", url) + "/models", "x", timeout=3))
        return [m["id"] for m in d.get("data", []) if _status(m) in ("loaded", "loading")]
    except Exception:
        return []


def router_active():
    try:
        r = subprocess.run(["systemctl", "is-active", "llama.service"], capture_output=True, text=True, timeout=10)
    except (OSError, subprocess.TimeoutExpired):
        return False
    return r.stdout.strip() == "active"


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("-t", "--target", action="append", required=True,
                    help="engine:model, engine = llamacpp | vllm (repeatable)")
    ap.add_argument("--url", action="append", default=[], help="engine=URL override, e.g. vllm=http://host:11348/v1")
    ap.add_argument("-k", "--key", default="no-key-required")
    ap.add_argument("-n", "--tokens", type=int, default=256)
    ap.add_argument("--temp", type=float, default=0.3)
    ap.add_argument("--ctx", type=int, default=0, help="also record decode at ~this many prompt tokens")
    ap.add_argument("--concurrency", type=int, default=0, help="also record aggregate t/s with N parallel requests")
    ap.add_argument("--think", action="store_true", help="leave the model's thinking ON (default: off on both engines)")
    ap.add_argument("--timeout", type=int, default=600)
    ap.add_argument("--retries", type=int, default=4); ap.add_argument("--retry-wait", type=int, default=15)
    ap.add_argument("--force", action="store_true", help="skip the two-engines-resident guards")
    ap.add_argument("--out", default=os.path.join(HERE, "bench-vs-results.md"))
    ap.add_argument("--history", default=os.path.join(HERE, "bench-vs-history.md"))
    a = ap.parse_args()

    urls = dict(URLS)
    for u in a.url:
        k, v = u.split("=", 1); urls[k] = v
    llcfg = parse_ini(os.path.join(HERE, "..", "llamacpp", "config.ini"))
    vlcfg = parse_ini(os.path.join(HERE, "vllm.ini"))
    total = sysfs("mem_info_vram_total")
    if total is None:
        print("  (amdgpu sysfs not readable here — VRAM columns are OFF; run this ON fl for them)", file=sys.stderr)

    filler = ("The quick brown fox jumps over the lazy dog. " * max(1, a.ctx // 9)) if a.ctx else None
    cols = ["engine", "model", "cfg ctx", "spec", "decode t/s", "prefill t/s", "TTFT s", "VRAM used", "VRAM free"]
    if a.ctx:
        cols.append(f"decode@{a.ctx // 1000}k")
    if a.concurrency:
        cols.append(f"agg t/s @{a.concurrency}")

    latest = open(a.out, "w", encoding="utf-8"); hist = open(a.history, "a", encoding="utf-8")

    def emit(line):
        print(line); latest.write(line + "\n"); latest.flush(); hist.write(line + "\n"); hist.flush()

    latest.write("# llama.cpp vs vLLM (latest run)\n\n"); hist.write(f"\n## run {datetime.now():%Y-%m-%d %H:%M}\n\n")
    emit(f"_Run {datetime.now():%Y-%m-%d %H:%M} · gen {a.tokens} tok · temp {a.temp} · thinking "
         f"{'ON' if a.think else 'OFF'} · VRAM total {gb(total)} · baseline used {gb(sysfs('mem_info_vram_used'))} · "
         f"decode = usage tokens / wall after first token on BOTH engines (llama.cpp's own timings in parentheses)_\n")
    emit("| " + " | ".join(cols) + " |"); emit("|" + "|".join(["---"] * len(cols)) + "|")

    for t in a.target:
        eng, _, model = t.partition(":")
        if eng not in urls or not model:
            emit(f"| {eng} | {model} | bad target — use engine:model with engine in {list(urls)} |"); continue
        base = urls[eng]
        # guards
        if not a.force:
            if eng == "llamacpp" and vllm_up(urls["vllm"]):
                emit(f"| llamacpp | {model} | ⚠ SKIPPED: vLLM is up on {urls['vllm']} — loading a router model on top = freeze. "
                     f"`vllmctl.py down` first |"); continue
            if eng == "vllm":
                loaded = router_loaded(urls["llamacpp"])
                if loaded or router_active():
                    emit(f"| vllm | {model} | ⚠ SKIPPED: router loaded={loaded or '-'} active={router_active()} — "
                         f"stop llama.service before benching vLLM |"); continue
        if eng == "llamacpp":
            c = llcfg.get(model, {})
            cfgctx, spec = c.get("ctx-size", "-"), (c.get("spec-type") or "-")
        else:
            c = vlcfg.get(model, {})
            cfgctx = c.get("max-model-len", "-")
            spec = "mtp" if c.get("speculative-config") else "-"
        cells = [eng, model, cfgctx, spec]
        try:
            with_retries(lambda: run(base, a.key, model, "hi", 8, a.temp, a.timeout, a.think), a.retries, a.retry_wait, model)
            used = sysfs("mem_info_vram_used")
            r = with_retries(lambda: run(base, a.key, model, TASK, a.tokens, a.temp, a.timeout, a.think), a.retries, a.retry_wait, model)
            dec = f"{r['tg']:.1f}" + (f" ({r['srv']:.1f})" if r["srv"] else "")
            cells += [dec, f"{r['pp']:.0f}" if r["pp"] else "-", f"{r['ttft']:.2f}", gb(used),
                      gb(total - used) if (total and used is not None) else "-"]
            if a.ctx:
                rc = run(base, a.key, model, filler + "\n\n" + TASK, a.tokens, a.temp, a.timeout, a.think)
                cells.append(f"{rc['tg']:.1f}")
            if a.concurrency:
                res = [None] * a.concurrency
                def worker(i):
                    res[i] = run(base, a.key, model, TASK + f" (variant {i})", a.tokens, a.temp, a.timeout, a.think)
                th = [threading.Thread(target=worker, args=(i,)) for i in range(a.concurrency)]
                for x in th: x.start()
                for x in th: x.join()
                ok = [x for x in res if x]
                wall = max(x["t1"] for x in ok) - min(x["t0"] for x in ok)
                cells.append(f"{sum(x['ctok'] for x in ok) / wall:.1f}" if wall > 0 else "-")
        except Exception as e:
            err = f"HTTP {e.code} {e.read().decode('utf-8', 'ignore')[:120]!r}" if isinstance(e, urllib.error.HTTPError) else f"ERROR {e}"
            cells = cells[:4] + [err] + [""] * (len(cols) - 5)   # same width whether it died early or in the --ctx/--concurrency pass
        emit("| " + " | ".join(str(x) for x in cells) + " |")

    latest.write("\n_llama.cpp rows run with the preset's MTP draft when `spec` says so; vLLM rows have no speculation "
                 "unless `spec` = mtp. `agg t/s @N` = N parallel requests, completion tokens over the wall time of "
                 "the whole batch (llama-server's default slot count serves them in parallel too)._\n")
    latest.close(); hist.close()
    print(f"\nwrote {a.out} (+ appended {a.history})")


if __name__ == "__main__":
    main()