"""Quick training-data builder: a few corpus shards at the eval-156 commit, rows with short
prompts, dumped as drawn.jsonl (the dataset.load format). No ranking needed: the QAT trains on
any diverse rows; the margin is verified on the official harness later."""
import json
import sys
import urllib.request
from pathlib import Path

sys.path.insert(0, "/root/ingot/src")
import pyarrow.parquet as pq

COMMIT = "70d2121e5dcbfc9f2c2e79c890563d96429eeb98"
REPO = "dendriteholdings/ingot_corpus"
BASE = f"https://huggingface.co/datasets/{REPO}/resolve/{COMMIT}/"

req = urllib.request.Request(
    f"https://huggingface.co/api/datasets/{REPO}/tree/{COMMIT}?recursive=1",
    headers={"User-Agent": "mk"},
)
tree = json.load(urllib.request.urlopen(req))
shards = sorted(
    t["path"] for t in tree if t["path"].startswith("data/chungus-") and t["path"].endswith(".parquet")
)
print("shards:", len(shards))

out = []
tmp = Path("/tmp/shards")
tmp.mkdir(exist_ok=True)
for shard in shards[::17]:  # ~15 shards, spread over the corpus
    url = BASE + shard
    local = tmp / Path(shard).name
    urllib.request.urlretrieve(url, local)
    parq = pq.ParquetFile(local)
    cols = ["dataset", "split", "entry_name", "domain", "language", "user_prompt"]
    if "conversation" in parq.schema_arrow.names:
        cols.append("conversation")
    df = parq.read(columns=cols).to_pylist()
    for row in df:
        if row.get("conversation"):
            conv = json.loads(row["conversation"])
            msgs, tools = conv["messages"], conv.get("tools")
        else:
            msgs = [{"role": "user", "content": row["user_prompt"]}]
            tools = None
        size = sum(len(m["content"].encode()) for m in msgs)
        if size > 3500:  # short-ish prompts only
            continue
        sid = f"{row['dataset']}/{row['split']}/{row['entry_name']}/{row['language']}"
        out.append({"id": sid, "messages": msgs, "tools": tools})
        if len(out) >= 160:
            break
    print(f"{shard}: {len(out)} kept", flush=True)
    if len(out) >= 160:
        break

Path("/root/drawn.jsonl").write_text(
    "".join(json.dumps(r, ensure_ascii=False) + "\n" for r in out)
)
print("wrote /root/drawn.jsonl:", len(out), "rows")