"""Pull a balanced scientific-domain human-vs-machine corpus from MAGE
(yaful/MAGE) via the HF datasets-server filter API — no full download.

MAGE: text / label (1=human, 0=machine) / src (domain+model). The 'sci_gen_*'
rows are scientific/academic text; the model name lives in src, giving us
leave-one-MODEL-out groups for honest validation. Length-filtered so every
doc has enough sentences for the burstiness/diversity features.

Output: data/calibration/mage_sci.jsonl  ({text, y, src_model})
Run:    python scripts/build_mage_corpus.py [per_model]   (default 120)
"""
import json, os, sys, time, urllib.parse, urllib.request

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "data", "calibration", "mage_sci.jsonl")
DS, CONFIG = "yaful/MAGE", "default"
MIN_CHARS = 800          # need a real document, not a snippet
PER_MODEL = int(sys.argv[1]) if len(sys.argv) > 1 else 120


def _get(url, tries=4):
    for t in range(tries):
        try:
            return json.load(urllib.request.urlopen(url, timeout=40))
        except Exception as e:
            if t == tries - 1:
                raise
            time.sleep(2 * (t + 1))


def filt(where, split, offset, length=100):
    u = (f"https://datasets-server.huggingface.co/filter?dataset={urllib.parse.quote(DS)}"
         f"&config={CONFIG}&split={split}&where={urllib.parse.quote(where)}"
         f"&offset={offset}&length={length}")
    return _get(u)


def collect(where, split, cap, by_model):
    """Page through a filtered view, bucket by src, length-filter, cap/model."""
    offset, total = 0, None
    while True:
        d = filt(where, split, offset)
        total = d.get("num_rows_total", 0)
        rows = d.get("rows", [])
        if not rows:
            break
        for r in rows:
            row = r["row"]
            txt, src = row.get("text", ""), row.get("src", "")
            if len(txt) < MIN_CHARS:
                continue
            bucket = by_model.setdefault(src, [])
            if len(bucket) < cap:
                bucket.append(txt)
        offset += len(rows)
        got = sum(min(len(v), cap) for v in by_model.values())
        print(f"  {split} offset={offset}/{total}  collected={got}", end="\r")
        if offset >= total or offset > 40000:
            break
    print()


def main():
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    human, machine = {}, {}
    for split in ("test", "validation", "train"):
        print(f"[human] {split}")
        collect("\"src\"='sci_gen_human'", split, PER_MODEL * 6, human)
        print(f"[machine] {split}")
        collect("\"src\" LIKE 'sci_gen_machine%'", split, PER_MODEL, machine)
        h = sum(len(v) for v in human.values())
        m = sum(len(v) for v in machine.values())
        if h >= PER_MODEL * 6 and len(machine) >= 3:
            break

    rows = []
    for txts in human.values():
        for t in txts[:PER_MODEL * 6]:
            rows.append({"text": t, "y": 1, "src_model": "human"})
    for src, txts in machine.items():
        model = src.replace("sci_gen_machine_continuation_", "").replace(
            "sci_gen_machine_", "")
        for t in txts:
            rows.append({"text": t, "y": 0, "src_model": "ai_" + model})
    with open(OUT, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")

    from collections import Counter
    c = Counter(r["src_model"] for r in rows)
    print(f"\nwrote {len(rows)} docs -> {OUT}")
    for k, v in sorted(c.items()):
        print(f"  {k:28s} {v}")


if __name__ == "__main__":
    main()