import os
import random

from datasets import load_dataset
from huggingface_hub import HfApi

from pipeline import export_dataset

REPO = "openbmb/Ultra-FineWeb-L3"
FRACTION = 0.1
SEED = 42
SAVE_ARROW = False

CONFIGS = {
    "Ultra-FineWeb-L3-en-QA-Synthetic": "data/ultrafineweb_en_l3/qa/",
    "Ultra-FineWeb-L3-zh-QA-Synthetic": "data/ultrafineweb_zh_l3/qa/",
}

HF_CACHE_DIR = "./cached_pt/ultra-fineweb-l3-qa-synthetic"
OUTPUT_DIR = "./dataset"


def process_func(input_dict: dict):
    return {"text": input_dict["content"]}


def main():
    api = HfApi()
    for config, prefix in CONFIGS.items():
        lang = "en" if "-en-" in config else "zh"

        shards = [
            f.path
            for f in api.list_repo_tree(
                REPO, path_in_repo=prefix, recursive=True, repo_type="dataset"
            )
            if f.path.endswith(".parquet")
        ]
        k = max(1, int(len(shards) * FRACTION))
        selected = random.Random(SEED).sample(shards, k)
        print(f"[{config}] total shards={len(shards)}, selected={k}", flush=True)

        dataset = load_dataset(
            REPO,
            data_files=selected,
            split="train",
            cache_dir=HF_CACHE_DIR,
        )
        print(f"[{config}] loaded {len(dataset)} rows", flush=True)

        if SAVE_ARROW:
            arrow_dir = os.path.join(
                HF_CACHE_DIR, f"arrow-{lang}"
            )
            dataset.save_to_disk(arrow_dir)
            print(f"[{config}] cached arrow to {arrow_dir}", flush=True)

        export_dataset(
            dataset=dataset,
            output_dir=OUTPUT_DIR,
            output_prefix=f"ultra-fineweb-l3-{lang}-qa-synthetic-10pct-pretrain",
            process_func=process_func,
        )


if __name__ == "__main__":
    main()