#!/usr/bin/env python3
"""Shuffle four named result dumps into A–D and write a sealed key file."""

from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

LABELS = ("A", "B", "C", "D")
CONFIGS = (
    "searxng",
    "searxng_brave",
    "searxng_brave_tavily",
    "hermes_tavily",
)


def load_queries(path: Path) -> list[str]:
    return [ln.strip() for ln in path.read_text().splitlines() if ln.strip()]


def load_targets(path: Path) -> list[str]:
    out: list[str] = []
    for ln in path.read_text().splitlines():
        s = ln.strip()
        if not s or s.startswith("#"):
            continue
        out.append(s)
    return out


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument(
        "--raw-dir",
        type=Path,
        required=True,
        help="Directory with {searxng,searxng_brave,searxng_brave_tavily,hermes_tavily}.json",
    )
    p.add_argument(
        "--out-dir",
        type=Path,
        default=Path("tests/testdata/bakeoff-out"),
    )
    p.add_argument(
        "--queries",
        type=Path,
        default=Path("tests/testdata/search_queries.txt"),
    )
    p.add_argument(
        "--targets",
        type=Path,
        default=Path("tests/testdata/search_targets.txt"),
    )
    p.add_argument("--seed", type=int, default=None)
    args = p.parse_args()

    queries = load_queries(args.queries)
    targets = load_targets(args.targets)
    if len(queries) != len(targets):
        print(
            f"query/target length mismatch: {len(queries)} vs {len(targets)}",
            file=sys.stderr,
        )
        return 1

    packs = {}
    for name in CONFIGS:
        f = args.raw_dir / f"{name}.json"
        packs[name] = json.loads(f.read_text())

    rng = random.Random(args.seed)
    order = list(CONFIGS)
    rng.shuffle(order)
    key = {lab: cfg for lab, cfg in zip(LABELS, order)}

    args.out_dir.mkdir(parents=True, exist_ok=True)
    for lab, cfg in key.items():
        payload = packs[cfg]
        if isinstance(payload, dict):
            payload = {k: v for k, v in payload.items() if k not in ("config", "sources")}
        (args.out_dir / f"{lab}.json").write_text(json.dumps(payload, indent=2) + "\n")

    key_path = args.out_dir / "key.json"
    key_path.write_text(json.dumps({"labels": key, "n": len(queries)}, indent=2) + "\n")
    print(f"wrote {args.out_dir}/{{A,B,C,D}}.json")
    print(f"key sealed at {key_path} — judges do not open this until scoring is done")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())