"""Run pinned JevBench public inputs through Instinct, using the upstream scorer.

uv run --no-project --with httpx==0.28.1 python benchmark.py \
  --jevbench ../jevbench --output run/jevbench/serial --key-file /path/to/key

No retries or truncation. Ground truth is used only by the local scorer.
Each output directory is exclusive; completed requests are durably recorded.
"""
import argparse
import asyncio
from collections import Counter
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import secrets
import subprocess
import sys
import time

import httpx


def now():
    return datetime.now(timezone.utc).isoformat()


def decision_request(task):
    q = {k: task.question[k] for k in ("type", "instructions", "criteria") if k in task.question}
    if q["type"] == "noul" and q.get("criteria"):
        # Autograd names the same binary descriptions yes/no. Preserve their text.
        q["criteria"] = {{"true": "yes", "false": "no"}.get(k, k): v for k, v in q["criteria"].items()}
    return {"model": "jev-lite", "messages": [{"role": "user", "content": json.dumps(
        {"state": task.state, "questions": {"decision": q}}, ensure_ascii=False)}], "stream": False}


def answer_probabilities(answer, task):
    if answer["type"] != task.question["type"]:
        raise ValueError("Native answer type mismatch")
    if answer["type"] == "noul":
        p = answer["noul"]
        if isinstance(p, bool) or not isinstance(p, (float, int)) or not 0 <= p <= 1:
            raise ValueError("Invalid boolean probability")
        return {"yes": p, "no": 1 - p}
    if answer["type"] == "choice" and answer["choice"] not in task.labels:
        raise ValueError("Invalid native choice")
    return answer["probabilities"]


async def run(args):
    expected_commit = "a1799db9673bf59010bc698753e729b04132431f"
    actual_commit = subprocess.check_output(["git", "-C", str(args.jevbench), "rev-parse", "HEAD"], text=True).strip()
    if actual_commit != expected_commit:
        raise ValueError("Use the documented JevBench commit: " + expected_commit)
    sys.path.insert(0, str(args.jevbench.resolve()))
    from jevbench.tasks import load_jsonl, dataset_hash
    from jevbench.scoring import score_task
    from jevbench.summarize import metric
    from jevbench.metrics import latency_summary

    key = args.key_file.read_text().strip() if args.key_file else os.environ["INSTINCT_API_KEY"]
    if not key:
        raise ValueError("Missing API key")
    tasks, tiers, hashes = [], {}, {}
    for tier in ("easy", "original", "hard"):
        path = args.jevbench / "datasets/public" / (tier + ".jsonl")
        rows = load_jsonl(path)
        hashes[tier] = hashlib.sha256(path.read_bytes()).hexdigest()
        tasks.extend(rows)
        tiers.update({t.id: tier for t in rows})
    if len({t.id for t in tasks}) != len(tasks):
        raise ValueError("Duplicate task IDs")
    if args.task_id:
        tasks = [t for t in tasks if t.id in args.task_id]
        if len(tasks) != len(set(args.task_id)):
            raise ValueError("Unknown task ID")
    planned = [tasks[i % len(tasks)] for i in range(args.count or len(tasks))]
    output = args.output.resolve()
    output.mkdir(parents=True, exist_ok=False)
    (output / "raw").mkdir()
    comparison = {}
    if args.compare:
        comparison = {r["task_id"]: r for r in map(json.loads, args.compare.read_text().splitlines())}
    manifest = {
        "started_utc": now(), "endpoint": args.endpoint, "model": "jev-lite",
        "benchmark_commit": subprocess.check_output(["git", "-C", str(args.jevbench), "rev-parse", "HEAD"], text=True).strip(),
        "driver_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "driver_origin_commit": "7468fb73a1789a5bca628151ee2079ede351299a",
        "dataset_sha256": hashes, "dataset_hash": dataset_hash(tasks),
        "n_unique_tasks": len(tasks), "n_planned": len(planned), "concurrency": args.concurrency,
        "timeout_s": args.timeout, "retries": 0, "truncation": False,
        "mapping": "one state/questions decision per request; noul true/false criteria renamed yes/no; choice option insertion order and score level order preserved",
        "probabilities": "native candidate distribution; noul mapped to yes=p and no=1-p",
        "scope": "public subset only; no official leaderboard composite or cost estimate",
        "repeat_note": "Repeated public inputs can reuse SGLang prefix cache; load phases are not independent quality measurements.",
        "measurement": "public HTTPS round trip from this client, pooled connections; no synthetic latency adjustment",
    }
    (output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    records, errors = [], 0
    pending = iter(enumerate(planned))
    stop = asyncio.Event()
    started = time.perf_counter()
    active = peak = 0
    headers = {"Authorization": "Bearer " + key, "Content-Type": "application/json", "User-Agent": "Instinct-JevBench/1.0"}
    limits = httpx.Limits(max_connections=args.concurrency, max_keepalive_connections=args.concurrency)
    with (output / "results.jsonl").open("x") as stream:
        async with httpx.AsyncClient(headers=headers, timeout=args.timeout, limits=limits, trust_env=False) as client:
            async def worker():
                nonlocal errors, active, peak
                while not stop.is_set():
                    item = next(pending, None)
                    if item is None:
                        return
                    index, task = item
                    body = decision_request(task)
                    trace_id = secrets.token_hex(16)
                    row = {"index": index, "task_id": task.id, "tier": tiers[task.id], "family": task.family,
                           "ts": time.time(), "started_utc": now(), "trace_id": trace_id, "model": "jev-lite",
                           "ok": False, "status": "failed", "valid": False, "strict_valid": False,
                           "renormalized": False, "correct": False, "predicted": None, "probs": None,
                           "probs_source": "native", "status_code": None, "error": None, "usage": {},
                           "cost_usd": None, "cost_basis": "self_hosted_compute_unpriced"}
                    raw = None
                    active += 1
                    peak = max(peak, active)
                    tick = time.perf_counter()
                    try:
                        response = await client.post(args.endpoint, json=body, headers={
                            "traceparent": f"00-{trace_id}-{secrets.token_hex(8)}-01"})
                        row["latency_s"] = time.perf_counter() - tick
                        row["status_code"] = response.status_code
                        try:
                            raw = response.json()
                        except ValueError:
                            raw = response.text
                        if response.status_code != 200:
                            raise ValueError(f"HTTP {response.status_code}: {str(raw)[:600]}")
                        content = json.loads(raw["choices"][0]["message"]["content"])
                        probs = answer_probabilities(content["answers"]["decision"], task)
                        scored = score_task(probs, task)
                        row.update(scored, ok=True, status="ok", probs_as_returned=probs,
                                   usage=raw.get("usage", {}), request_id=raw.get("id"),
                                   inference_seconds=content.get("inference_seconds"))
                        if not scored["valid"]:
                            row.update(ok=False, status="failed")
                        reference = comparison.get(task.id)
                        if reference and row["valid"] and reference.get("valid"):
                            row["baseline_decision_matches"] = row["predicted"] == reference["predicted"]
                            row["baseline_max_probability_delta"] = max(abs(row["probs"][k] - reference["probs"][k]) for k in task.labels)
                    except (httpx.HTTPError, ValueError, KeyError, TypeError, IndexError) as exc:
                        row["error"] = str(exc).replace(key, "[REDACTED]")[:700]
                    finally:
                        row.setdefault("latency_s", time.perf_counter() - tick)
                        active -= 1
                    raw_bytes = json.dumps({"request": body, "response": raw, "http_status": row["status_code"]}, ensure_ascii=False, allow_nan=False).replace(key, "[REDACTED]").encode()
                    (output / "raw" / f"{index:06d}.json").write_bytes(raw_bytes)
                    row["raw_sha256"] = hashlib.sha256(raw_bytes).hexdigest()
                    records.append(row)
                    stream.write(json.dumps(row, allow_nan=False) + "\n")
                    stream.flush()
                    os.fsync(stream.fileno())
                    infrastructure_error = not row["ok"] and (row["status_code"] is None or row["status_code"] >= 500)
                    errors = errors + 1 if infrastructure_error else 0
                    if row["status_code"] in (401, 403, 429) or errors >= 3:
                        stop.set()
                    if len(records) % 25 == 0 or not row["ok"]:
                        print(json.dumps({"completed": len(records), "planned": len(planned), "failed": sum(not r["ok"] for r in records), "last_id": task.id, "last_error": row["error"]}), flush=True)
            await asyncio.gather(*(worker() for _ in range(args.concurrency)))
    elapsed = time.perf_counter() - started
    summary = metric(planned, records)
    summary.update(elapsed_s=elapsed, concurrency=args.concurrency, client_peak_inflight=peak,
                   attempted_requests_per_second=len(records) / elapsed,
                   successful_requests_per_second=sum(r["ok"] for r in records) / elapsed,
                   complete=len(records) == len(planned), all_requests_succeeded=all(r["ok"] for r in records),
                   http_status_counts=dict(Counter(str(r["status_code"]) for r in records)),
                   successful_latency=latency_summary([r["latency_s"] for r in records if r["ok"]]),
                   input_tokens=sum(r["usage"].get("prompt_tokens", 0) for r in records),
                   per_tier={tier: metric([t for t in planned if tiers[t.id] == tier], records) for tier in sorted(set(tiers.values()))},
                   per_family={family: metric([t for t in planned if t.family == family], records) for family in sorted({t.family for t in planned})})
    compared = [r for r in records if "baseline_decision_matches" in r]
    if compared:
        summary["baseline_comparison"] = {"n": len(compared), "decision_mismatches": sum(not r["baseline_decision_matches"] for r in compared),
                                          "max_probability_delta": max(r["baseline_max_probability_delta"] for r in compared)}
    (output / "summary.json").write_text(json.dumps(summary, indent=2, allow_nan=False) + "\n")
    manifest.update(finished_utc=now(), n_attempted=len(records))
    (output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps({"output": str(output), **{k: summary[k] for k in ("n_attempted", "n_correct", "accuracy", "complete", "all_requests_succeeded", "elapsed_s", "successful_requests_per_second", "successful_latency", "http_status_counts")}}), flush=True)
    return 0 if summary["complete"] and summary["all_requests_succeeded"] else 1


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--jevbench", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--key-file", type=Path)
    parser.add_argument("--endpoint", default="https://instinct.zoowork.ai/v1/chat/completions")
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument("--count", type=int)
    parser.add_argument("--task-id", action="append")
    parser.add_argument("--timeout", type=float, default=150)
    parser.add_argument("--compare", type=Path)
    args = parser.parse_args()
    if args.concurrency < 1 or args.count is not None and args.count < 1:
        parser.error("concurrency and count must be positive")
    return asyncio.run(run(args))


if __name__ == "__main__":
    sys.exit(main())