#!/usr/bin/env python3
"""Independent Sidekick Cortex Pro (formerly Axon) benchmark runner.
Uses Ollama /api/chat with the dataset's persona system message (temp 0).

Usage:
  python tests/run_benchmark.py lastbrowser/sidekick-cortex-pro:latest \\
      data/axon_handwritten_testset_50.jsonl outputs_hw50.jsonl
  python tests/score.py outputs_hw50.jsonl
"""
import json, time, sys, requests

def main():
    if len(sys.argv) < 4:
        print(__doc__)
        sys.exit(2)
    model, src, out = sys.argv[1], sys.argv[2], sys.argv[3]
    base = sys.argv[4] if len(sys.argv) > 4 else "http://localhost:11434"
    samples = [json.loads(l) for l in open(src, encoding="utf-8") if l.strip()]
    with open(out, "w", encoding="utf-8") as fo:
        for i, s in enumerate(samples, 1):
            msgs = [m for m in s["messages"] if m["role"] != "assistant"]
            t0 = time.time()
            r = requests.post(
                f"{base}/api/chat",
                json={
                    "model": model,
                    "messages": msgs,
                    "stream": False,
                    "options": {"temperature": 0, "num_ctx": 4096, "num_predict": 256},
                },
                timeout=300,
            ).json()
            exp = next(m["content"] for m in s["messages"] if m["role"] == "assistant")
            fo.write(
                json.dumps(
                    {
                        "id": s.get("id", f"sample_{i}"),
                        "prompt": msgs[-1]["content"],
                        "expected": exp,
                        "raw": r.get("message", {}).get("content", ""),
                        "eval_count": r.get("eval_count"),
                        "eval_duration_ns": r.get("eval_duration"),
                        "wall_s": time.time() - t0,
                    },
                    ensure_ascii=False,
                )
                + "\n"
            )
            fo.flush()
            print(f"[{i}/{len(samples)}] {s.get('id', i)} done", flush=True)

if __name__ == "__main__":
    main()