#!/usr/bin/env python3
"""Build an annotation queue from rafmacalaba/datause-ner via datasets-server.

Stdlib only (urllib) so it runs anywhere, including Spaces builds.

    uv run python human_labeling/build_pool_queue.py [--config probe_candidates] [--split pool]
        [--limit 500] [--out human_labeling/queue.json]

Defaults target the annotatable config: probe_candidates/pool is the only
one with an embedded `passage`. --config probe_splits yields surface-only
triage items flagged ctx_missing (no passage on the Hub; see README).
Stratified sample: round-robin over origin x score-band so the queue
covers keep/confusion/drop and every origin instead of the head-heavy
random draw (proxy bands on the Hub score; true tags come from rescore).
"""

import argparse
import json
import urllib.parse
import urllib.request
from pathlib import Path

REPO = Path(__file__).resolve().parent
DATASET = "rafmacalaba/datause-ner"
API = "https://datasets-server.huggingface.co/rows"

BANDS = (("drop", 0.0, 0.05), ("confusion", 0.05, 0.9), ("keep", 0.9, 1.01))
# Proxy strata on the Hub (v3-era) score, used ONLY to stratify the sample
# across the eventual decision zones (same 0.05/0.9 edges as the audit
# deck). True keep/confusion/drop labels are assigned later by
# rescore_singlepass.probe_labels.decide; the written item band is always
# "unscored" until then.
def band_of(score: float) -> str:
    for name, lo, hi in BANDS:
        if lo <= score < hi:
            return name
    return "confusion"


def fetch_rows(config: str, split: str, offset: int, length: int) -> list[dict]:
    import time
    import urllib.error
    q = urllib.parse.urlencode(
        {"dataset": DATASET, "config": config, "split": split,
         "offset": offset, "length": length}
    )
    last = None
    for attempt in range(6):
        try:
            with urllib.request.urlopen(f"{API}?{q}", timeout=60) as r:
                payload = json.load(r)
            return [d["row"] for d in payload.get("rows", [])]
        except urllib.error.HTTPError as e:
            last = e
            if e.code not in (429, 500, 502, 503):
                raise
            time.sleep(2 ** attempt)
    raise last
def to_item(row: dict, config: str, split: str) -> dict:
    score = float(row.get("head_score") or 0.0)
    passage = row.get("passage")
    start, end = row.get("start"), row.get("end")
    item = {
        "key": row.get("key"),
        "surface": row.get("surface"),
        "ctx": passage,  # None for probe_splits rows (no passage on Hub)
        "ctx_missing": passage is None,
        "head_score": score,  # Hub v3-era placeholder; rescore overwrites
        "start": start, "end": end,
        "band": "unscored",  # rescore_singlepass tags keep/confusion/drop
        "origin": row.get("origin"),
        "specificity": row.get("extractor_label") or row.get("stratum") or "",
        "split": row.get("split", split),
        "queue": f"{config}/{split}",
    }
    return item


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="probe_candidates")
    ap.add_argument("--split", default="pool")
    ap.add_argument("--limit", type=int, default=500)
    ap.add_argument("--fetch", type=int, default=5000,
                    help="rows scanned from the Hub for stratification")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--out", default=str(REPO / "queue.json"))
    a = ap.parse_args()

    scanned: list[dict] = []
    offset = 0
    while len(scanned) < a.fetch:
        batch = fetch_rows(a.config, a.split, offset, min(100, a.fetch - len(scanned)))
        if not batch:
            break
        scanned.extend(batch)
        offset += len(batch)

    # stratify: round-robin over (origin, band)
    import random
    rng = random.Random(a.seed)
    buckets: dict[tuple, list[dict]] = {}
    for row in scanned:
        item = to_item(row, a.config, a.split)
        buckets.setdefault((item["origin"], item["band"]), []).append(item)
    for b in buckets.values():
        rng.shuffle(b)
    queue: list[dict] = []
    while len(queue) < min(a.limit, len(scanned)) and buckets:
        for k in sorted(buckets):
            if len(queue) >= a.limit:
                break
            if buckets[k]:
                queue.append(buckets[k].pop())
        buckets = {k: v for k, v in buckets.items() if v}
    rng.shuffle(queue)

    out = Path(a.out)
    out.write_text("\n".join(json.dumps(q) for q in queue) + "\n")
    origins = sorted({q["origin"] for q in queue})
    missing = sum(1 for q in queue if q["ctx_missing"])
    print(f"scanned={len(scanned)} queued={len(queue)} origins={origins} "
          f"ctx_missing={missing} -> {out}")


if __name__ == "__main__":
    main()