#!/usr/bin/env python3
"""Nex-N2.5-mini speed + gates. Every 4-bit file x backend, no draft head (the checkpoint ships no MTP weights),
same house protocol as Agnes (nex_harness.py). Then gates on STRIX_LEAN: aligned prompt-cache identity, tool calls
(thinking on + off), vision with -fa on and off. Sequential GPU; results/nex_bench.jsonl."""
import fcntl, os, subprocess, time
W = "/mnt/models/nex-n2.5-mini"; N = "Nex-N2.5-mini"; A = "/mnt/models/agnes-3.0-flash"
J = f"{W}/results/nex_bench.jsonl"
OB = "/opt/llama-rocm/rocmfpx-724/build-hipvk/bin"
T = {"q106": f"{W}/out/{N}-Q4_0_ROCMFP4_STRIX_LEAN.gguf", "q102": f"{W}/out/{N}-Q4_0_ROCMFP4_COHERENT.gguf",
     "q103": f"{W}/out/{N}-Q4_0_ROCMFP4_FAST.gguf",
     "q106i": f"{W}/out-imat/{N}-imatrix-Q4_0_ROCMFP4_STRIX_LEAN.gguf",
     "q102i": f"{W}/out-imat/{N}-imatrix-Q4_0_ROCMFP4_COHERENT.gguf",
     "q103i": f"{W}/out-imat/{N}-imatrix-Q4_0_ROCMFP4_FAST.gguf"}
ENV = dict(os.environ, AGNES_BIN=OB)


def log(m): print(f"[{time.strftime('%FT%TZ', time.gmtime())}] {m}", flush=True)


def run(argv, label):
    log(label)
    r = subprocess.run(argv, capture_output=True, text=True, env=ENV)
    print(r.stdout.strip()[-800:], flush=True)
    if r.returncode != 0:
        print("  RC", r.returncode, r.stderr.strip()[-600:], flush=True)
    return r.returncode


def harness(cmd, label, model, dev="ROCm0", reps=3, workload="code", extra=()):
    return run(["python3", f"{W}/nex_harness.py", cmd, "--model", model, "--dev", dev, "--label", label,
                "--reps", str(reps), "--jsonl", J, "--serverlog", f"{W}/logs/b_{label}.log", "--workload", workload,
                *extra], f"{cmd} {label}")


# One quiet box: publish_v2.sh upload takes the same lock, so an HF upload never overlaps the speed runs.
QUIET = open("/mnt/models/.quiet-box.lock", "w")
log("waiting for the quiet-box lock (no HF upload during speed runs)")
fcntl.flock(QUIET, fcntl.LOCK_EX)
log("quiet-box lock held")
for tag in ("q106", "q102", "q103", "q106i", "q102i", "q103i"):
    for dev, sh in (("ROCm0", "rocm"), ("Vulkan0", "vk")):
        harness("bench", f"n-{tag}-{sh}", T[tag], dev=dev)
harness("bench", "n-q106-rocm-prose", T["q106"], workload="prose")
harness("bench", "n-q106-vk-prose", T["q106"], dev="Vulkan0", workload="prose")
run(["python3", f"{W}/nex_cachegate.py", "--model", T["q106"], "--label", "n-c3-q106", "--reps", "5",
     "--jsonl", J, "--serverlog", f"{W}/logs/b_n-c3-q106.log"], "cachegate n-c3-q106")
harness("tools", "n-tools-q106", T["q106"])
for fa in ("on", "off"):
    harness("vision", f"n-vision-q106-fa{fa}", T["q106"],
            extra=("--fa", fa, "--mmproj", f"{W}/out/mmproj-{N}-BF16.gguf", "--image", f"{A}/calib/vision_probe.png",
                   "--question", "Describe the two shapes in this image and their colors.",
                   "--expect", "red,blue,circle,square"))
log("NEX_BENCH_DONE")