#!/usr/bin/env python3
"""rigbench_submit.py: measure local LLM speed on this machine and prepare a RigBench submission.

  python3 rigbench_submit.py --llama-bench ./llama-bench --model ./llama-2-7b.Q4_0.gguf
  python3 rigbench_submit.py --ollama llama3.1:8b
  python3 rigbench_submit.py --parse bench.json --model llama-2-7b.Q4_0.gguf   # existing `llama-bench -o json` output

Writes rigbench-result.json. Nothing is uploaded unless you add --submit, which POSTs that JSON
(and nothing else) to the RigBench site for review. Python 3.8+, standard library only.
License: MIT.
"""
import argparse, datetime, json, os, platform, re, shutil, subprocess, sys, urllib.error, urllib.parse, urllib.request

VERSION = "1.0.0"
DEFAULT_SITE = "https://rigbench.pages.dev"
QUANT_RE = re.compile(r"(?<![A-Za-z0-9])((?:UD-)?(?:IQ[1-4]_(?:XXS|XS|S|M|NL)|Q[2-8]_K_(?:XL|L|M|S)|Q[2-8]_K|Q[4-8]_[01]|MXFP4(?:_MOE)?|BF16|F16|F32))(?![A-Za-z0-9])", re.I)


def sh(cmd, timeout=30):
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout).stdout.strip()
    except Exception:
        return ""


def detect_hardware():
    hw = {"name": None, "vendor": None, "memory_gb": None, "driver": None, "gpu_cores": None,
          "os": f"{platform.system()} {platform.release()}", "cpu": platform.processor() or platform.machine()}
    sysname = platform.system()
    if sysname == "Darwin":
        chip = sh(["sysctl", "-n", "machdep.cpu.brand_string"])
        mem = sh(["sysctl", "-n", "hw.memsize"])
        hw.update(name=chip or "Apple Silicon", vendor="Apple", cpu=chip)
        if mem.isdigit():
            hw["memory_gb"] = round(int(mem) / 2**30)
        try:
            d = json.loads(sh(["system_profiler", "SPDisplaysDataType", "-json"], 60) or "{}")
            for g in d.get("SPDisplaysDataType", []):
                if g.get("sppci_cores"):
                    hw["gpu_cores"] = int(g["sppci_cores"])
        except Exception:
            pass
        if hw["gpu_cores"]:
            hw["name"] = f"{hw['name']} ({hw['gpu_cores']}-core GPU)"
        return hw
    if shutil.which("nvidia-smi"):
        out = sh(["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader,nounits"])
        if out:
            name, mem, drv = [x.strip() for x in out.splitlines()[0].split(",")[:3]]
            hw.update(name=name, vendor="NVIDIA", memory_gb=round(float(mem) / 1024), driver=drv)
            return hw
    if shutil.which("rocm-smi"):
        out = sh(["rocm-smi", "--showproductname", "--json"])
        try:
            d = json.loads(out)
            card = next(iter(d.values()))
            name = card.get("Card Series") or card.get("Card series") or card.get("Card Model")
            hw.update(name=name, vendor="AMD")
            m = sh(["rocm-smi", "--showmeminfo", "vram", "--json"])
            v = next(iter(json.loads(m).values()))
            tot = v.get("VRAM Total Memory (B)")
            if tot:
                hw["memory_gb"] = round(int(tot) / 2**30)
            return hw
        except Exception:
            pass
    if sysname == "Windows":
        out = sh(["powershell", "-NoProfile", "-Command", "(Get-CimInstance Win32_VideoController | Select-Object -First 1).Name"])
        if out:
            hw["name"] = out
    elif shutil.which("lspci"):
        for line in sh(["lspci"]).splitlines():
            if re.search(r"VGA|3D|Display", line):
                hw["name"] = line.split(":", 2)[-1].strip()
                break
    return hw


def model_key_from_file(path):
    base = os.path.basename(path)
    base = re.sub(r"-\d{5}-of-\d{5}", "", base)
    base = re.sub(r"\.gguf$", "", base, flags=re.I)
    q = QUANT_RE.search(base)
    name = base[: q.start()] if q else base
    name = re.sub(r"[-_.]+$", "", name)
    key = re.sub(r"[^a-z0-9.]+", "-", name.lower()).strip("-")
    return key, (q.group(1).upper() if q else None), name


def parse_llama_bench(entries, model_path):
    key, quant, name = model_key_from_file(model_path or (entries[0].get("model_filename") or "model"))
    results = {}
    for e in entries:
        fa = bool(e.get("flash_attn"))
        r = results.setdefault(fa, dict(model_key=key, model_name=name, model_file=os.path.basename(model_path or e.get("model_filename", "")),
                                        model_size_bytes=e.get("model_size"), params=e.get("model_n_params"),
                                        quant=quant or e.get("model_type", "").split()[-1], backend=e.get("backends") or e.get("backend"),
                                        flash_attn=fa, ngl=e.get("n_gpu_layers"), build=f'{e.get("build_commit", "")} ({e.get("build_number", "")})',
                                        gpu_info=e.get("gpu_info"), pp_tps=None, pp_test=None, tg_tps=None, tg_test=None))
        ts = e.get("avg_ts")
        if e.get("n_prompt", 0) > 0 and e.get("n_gen", 0) == 0:
            r["pp_tps"], r["pp_test"], r["pp_stddev"] = round(ts, 2), f'pp{e["n_prompt"]}', round(e.get("stddev_ts", 0), 2)
        elif e.get("n_gen", 0) > 0 and e.get("n_prompt", 0) == 0:
            r["tg_tps"], r["tg_test"], r["tg_stddev"] = round(ts, 2), f'tg{e["n_gen"]}', round(e.get("stddev_ts", 0), 2)
    return list(results.values())


def run_llama_bench(binary, model, ngl, fa):
    cmd = [binary, "-m", model, "-ngl", str(ngl), "-fa", fa, "-o", "json"]
    print("running:", " ".join(cmd), file=sys.stderr)
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        sys.exit(f"llama-bench failed:\n{p.stderr[-2000:]}")
    return json.loads(p.stdout)


def ollama_req(url, path, payload=None):
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(url.rstrip("/") + path, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=600) as r:
        return json.loads(r.read())


def run_ollama(model, url, runs=3):
    show = ollama_req(url, "/api/show", {"model": model})
    det = show.get("details", {})
    size = None
    for t in ollama_req(url, "/api/tags").get("models", []):
        if t.get("name") == model or t.get("model") == model:
            size = t.get("size")
    prompt = "Explain, in about 300 words, how a graphics card's memory bandwidth limits how fast a language model can write text."
    tgs, pps = [], []
    for i in range(runs + 1):
        r = ollama_req(url, "/api/generate", {"model": model, "prompt": prompt, "stream": False,
                                               "options": {"num_predict": 256, "temperature": 0, "seed": 1}})
        if i == 0:
            continue  # warm-up (load into memory)
        if r.get("eval_duration"):
            tgs.append(r["eval_count"] / (r["eval_duration"] / 1e9))
        if r.get("prompt_eval_duration") and r.get("prompt_eval_count"):
            pps.append(r["prompt_eval_count"] / (r["prompt_eval_duration"] / 1e9))
        print(f"run {i}: {tgs[-1]:.1f} tok/s", file=sys.stderr)
    tgs.sort(); pps.sort()
    key = re.sub(r"[^a-z0-9.]+", "-", model.lower()).strip("-")
    return [dict(model_key=key, model_name=model, model_file=None, model_size_bytes=size, params=det.get("parameter_size"),
                 quant=(det.get("quantization_level") or "").upper() or None, backend="ollama", flash_attn=None, ngl=None,
                 build=sh(["ollama", "--version"]) or None, tg_tps=round(tgs[len(tgs) // 2], 2) if tgs else None,
                 tg_test="ollama 256-token generation (median of %d)" % len(tgs),
                 pp_tps=round(pps[len(pps) // 2], 2) if pps else None, pp_test="ollama short prompt" if pps else None)]


def submit(site, result):
    req = urllib.request.Request(site.rstrip("/") + "/api/submit", method="POST", data=json.dumps(result).encode(),
                                 headers={"Content-Type": "application/json", "User-Agent": "rigbench_submit/" + VERSION})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        try:
            return json.loads(e.read())
        except Exception:
            return {"ok": False, "error": f"HTTP {e.code}"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--llama-bench", help="path to llama.cpp's llama-bench binary")
    ap.add_argument("--model", help="path to the .gguf model file (llama-bench mode)")
    ap.add_argument("--ngl", default=99, type=int, help="GPU layers (default 99 = all)")
    ap.add_argument("--fa", default="0,1", help="flash attention settings to test (default 0,1)")
    ap.add_argument("--ollama", metavar="MODEL", help="benchmark an Ollama model, e.g. llama3.1:8b")
    ap.add_argument("--ollama-url", default=os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434"))
    ap.add_argument("--parse", metavar="JSON", help="use existing `llama-bench -o json` output instead of running")
    ap.add_argument("--gpu-name", help="override the detected GPU name")
    ap.add_argument("--submit", action="store_true", help="send the result to RigBench for review")
    ap.add_argument("--site", default=DEFAULT_SITE, help="RigBench site that receives submissions")
    ap.add_argument("--out", default="rigbench-result.json")
    a = ap.parse_args()

    hw = detect_hardware()
    if a.gpu_name:
        hw["name"] = a.gpu_name
    if a.parse:
        results = parse_llama_bench(json.load(open(a.parse)), a.model)
    elif a.llama_bench:
        if not a.model:
            sys.exit("--model is required with --llama-bench")
        results = parse_llama_bench(run_llama_bench(a.llama_bench, a.model, a.ngl, a.fa), a.model)
    elif a.ollama:
        results = run_ollama(a.ollama, a.ollama_url)
    else:
        ap.print_help()
        sys.exit(1)
    if not hw.get("name"):
        gi = next((r.get("gpu_info") for r in results if r.get("gpu_info")), None)
        hw["name"] = gi or "Unknown GPU (edit before submitting)"
    out = dict(schema="rigbench-result/1", created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
               tool=dict(name="rigbench_submit.py", version=VERSION), hardware=hw, results=results)
    with open(a.out, "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))
    print(f"\nSaved {a.out}.")
    if a.submit:
        r = submit(a.site, out)
        print(("Submitted for review. Reference: " + r.get("id", "")) if r.get("ok") else
              f"Not submitted ({r.get('message') or r.get('error')}). Your result is saved in {a.out}; you can upload it at {a.site}/submit/")
    else:
        print(f"Review it, then send it with --submit, or upload the file at {a.site}/submit/")


if __name__ == "__main__":
    main()