#!/usr/bin/env python3
"""Ballot test runner: one contest x one arm -> runs/<arm>/<contest>.md + .json (meta).

Arms (pre-registered in ../TEST-PLAN.md):
  A  Hermes, Qwen3.8-27B dense (local MLX), live web search, thinking OFF (deviation 1, 10/28 10:25 AM ET)
  B  Hermes, Qwen3-30B-A3B MoE (local MLX), live web search
  C  Copilot CLI, Claude Opus 5.5, live web search
  D  Copilot CLI, GPT-6.1-sol, live web search
  E  Hermes, Qwen3.8-27B dense, NO tools, fixed source packet (packets/<contest>.md)
"""
import argparse, datetime as dt, json, os, subprocess, sys, time
from pathlib import Path

HERE = Path(__file__).resolve().parent
HERMES = os.path.expanduser("~/.local/bin/hermes")
ARMS = {
    "A": {"harness": "hermes", "model": "mlx-community/Qwen3.8-27B-4bit",
          "provider": "custom:Qwen3.8-27B Local MLX No-Think", "tools": "web", "timeout": 3600},
    "B": {"harness": "hermes", "model": "mlx-community/Qwen3-30B-A3B-4bit",
          "provider": "custom:Qwen3-30B-A3B MoE Local MLX", "tools": "web", "timeout": 2400},
    "C": {"harness": "copilot", "model": "claude-opus-5.5", "timeout": 1800},
    "D": {"harness": "copilot", "model": "gpt-6.1-sol", "timeout": 1800},
    "E": {"harness": "hermes", "model": "mlx-community/Qwen3.8-27B-4bit",
          "provider": "custom:Qwen3.8-27B Local MLX No-Think", "tools": "", "timeout": 3600},
}


def contest_block(c: dict) -> str:
    lines = [f"Title: {c['title']}", f"Who votes on this: {c['scope']}"]
    if c["kind"] == "question":
        lines += ["Type: BALLOT QUESTION (use the ballot-question format)",
                  f"Official ballot text: {c['official_text']}"]
        if c.get("official_summary"):
            lines.append(f"Official State summary: {c['official_summary']}")
    else:
        lines += ["Type: RACE (use the race format)", f"Vote for up to: {c['vote_for']}",
                  "Candidates in ballot order:"]
        lines += [f"  - {x['name']} ({x['party']})" for x in c["candidates"]]
        if c.get("note"):
            lines.append(f"Note: {c['note']}")
    return "\n".join(lines)


def build_prompt(c: dict, arm: str) -> str:
    tpl = (HERE / "PROMPT.md").read_text()
    p = (tpl.replace("{today}", dt.date.today().strftime("%A, %B %-d, %Y"))
            .replace("{contest_block}", contest_block(c)).replace("{title}", c["title"]))
    if arm == "E":
        packet = (HERE / "packets" / f"{c['id']}.md").read_text()
        p = p.replace("1. RESEARCH FIRST.", "1. NO TOOLS: you have no web access. Use ONLY the source "
                      "packet at the end of this prompt and cite its numbered sources. IGNORE: ")
        p += "\n\n## Source packet (the only sources you may cite)\n" + packet
    return p


def run(arm: str, c: dict, out_dir: Path) -> dict:
    cfg = ARMS[arm]
    prompt = build_prompt(c, arm)
    out_md = out_dir / f"{c['id']}.md"
    env = dict(os.environ)
    if cfg["harness"] == "hermes":
        usage = out_dir / f"{c['id']}.usage.json"
        cmd = [HERMES, "--ignore-rules", "-m", cfg["model"], "--provider", cfg["provider"],
               "--usage-file", str(usage)]
        cmd += ["-t", cfg["tools"]] if cfg["tools"] else ["-t", "clarify"]
        cmd += ["-z", prompt]
    else:
        usage = out_dir / f"{c['id']}.usage.json"
        cmd = ["copilot", "-p", prompt, "--model", cfg["model"], "--no-custom-instructions",
               "--disable-builtin-mcps", "--available-tools", "web_search", "web_fetch",
               "--allow-all-tools", "--allow-all-urls", "-s", "--no-color",
               "--usage-output-file", str(usage)]
        env.update({"OTEL_EXPORTER_OTLP_ENDPOINT": "http://192.168.68.58:4318",
                    "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf",
                    "OTEL_RESOURCE_ATTRIBUTES": "openinference.project.name=ballot-2026,"
                    f"agent.name=ballot-arm-{arm},deployment.environment=lab,host.name=macbook-pro"})
    t0 = time.time()
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=cfg["timeout"],
                           cwd="/tmp", env=env)
        rc, out, err = p.returncode, p.stdout, p.stderr
    except subprocess.TimeoutExpired as e:
        rc, out, err = -9, (e.stdout or b"").decode() if isinstance(e.stdout, bytes) else (e.stdout or ""), "TIMEOUT"
    wall = round(time.time() - t0, 1)
    text = out.strip()
    if "# " in text and not text.startswith("# "):
        text = text[text.index("# "):]
    out_md.write_text(text + "\n")
    meta = {"arm": arm, "contest": c["id"], "model": cfg["model"], "harness": cfg["harness"],
            "wall_seconds": wall, "exit_code": rc, "words": len(text.split()),
            "started": dt.datetime.fromtimestamp(t0).isoformat(timespec="seconds"),
            "stderr_tail": err[-1500:] if err else ""}
    (out_dir / f"{c['id']}.json").write_text(json.dumps(meta, indent=1))
    return meta


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=list(ARMS))
    ap.add_argument("--contests", default="all", help="comma list of ids or 'all'")
    ap.add_argument("--skip-done", action="store_true")
    a = ap.parse_args()
    data = json.loads((HERE / "contests.json").read_text())
    ids = None if a.contests == "all" else set(a.contests.split(","))
    out_dir = HERE.parent / "runs" / a.arm
    out_dir.mkdir(parents=True, exist_ok=True)
    for c in data["contests"]:
        if ids and c["id"] not in ids:
            continue
        if a.skip_done and (out_dir / f"{c['id']}.json").exists():
            if json.loads((out_dir / f"{c['id']}.json").read_text()).get("exit_code") == 0:
                continue
        m = run(a.arm, c, out_dir)
        print(f"{dt.datetime.now():%H:%M} arm {a.arm} {c['id']}: rc={m['exit_code']} "
              f"{m['wall_seconds']}s {m['words']} words", flush=True)


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