"""Pluggable eval sets for the head-to-head bench.

A set yields records ``{"uid", "context", "schema", "value"}`` (``value`` = gold JSON).
Names:

* ``datagen-val`` / ``datagen-test``: the datagen ``val`` / ``test`` (held-out domains) splits
  of :data:`deg.train.DEG_RECORDS` (or ``--records``);
* ``extractbench``: ``/data/deg/extractbench/records.jsonl`` (the ExtractBench loader writes
  records in this same shape; human gold, eval only);
* ``jevbench``: the 231 public JevBench decisions as one-question decide records
  (:mod:`deg.bench.jevbench`; eval only — its own runner scores them by JevBench's rules);
* ``jsonl:<path>[@<split>]``: any JSONL file of records in this shape.

New sets register with :func:`register`. The sample is fixed by a manifest: the first run
writes ``<out>/sets/<name>.uids`` (all records then present, ordered by sha1(uid)); every
run takes the first ``n`` of it, so every system scores the same records even while
datagen keeps appending.
"""

from __future__ import annotations

import hashlib
import json
from collections.abc import Callable
from pathlib import Path
from typing import Any

Record = dict[str, Any]
Loader = Callable[[], list[Record]]

EXTRACTBENCH = Path("/data/deg/extractbench/records.jsonl")
_REGISTRY: dict[str, Callable[[dict[str, Any]], Loader]] = {}


def register(name: str, factory: Callable[[dict[str, Any]], Loader]) -> None:
    """``factory(options) -> loader``; options carry CLI paths such as ``records``."""
    _REGISTRY[name] = factory


def read_records(path: Path, split: str | None) -> list[Record]:
    out = []
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            if not line.strip():
                continue
            rec = json.loads(line)
            if split is None or rec.get("split") == split:
                out.append({"uid": rec["uid"], "context": rec["context"],
                            "schema": rec["schema"], "value": rec["value"],
                            "domain": rec.get("domain")})
    return out


def _datagen(split: str) -> Callable[[dict[str, Any]], Loader]:
    def factory(opts: dict[str, Any]) -> Loader:
        from deg.train import DEG_RECORDS

        path = Path(opts.get("records") or DEG_RECORDS)
        return lambda: read_records(path, split)
    return factory


register("datagen-val", _datagen("val"))
register("datagen-test", _datagen("test"))
register("extractbench", lambda opts: lambda: read_records(
    Path(opts.get("extractbench") or EXTRACTBENCH), None))


def _jevbench(opts: dict[str, Any]) -> Loader:
    from deg.bench import jevbench

    return lambda: jevbench.records(Path(opts.get("jevbench") or jevbench.REPO))


register("jevbench", _jevbench)


def loader(name: str, opts: dict[str, Any] | None = None) -> Loader:
    opts = opts or {}
    if name.startswith("jsonl:"):
        spec = name[len("jsonl:"):]
        path, _, split = spec.partition("@")
        return lambda: read_records(Path(path), split or None)
    if name not in _REGISTRY:
        raise KeyError(f"unknown set {name!r}; known: {sorted(_REGISTRY)} or jsonl:<path>")
    return _REGISTRY[name](opts)


def safe_name(name: str) -> str:
    return "".join(c if c.isalnum() or c in "-_." else "_" for c in name)


def sample(name: str, n: int | None, out: Path, opts: dict[str, Any] | None = None
           ) -> list[Record]:
    """The fixed sample for ``name``: manifest uids if present, else first ``n`` by sha1."""
    recs = loader(name, opts)()
    by_uid = {r["uid"]: r for r in recs}
    manifest = out / "sets" / f"{safe_name(name)}.uids"
    if manifest.exists():
        uids = [u for u in manifest.read_text().split() if u]
        missing = [u for u in uids if u not in by_uid]
        if missing:
            raise RuntimeError(f"{manifest}: {len(missing)} uids no longer in the set")
    else:
        uids = sorted(by_uid, key=lambda u: hashlib.sha1(u.encode()).hexdigest())
        manifest.parent.mkdir(parents=True, exist_ok=True)
        manifest.write_text("\n".join(uids) + "\n")
    picked = [by_uid[u] for u in uids]
    return picked[:n] if n else picked