"""Streaming document sources for the GPT-U-20M data mix.

Four sub-streams feed the mix: `dclm`, `edu`, and the Stack v3 slice split by file language into `code_py` and
`code_other`, so Python can be held at ~27% of the code tokens (it is only ~4% of Stack v3 naturally).
Each source's file list is split once: its last file is the holdout file and never feeds the tokenizer corpus
or the training shards.
"""
import json
import random
from pathlib import Path

import pyarrow as pa
import pyarrow.compute as pc
from datasets import Features, Value, load_dataset
from huggingface_hub import HfFileSystem

SEED = 20_453_760
MIX = {"dclm": 0.45, "edu": 0.35, "code": 0.20}
PYTHON_SHARE = 0.27  # spec: ~25-30% Python inside the code slice
STREAMS = ("dclm", "edu", "code_py", "code_other")
STREAM_DOMAIN = {"dclm": "dclm", "edu": "edu", "code_py": "code", "code_other": "code"}
STREAM_SHARE = {
    "dclm": MIX["dclm"],
    "edu": MIX["edu"],
    "code_py": MIX["code"] * PYTHON_SHARE,
    "code_other": MIX["code"] * (1 - PYTHON_SHARE),
}
SOURCES = {  # domain -> (dataset repo, file glob inside the repo)
    "dclm": ("mlfoundations/dclm-baseline-1.0", "global-shard_*/local-shard_*/*.jsonl.zst"),
    "edu": ("HuggingFaceFW/fineweb-edu", "sample/10BT/*.parquet"),  # uniform sample over all CC dumps
    "code": ("HuggingFaceCode/stack-v3-train", "data/*.parquet"),
}
FEATURES = Features({"text": Value("string"), "src": Value("string")})


def file_splits(cache_path: Path) -> dict:
    """{domain: {"train": [...], "holdout": [...]}}. Train order is a fixed shuffle, cached so every script agrees."""
    if cache_path.exists():
        return json.loads(cache_path.read_text())
    fs = HfFileSystem()
    splits = {}
    for domain, (repo, pattern) in SOURCES.items():
        prefix = f"datasets/{repo}/"
        files = sorted(p[len(prefix):] for p in fs.glob(prefix + pattern))
        train = files[:-1]
        random.Random(SEED).shuffle(train)
        splits[domain] = {"train": train, "holdout": files[-1:]}
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    cache_path.write_text(json.dumps(splits, indent=1))
    return splits


def _code_files(table: pa.Table, python: bool, src: str) -> pa.Table:
    """Repository rows -> one row per non-vendor, non-empty file of the requested language group."""
    files = pc.list_flatten(table["files"])
    content = pc.struct_field(files, "content")
    is_py = pc.equal(pc.struct_field(files, "language"), "Python")
    keep = pc.and_(pc.invert(pc.struct_field(files, "is_vendor")), is_py if python else pc.invert(is_py))
    keep = pc.and_(keep, pc.greater(pc.binary_length(content), 0))
    text = pc.filter(content, keep)
    return pa.table({"text": text, "src": pa.array([src] * len(text), pa.string())})


def _tag(table: pa.Table, src: str) -> pa.Table:
    return pa.table({"text": table["text"], "src": pa.array([src] * len(table), pa.string())})


def open_stream(name: str, files: list[str]):
    """Streaming IterableDataset with columns `text` and `src` for one sub-stream, reading `files` in order."""
    domain = STREAM_DOMAIN[name]
    repo = SOURCES[domain][0]
    if domain == "code":
        ds = load_dataset(repo, data_files=files, split="train", streaming=True, columns=["files"])
        ds = ds.with_format("arrow").map(_code_files, batched=True, batch_size=64, remove_columns=["files"],
                                         features=FEATURES, fn_kwargs={"python": name == "code_py", "src": name})
    else:
        kwargs = {"columns": ["text"]} if domain == "edu" else {}
        ds = load_dataset(repo, data_files=files, split="train", streaming=True, **kwargs).select_columns(["text"])
        ds = ds.with_format("arrow").map(_tag, batched=True, batch_size=1000, features=FEATURES, fn_kwargs={"src": name})
    return ds.with_format(None)