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

Blinded packs are a list of hit-lists (one per query, same order as
search_queries.txt). Each hit is only {rank, url, title, snippet}.
"""

from __future__ import annotations

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

LABELS = ("A", "B", "C", "D")
CONFIGS = (
    "searxng",
    "searxng_brave",
    "searxng_brave_tavily",
    "hermes_tavily",
)
HIT_KEYS = ("rank", "url", "title", "snippet")
LEAK_RE = re.compile(r"searxng|brave|tavily|hermes", re.I)
EXEMPT_FIELDS = frozenset({"url", "snippet"})


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 _hit_list(block: object) -> list:
    if isinstance(block, list):
        if block and all(isinstance(x, dict) and ("url" in x or "title" in x) for x in block):
            return block
        return []
    if not isinstance(block, dict):
        return []
    for key in ("results", "web", "hits"):
        v = block.get(key)
        if isinstance(v, list):
            return v
    data = block.get("data")
    if isinstance(data, list):
        return data
    if isinstance(data, dict):
        for key in ("web", "results", "hits"):
            v = data.get(key)
            if isinstance(v, list):
                return v
    return []


def _items(raw: object) -> list:
    if isinstance(raw, list):
        return raw
    if isinstance(raw, dict):
        for key in ("queries", "results"):
            v = raw.get(key)
            if isinstance(v, list) and v and isinstance(v[0], (dict, list)):
                if key == "results" and v and isinstance(v[0], dict) and "url" in v[0]:
                    return [raw]
                return v
        return [raw]
    raise ValueError(f"raw dump must be a list or object, got {type(raw).__name__}")


def normalize(raw: object, n: int) -> list[list[dict]]:
    items = _items(raw)
    if len(items) != n:
        raise ValueError(f"expected {n} queries in dump, got {len(items)}")
    out: list[list[dict]] = []
    for item in items:
        hits = []
        if isinstance(item, list):
            raw_hits = item
        else:
            raw_hits = _hit_list(item)
        for i, h in enumerate(raw_hits, 1):
            if not isinstance(h, dict):
                continue
            rank = h.get("rank", h.get("position", i))
            try:
                rank_i = int(rank)
            except (TypeError, ValueError):
                rank_i = i
            hits.append(
                {
                    "rank": rank_i,
                    "url": str(h.get("url") or ""),
                    "title": str(h.get("title") or ""),
                    "snippet": str(
                        h.get("snippet") or h.get("description") or h.get("content") or ""
                    ),
                }
            )
        out.append(hits)
    return out


def leak_hits(obj: object, path: str = "$"):
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k in EXEMPT_FIELDS and isinstance(v, str):
                continue
            yield from leak_hits(v, f"{path}.{k}")
    elif isinstance(obj, list):
        for i, v in enumerate(obj):
            yield from leak_hits(v, f"{path}[{i}]")
    elif isinstance(obj, str) and LEAK_RE.search(obj):
        yield path, obj


def assert_blind(pack: object) -> None:
    if not isinstance(pack, list):
        raise ValueError("blind pack must be a list of hit lists")
    for qi, hits in enumerate(pack):
        if not isinstance(hits, list):
            raise ValueError(f"query {qi}: hits must be a list")
        for hi, h in enumerate(hits):
            if not isinstance(h, dict) or set(h) != set(HIT_KEYS):
                raise ValueError(f"query {qi} hit {hi}: keys must be {HIT_KEYS}, got {h!r}")
            if not isinstance(h["rank"], int):
                raise ValueError(f"query {qi} hit {hi}: rank must be int")
            for k in ("url", "title", "snippet"):
                if not isinstance(h[k], str):
                    raise ValueError(f"query {qi} hit {hi}: {k} must be str")
    found = list(leak_hits(pack))
    if found:
        bits = "; ".join(f"{p}={v!r}" for p, v in found[:8])
        raise ValueError(f"vendor token outside url/snippet: {bits}")


def write_packs(packs: dict[str, list], out_dir: Path, n: int, seed: int | None) -> Path:
    rng = random.Random(seed)
    order = list(CONFIGS)
    rng.shuffle(order)
    key = {lab: cfg for lab, cfg in zip(LABELS, order)}
    out_dir.mkdir(parents=True, exist_ok=True)
    for lab, cfg in key.items():
        pack = packs[cfg]
        assert_blind(pack)
        (out_dir / f"{lab}.json").write_text(json.dumps(pack, indent=2) + "\n")
    key_path = out_dir / "key.json"
    key_path.write_text(json.dumps({"labels": key, "n": n}, indent=2) + "\n")
    return key_path


def self_test() -> int:
    n = 2
    sift_q = {
        "query": "example docs",
        "results": [
            {
                "url": "https://docs.example.com/search",
                "title": "Search method",
                "snippet": "POST /v2/search",
                "score": 0.91,
                "source": "searxng",
            }
        ],
        "warnings": ["tavily: 429"],
        "timings": {"search_ms": 40},
        "credits": 1,
        "config": "searxng",
        "sources": ["searxng"],
    }
    hermes_q = {
        "query": "example docs",
        "success": True,
        "data": {
            "web": [
                {
                    "url": "https://api.tavily.com/docs",
                    "title": "Vendor docs",
                    "description": "search_depth basic",
                    "position": 1,
                    "score": 0.4,
                }
            ]
        },
        "warnings": ["brave: timeout"],
    }
    dumps = {
        "searxng": [sift_q, sift_q],
        "searxng_brave": {
            "config": "searxng_brave",
            "queries": [sift_q, sift_q],
        },
        "searxng_brave_tavily": [sift_q, sift_q],
        "hermes_tavily": [hermes_q, hermes_q],
    }
    with tempfile.TemporaryDirectory() as tmp:
        raw = Path(tmp) / "raw"
        out = Path(tmp) / "out"
        raw.mkdir()
        for name, payload in dumps.items():
            (raw / f"{name}.json").write_text(json.dumps(payload) + "\n")
        packs = {name: normalize(payload, n) for name, payload in dumps.items()}
        for pack in packs.values():
            assert_blind(pack)
        key_path = write_packs(packs, out, n, seed=1)
        for lab in LABELS:
            blob = json.loads((out / f"{lab}.json").read_text())
            assert_blind(blob)
            assert len(blob) == n
            assert blob[0][0]["url"].startswith("https://")
            keys = set()

            def walk(o):
                if isinstance(o, dict):
                    keys.update(o)
                    for v in o.values():
                        walk(v)
                elif isinstance(o, list):
                    for v in o:
                        walk(v)

            walk(blob)
            if keys != {"rank", "url", "title", "snippet"}:
                raise AssertionError(f"{lab} keys {keys}")
        # snippet/url may name vendors; title must not
        bad = [[{"rank": 1, "url": "https://x.example", "title": "tavily guide", "snippet": "ok"}]]
        try:
            assert_blind(bad)
        except ValueError as e:
            if "vendor token" not in str(e):
                raise
        else:
            raise AssertionError("title leak not caught")
        ok_snip = [
            [{"rank": 1, "url": "https://api.tavily.com/x", "title": "Docs", "snippet": "uses tavily"}]
        ]
        assert_blind(ok_snip)
        print("self-test ok", key_path)
    return 0


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--self-test", action="store_true")
    p.add_argument(
        "--raw-dir",
        type=Path,
        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()
    if args.self_test:
        return self_test()
    if args.raw_dir is None:
        p.error("--raw-dir is required unless --self-test")

    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] = normalize(json.loads(f.read_text()), len(queries))

    try:
        key_path = write_packs(packs, args.out_dir, len(queries), args.seed)
    except ValueError as e:
        print(e, file=sys.stderr)
        return 1
    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())