#!/usr/bin/env python3
"""Stream FineWeb-Edu and write GPT-2 token .bin files for nanoGPT training."""

import argparse
import json
import os
import pickle
import shutil
import time
from typing import Optional

import numpy as np
import tiktoken
from datasets import load_dataset
from tqdm import tqdm


GPT2_VOCAB_SIZE = 50257


def parse_token_count(value: Optional[str]) -> Optional[int]:
    if value is None:
        return None
    raw = str(value).strip().replace("_", "").lower()
    if raw in {"none", "all", "full", "0"}:
        return None

    multipliers = {
        "k": 1_000,
        "m": 1_000_000,
        "b": 1_000_000_000,
    }
    suffix = raw[-1]
    if suffix in multipliers:
        return int(float(raw[:-1]) * multipliers[suffix])
    return int(raw)


def safe_output_paths(out_dir: str, overwrite: bool) -> tuple[str, str, str]:
    os.makedirs(out_dir, exist_ok=True)
    train_path = os.path.join(out_dir, "train.bin")
    val_path = os.path.join(out_dir, "val.bin")
    meta_path = os.path.join(out_dir, "meta.pkl")

    existing = [p for p in (train_path, val_path, meta_path) if os.path.exists(p)]
    if existing and not overwrite:
        joined = "\n  ".join(existing)
        raise FileExistsError(
            "Refusing to overwrite existing output files. Remove them or pass "
            f"--overwrite:\n  {joined}"
        )
    return train_path, val_path, meta_path


def write_tokens(handle, tokens: list[int]) -> int:
    if not tokens:
        return 0
    arr = np.asarray(tokens, dtype=np.uint16)
    arr.tofile(handle)
    return int(arr.size)


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Stream HuggingFaceFW/fineweb-edu sample-10BT, tokenize with the "
            "GPT-2 BPE tokenizer, and write data/fineweb/train.bin and val.bin."
        )
    )
    parser.add_argument("--out_dir", default="data/fineweb")
    parser.add_argument("--dataset", default="HuggingFaceFW/fineweb-edu")
    parser.add_argument("--name", default="sample-10BT")
    parser.add_argument("--split", default="train")
    parser.add_argument(
        "--max_tokens",
        default=None,
        help=(
            "Maximum total tokens to write, e.g. 100M, 1B, 10B. "
            "Use all/full/0 or omit for the full streamed split."
        ),
    )
    parser.add_argument(
        "--val_tokens",
        default="5M",
        help="Tokens reserved from the beginning of the stream for training-time val.bin.",
    )
    parser.add_argument(
        "--min_score",
        type=float,
        default=None,
        help="Optional FineWeb-Edu quality score filter, e.g. 3.0.",
    )
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()

    max_tokens = parse_token_count(args.max_tokens)
    val_tokens = parse_token_count(args.val_tokens)
    if val_tokens is None or val_tokens <= 0:
        raise ValueError("--val_tokens must be a positive token count")
    if max_tokens is not None and max_tokens <= val_tokens:
        raise ValueError("--max_tokens must be larger than --val_tokens")

    train_path, val_path, meta_path = safe_output_paths(args.out_dir, args.overwrite)
    train_part = train_path + ".part"
    val_part = val_path + ".part"

    for path in (train_part, val_part):
        if os.path.exists(path):
            if args.overwrite:
                os.remove(path)
            else:
                raise FileExistsError(
                    f"Partial output exists: {path}. Remove it or pass --overwrite."
                )

    enc = tiktoken.get_encoding("gpt2")
    eot = enc.eot_token

    ds = load_dataset(
        args.dataset,
        name=args.name,
        split=args.split,
        streaming=True,
    )

    total_written = 0
    val_written = 0
    train_written = 0
    docs_seen = 0
    docs_used = 0
    start = time.time()

    progress_total = max_tokens if max_tokens is not None else None
    pbar = tqdm(total=progress_total, unit="tok", smoothing=0.05)

    with open(val_part, "wb") as val_f, open(train_part, "wb") as train_f:
        for doc in ds:
            docs_seen += 1
            if args.min_score is not None and doc.get("score") is not None:
                if float(doc["score"]) < args.min_score:
                    continue

            text = doc.get("text")
            if not text:
                continue

            tokens = [eot] + enc.encode_ordinary(text)
            if max_tokens is not None:
                remaining = max_tokens - total_written
                if remaining <= 0:
                    break
                tokens = tokens[:remaining]

            cursor = 0
            if val_written < val_tokens:
                take = min(val_tokens - val_written, len(tokens))
                val_written += write_tokens(val_f, tokens[:take])
                cursor = take

            if cursor < len(tokens):
                train_written += write_tokens(train_f, tokens[cursor:])

            docs_used += 1
            total_written = val_written + train_written
            pbar.update(len(tokens))

            if max_tokens is not None and total_written >= max_tokens:
                break

    pbar.close()

    os.replace(train_part, train_path)
    os.replace(val_part, val_path)

    metadata = {
        "vocab_size": GPT2_VOCAB_SIZE,
        "tokenizer": "tiktoken:gpt2",
        "dataset": args.dataset,
        "name": args.name,
        "split": args.split,
        "max_tokens": max_tokens,
        "val_tokens_requested": val_tokens,
        "val_tokens_written": val_written,
        "train_tokens_written": train_written,
        "docs_seen": docs_seen,
        "docs_used": docs_used,
        "min_score": args.min_score,
        "elapsed_sec": round(time.time() - start, 2),
    }
    with open(meta_path, "wb") as f:
        pickle.dump({"vocab_size": GPT2_VOCAB_SIZE}, f)
    with open(os.path.join(args.out_dir, "data_info.json"), "w") as f:
        json.dump(metadata, f, indent=2)

    total_size = os.path.getsize(train_path) + os.path.getsize(val_path)
    print("\nFinished FineWeb-Edu data preparation")
    print(f"  train tokens: {train_written:,} -> {train_path}")
    print(f"  val tokens:   {val_written:,} -> {val_path}")
    print(f"  disk size:    {total_size / 1024**3:.2f} GiB")
    print(f"  meta:         {meta_path}")
    print(f"  info:         {os.path.join(args.out_dir, 'data_info.json')}")

    if shutil.disk_usage(args.out_dir).free < 5 * 1024**3:
        print("WARNING: less than 5 GiB free space remains in the output directory.")


if __name__ == "__main__":
    main()