#!/usr/bin/env python3
"""Will It Run? auto-benchmark (v1) - https://willitrunllm.com/submit/

Runs llama.cpp's llama-bench (pp512 + tg128) on a GGUF model, detects your hardware,
shows you EXACTLY what would be sent, and only submits if you type "y".
Submissions are reviewed by hand before they appear on willitrunllm.com. No email needed,
nothing is sent without your confirmation, no telemetry.

Usage:
  curl -fsSL https://willitrunllm.com/bench.py | python3 - --model ~/models/Qwen3-8B-Q4_K_M.gguf
  python3 bench.py --model /path/model.gguf [--llama-bench /path/llama-bench] [--ngl 99]
  python3 bench.py --hf-repo bartowski/Qwen_Qwen3-8B-GGUF --hf-file Qwen_Qwen3-8B-Q4_K_M.gguf
  python3 bench.py --ollama llama3.1:8b           # lower-fidelity Ollama "eval rate" instead of llama-bench
  python3 bench.py --model m.gguf --dry-run       # run + show payload, never send
  python3 bench.py --sample-json out.json --dry-run   # parse an existing `llama-bench -o json` file
Only the Python standard library is required (huggingface_hub optional, for --hf-repo downloads).
"""
import argparse, json, os, platform, re, shutil, subprocess, sys, time, urllib.request, urllib.error

VERSION = "bench.py/1.0"
SITE = os.environ.get("WIR_SITE", "https://willitrunllm.com")
START_MS = int(time.time() * 1000)


def run(cmd, timeout=30):
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        return p.stdout if p.returncode == 0 else ""
    except Exception:
        return ""


# ---------------- hardware detection ----------------
def ram_gb():
    try:
        if sys.platform.startswith("linux"):
            for l in open("/proc/meminfo"):
                if l.startswith("MemTotal:"):
                    return round(int(l.split()[1]) / 1048576, 1)
        if sys.platform == "darwin":
            return round(int(run(["sysctl", "-n", "hw.memsize"]).strip()) / 2**30, 1)
        if sys.platform == "win32":
            import ctypes
            class M(ctypes.Structure):
                _fields_ = [("l", ctypes.c_ulong), ("m", ctypes.c_ulong), ("t", ctypes.c_ulonglong), ("a", ctypes.c_ulonglong),
                            ("tp", ctypes.c_ulonglong), ("ap", ctypes.c_ulonglong), ("tv", ctypes.c_ulonglong), ("av", ctypes.c_ulonglong), ("e", ctypes.c_ulonglong)]
            m = M(); m.l = ctypes.sizeof(M); ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(m))
            return round(m.t / 2**30, 1)
    except Exception:
        pass
    return None


def cpu_name():
    try:
        if sys.platform.startswith("linux"):
            for l in open("/proc/cpuinfo"):
                if l.lower().startswith("model name"):
                    return l.split(":", 1)[1].strip()
        if sys.platform == "darwin":
            return run(["sysctl", "-n", "machdep.cpu.brand_string"]).strip()
        if sys.platform == "win32":
            out = run(["powershell", "-NoProfile", "-Command", "(Get-CimInstance Win32_Processor).Name"])
            if out.strip():
                return out.strip().splitlines()[0]
    except Exception:
        pass
    return platform.processor() or None


def gpus():
    out = []
    if shutil.which("nvidia-smi"):
        for l in run(["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"]).splitlines():
            if "," in l:
                n, mem = l.rsplit(",", 1)
                try: out.append({"vendor": "nvidia", "name": n.strip(), "vramGB": round(float(mem) / 1024, 1)})
                except ValueError: pass
    if shutil.which("rocm-smi"):
        try:
            j = json.loads(run(["rocm-smi", "--showproductname", "--showmeminfo", "vram", "--json"]) or "{}")
            for k, v in j.items():
                if not k.startswith("card"): continue
                name = v.get("Card Series") or v.get("Card series") or v.get("Card Model") or "AMD GPU"
                tot = v.get("VRAM Total Memory (B)")
                out.append({"vendor": "amd", "name": name, "vramGB": round(int(tot) / 2**30, 1) if tot else None})
        except Exception:
            pass
    if sys.platform == "darwin":
        try:
            j = json.loads(run(["system_profiler", "SPDisplaysDataType", "-json"], 60) or "{}")
            for d in j.get("SPDisplaysDataType", []):
                out.append({"vendor": "apple", "name": d.get("sppci_model", "Apple GPU"), "vramGB": None, "cores": d.get("sppci_cores")})
        except Exception:
            pass
    return out


def detect():
    s = {"os": f"{platform.system()} {platform.release()} ({platform.machine()})", "cpu": cpu_name(), "ramGB": ram_gb(), "gpus": gpus(), "apu": None}
    c = (s["cpu"] or "").lower()
    if "ryzen ai max" in c:
        s["apu"] = s["cpu"]  # AMD Strix Halo: unified memory, GPU = Radeon 8060S / 8050S
        s["gpus"] = [g for g in s["gpus"] if g["vendor"] != "amd"] + [{"vendor": "amd", "name": "Radeon 8060S" if "395" in c else "Radeon 8050S", "vramGB": None, "unified": True}]
    return s


def get_json(url):
    req = urllib.request.Request(url, headers={"User-Agent": VERSION})
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode())


def norm(s):
    return re.sub(r"[^a-z0-9]+", " ", (s or "").lower()).strip()


def match_hw(sysinfo, hw):
    """Map detected hardware to a willitrunllm.com hardware id. Returns (id, gpuCount) or (None, n)."""
    ram = sysinfo.get("ramGB") or 0
    if sysinfo.get("apu"):
        c = sysinfo["apu"].lower(); plus = "395" in c
        os_ = "windows" if sys.platform == "win32" else "linux"
        cands = [h for h in hw if h["type"] == "unified-apu" and ("395" in h["name"]) == plus and os_ in h["name"].lower()]
        if cands:
            return min(cands, key=lambda h: abs(h["memGB"] - ram))["id"], 1
    if sys.platform == "darwin":
        chip = norm(sysinfo.get("cpu", "")).replace("apple ", "")
        cands = [h for h in hw if h["type"] == "mac" and re.sub(r" \d+ gb$", "", norm(h["name"]).replace("apple ", "")) == chip]
        if cands:
            return min(cands, key=lambda h: abs(h["memGB"] - ram))["id"], 1
    gs = [g for g in sysinfo["gpus"] if g["vendor"] in ("nvidia", "amd") and not g.get("unified")]
    if not gs:
        return None, 0
    g = gs[0]; gn = norm(g["name"]).replace("nvidia ", "").replace("geforce ", "").replace("amd ", "")
    best = None
    for h in hw:
        if h["type"] != "gpu": continue
        hn = norm(h["name"]).replace("nvidia ", "").replace("geforce ", "").replace("amd ", "").replace("intel ", "")
        base = re.sub(r" \d+gb$", "", hn)
        if base and (" " + base + " ") in (" " + gn + " "):
            score = len(base) - (abs(h["memGB"] - g["vramGB"]) if g.get("vramGB") else 0)
            if best is None or score > best[0]:
                best = (score, h["id"])
    same = sum(1 for x in gs if x["name"] == g["name"])
    return (best[1] if best else None), max(1, same)


# ---------------- benchmark runners ----------------
def find_llama_bench(flag):
    if flag:
        return flag if os.path.exists(flag) else None
    for n in ("llama-bench", "llama-bench.exe"):
        p = shutil.which(n)
        if p: return p
    for d in ("~/llama.cpp/build/bin", "~/llama.cpp", "./build/bin", "."):
        p = os.path.expanduser(os.path.join(d, "llama-bench"))
        if os.path.exists(p): return p
    return None


def resolve_model(a):
    if a.model:
        p = os.path.expanduser(a.model)
        if not os.path.exists(p): sys.exit(f"Model file not found: {p}")
        return p, None
    if a.hf_repo:
        if not a.hf_file: sys.exit("--hf-repo needs --hf-file (the exact .gguf file name)")
        try:
            from huggingface_hub import hf_hub_download
        except ImportError:
            sys.exit(f"huggingface_hub isn't installed. Either `pip install huggingface_hub` and re-run, or download it yourself:\n"
                     f"  https://huggingface.co/{a.hf_repo}/resolve/main/{a.hf_file}\nthen run with --model /path/{a.hf_file}")
        print(f"Downloading {a.hf_repo}/{a.hf_file} via huggingface_hub ...")
        return hf_hub_download(a.hf_repo, a.hf_file), f"{a.hf_repo}/{a.hf_file}"
    sys.exit("Pass --model /path/to/model.gguf, --hf-repo REPO --hf-file FILE, or --ollama MODEL. See --help.")


def parse_llama_bench_json(text):
    """llama-bench -o json prints a JSON array (sometimes after log lines)."""
    i = text.find("[")
    if i < 0: raise ValueError("no JSON array in llama-bench output")
    recs = json.loads(text[i:text.rindex("]") + 1])
    keep = ("model_type", "model_size", "model_n_params", "backends", "n_gpu_layers", "n_prompt", "n_gen", "avg_ts", "stddev_ts",
            "build_commit", "build_number", "gpu_info", "cpu_info", "n_threads", "flash_attn", "type_k", "type_v", "n_batch", "n_ubatch")
    out = [{k: r.get(k) for k in keep if k in r} for r in recs if isinstance(r, dict)]
    if not out: raise ValueError("llama-bench JSON had no results")
    return out


def run_llama_bench(lb, model, a):
    cmd = [lb, "-m", model, "-p", "512", "-n", "128", "-ngl", str(a.ngl), "-r", str(a.reps), "-o", "json"] + (a.extra or [])
    print("Running:", " ".join(cmd)); print("(this usually takes 1-5 minutes)")
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        sys.exit("llama-bench failed:\n" + (p.stderr or p.stdout)[-3000:])
    return parse_llama_bench_json(p.stdout)


def run_ollama(model):
    if not shutil.which("ollama"): sys.exit("ollama not found on PATH.")
    show = run(["ollama", "show", model], 60)
    q = re.search(r"quantization\s+(\S+)", show); par = re.search(r"parameters\s+(\S+)", show)
    prompt = "Write a detailed, 300-word explanation of how a CPU cache works."
    print(f"Running: ollama run {model} --verbose  (one generation)")
    p = subprocess.run(["ollama", "run", model, "--verbose", prompt], capture_output=True, text=True, timeout=900)
    return parse_ollama(p.stderr + p.stdout, model, q.group(1) if q else None, par.group(1) if par else None)


def parse_ollama(text, model, quant=None, params=None):
    ev = re.findall(r"^\s*eval rate:\s*([\d.]+)\s*tokens/s", text, re.M)
    pev = re.findall(r"prompt eval rate:\s*([\d.]+)\s*tokens/s", text)
    cnt = re.findall(r"^\s*eval count:\s*(\d+)", text, re.M)
    if not ev: raise ValueError("couldn't find 'eval rate' in ollama --verbose output")
    return {"model": model, "quantization": quant, "parameters": params, "evalRate": float(ev[-1]), "promptEvalRate": float(pev[-1]) if pev else None,
            "evalCount": int(cnt[-1]) if cnt else None, "note": "Ollama --verbose eval rate: lower fidelity than llama-bench (single run, prompt-dependent)."}


# ---------------- consent + submit ----------------
def ask(prompt):
    """Read the answer from the terminal even when the script itself came in on stdin (curl | python3 -)."""
    for dev in ("/dev/tty", "CON"):
        try:
            with open(dev, "r+") as t:
                t.write(prompt); t.flush()
                return t.readline().strip()
        except OSError:
            continue
    if sys.stdin.isatty():
        return input(prompt).strip()
    return ""


def table(rows):
    print("\n  {:<34} {:<10} {:<8} {:>12}".format("model", "backend", "test", "tok/s"))
    for r in rows:
        print("  {:<34} {:<10} {:<8} {:>12}".format(str(r.get("model_type", ""))[:34], str(r.get("backends", ""))[:10],
              ("pp%d" % r["n_prompt"]) if r.get("n_prompt") else ("tg%d" % r.get("n_gen", 0)), "%.2f" % r.get("avg_ts", 0)))


def main():
    ap = argparse.ArgumentParser(description="Will It Run? auto-benchmark: run llama-bench and (optionally) submit for review.")
    ap.add_argument("--model", help="path to a .gguf file")
    ap.add_argument("--hf-repo"); ap.add_argument("--hf-file")
    ap.add_argument("--ollama", metavar="MODEL", help="use `ollama run --verbose` instead (lower fidelity)")
    ap.add_argument("--llama-bench", help="path to llama-bench (default: search PATH)")
    ap.add_argument("--ngl", type=int, default=99, help="GPU layers (default 99 = all)")
    ap.add_argument("--reps", type=int, default=3)
    ap.add_argument("--hw", help="willitrunllm.com hardware id (skip auto-match), see /api/v1/hardware")
    ap.add_argument("--notes", default="", help="optional notes (driver, OS, build)")
    ap.add_argument("--sample-json", help="parse an existing `llama-bench -o json` output file instead of running")
    ap.add_argument("--dry-run", action="store_true", help="never submit; just show what would be sent")
    ap.add_argument("--offline", action="store_true", help="don't fetch the hardware list (implies --dry-run unless --hw given)")
    ap.add_argument("extra", nargs="*", help="extra llama-bench args after --, e.g. -- -fa 1")
    a = ap.parse_args()

    print(f"Will It Run? auto-benchmark ({VERSION}) - {SITE}\n")
    sysinfo = detect()
    print("Detected:"); print(f"  OS:  {sysinfo['os']}\n  CPU: {sysinfo['cpu']}\n  RAM: {sysinfo['ramGB']} GB")
    for g in sysinfo["gpus"] or [{"name": "none found", "vramGB": None}]:
        print(f"  GPU: {g['name']}" + (f" ({g['vramGB']} GB)" if g.get("vramGB") else ""))

    hw_id, n = a.hw, len([g for g in sysinfo["gpus"] if g.get("vendor") in ("nvidia", "amd") and not g.get("unified")]) or 1
    sysinfo["detectedHw"] = None
    if not a.offline:
        try:
            hw = get_json(SITE + "/api/v1/hardware")["hardware"]
            auto, n2 = match_hw(sysinfo, hw)
            sysinfo["detectedHw"] = auto
            if not hw_id: hw_id, n = auto, n2
        except Exception as e:
            print(f"  (couldn't fetch the hardware list: {e})")
    print(f"  Will It Run? hardware id: {hw_id or 'NOT MATCHED (use --hw <id>, see ' + SITE + '/api/v1/hardware)'}" + (f"  x{n} GPUs" if n > 1 else ""))

    hf_model = None
    if a.sample_json:
        bench = parse_llama_bench_json(open(a.sample_json).read()); kind = "llama-bench"
    elif a.ollama:
        oll = run_ollama(a.ollama); kind = "ollama"
    else:
        lb = find_llama_bench(a.llama_bench)
        if not lb:
            sys.exit("llama-bench not found. Install llama.cpp (https://github.com/ggml-org/llama.cpp/releases or `brew install llama.cpp`) "
                     "and put llama-bench on PATH, or pass --llama-bench /path/to/llama-bench.")
        model, hf_model = resolve_model(a)
        bench = run_llama_bench(lb, model, a); kind = "llama-bench"

    if kind == "llama-bench":
        print("\nResults (llama-bench):"); table(bench)
        if all(str(r.get("backends", "")).upper() in ("CPU", "BLAS", "") for r in bench) and not (hw_id or "").startswith("apple"):
            print("\n  WARNING: these results ran on the CPU only (no GPU backend). Use a CUDA/ROCm/Vulkan/Metal build of llama.cpp,\n  otherwise the numbers don't describe your GPU. Reviewers will reject CPU-only runs labelled as a GPU.")
    else:
        print(f"\nResults (Ollama, LOWER FIDELITY): eval rate {oll['evalRate']} tok/s, prompt eval {oll['promptEvalRate']} tok/s, quant {oll['quantization']}")

    payload = {"kind": kind, "gpu": hw_id, "gpuCount": n, "system": sysinfo, "hfModel": hf_model, "notes": a.notes[:1000], "client": VERSION, "t": START_MS}
    if kind == "llama-bench": payload["bench"] = bench
    else: payload["ollama"] = oll
    print("\n==== This is EXACTLY what would be sent to " + SITE + "/api/submit-bench ====")
    print(json.dumps(payload, indent=2))
    print("==== (no email, no file paths, no IP logging beyond normal web server logs) ====\n")
    if a.dry_run or (a.offline and not a.hw):
        print("Dry run: nothing was sent."); return
    if not hw_id:
        print("Can't submit without a hardware id. Re-run with --hw <id> (see " + SITE + "/api/v1/hardware)."); return
    if n > 1:
        print("Note: multi-GPU results are kept for review but aren't shown on the site yet.")
    ans = ask("Submit these results for manual review? [y/N] ")
    if ans.lower() not in ("y", "yes"):
        print("Not sent."); return
    if int(time.time() * 1000) - START_MS < 3500: time.sleep(3.5)  # server spam trap expects >3 s
    req = urllib.request.Request(SITE + "/api/submit-bench", data=json.dumps(payload).encode(), method="POST",
                                 headers={"Content-Type": "application/json", "User-Agent": VERSION})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            j = json.loads(r.read().decode())
    except urllib.error.HTTPError as e:
        j = json.loads(e.read().decode() or "{}")
    print("Thanks! Submitted (id %s), pending review." % j.get("id") if j.get("ok") else "Submission failed: %s" % j.get("error"))


if __name__ == "__main__":
    main()