"""Phase 2.2 — pull open RU datasets → data/cpt_raw/<bucket>/*.jsonl

Streams from HuggingFace so we never materialize full dumps. Each bucket
over-pulls ~1.5x its token target (cleaning drops 20-40%). Writes newline JSONL
{"text": "..."} shards; corpus_clean.py consumes them.

Usage:
    huggingface-cli login                 # accept CulturaX/OSCAR licenses first
    python corpus_fetch.py culturax       # one bucket
    python corpus_fetch.py all            # every open bucket (not synth/logs)

Buckets synth (corpus_synth.py) and logs (corpus_logs.py) are handled separately.
"""
import sys
import json

from datasets import load_dataset

from corpus_common import RAW, TARGET_TOKENS, BUCKET_SHARE

# rough chars-per-token for RU under the Qwen tokenizer; only used to size pulls.
CHARS_PER_TOK = 3.5
OVERPULL = 1.5

# bucket -> (loader kwargs, function mapping a row to its text string)
SOURCES = {
    "culturax": (
        dict(path="allenai/c4", name="ru", split="train", streaming=True),
        lambda r: r.get("text", ""),
    ),
    "wiki": (
        dict(path="wikimedia/wikipedia", name="20231101.ru", split="train", streaming=True),
        lambda r: r.get("text", ""),
    ),
    "books": (
        dict(path="IlyaGusev/gazeta", split="train", streaming=True),
        lambda r: r.get("text", ""),
    ),
}


def target_chars(bucket: str) -> int:
    return int(TARGET_TOKENS * BUCKET_SHARE[bucket] * CHARS_PER_TOK * OVERPULL)


def fetch(bucket: str):
    if bucket not in SOURCES:
        raise SystemExit(f"unknown bucket {bucket}; known: {list(SOURCES)}")
    budget = target_chars(bucket)
    out_dir = RAW / bucket
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "shard-000.jsonl"

    kw, getter = SOURCES[bucket]
    ds = load_dataset(**kw)

    written = 0
    n = 0
    with open(out_path, "w", encoding="utf-8") as f:
        for row in ds:
            text = (getter(row) or "").strip()
            if len(text) < 200:            # skip fragments early
                continue
            f.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
            written += len(text)
            n += 1
            if n % 5000 == 0:
                print(f"[{bucket}] {n} docs, {written/1e6:.1f}M chars / {budget/1e6:.1f}M target")
            if written >= budget:
                break
    print(f"[{bucket}] DONE {n} docs, {written/1e6:.1f}M chars -> {out_path}")


if __name__ == "__main__":
    args = sys.argv[1:] or ["all"]
    buckets = list(SOURCES) if args == ["all"] else args
    for b in buckets:
        fetch(b)