"""Offline data preparation (parallel tokenization).

Streams raw text from a Hugging Face dataset, tokenizes it across many CPU
worker processes, and appends the token ids to a per-source .bin file (uint32).
Documents are separated by the tokenizer's end-of-text id.

Code sources are read one language at a time (e.g. "stack_python" reads only
Python) and pass a quality filter that drops data dumps and minified files.

--skip_docs skips documents an earlier run already used, so batches never
overlap. Each run prints the value to pass as the next run's --skip_docs.
"""

import os

# each worker is its own process, so the tokenizer must NOT also spawn internal
# threads — those Rust threads race Python's shutdown and crash on exit
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import argparse
import itertools
import multiprocessing as mp
import re

import numpy as np
from datasets import load_dataset

from .tokenizer import Tokenizer

# stop cleanly (keeping everything written) if the shared disk gets this full
MIN_FREE_BYTES = 300 * 1024**3
DTYPE = np.uint32                      # vocab > 65535, so 2 bytes is not enough


def free_bytes(path: str) -> int:
    st = os.statvfs(path)
    return st.f_bavail * st.f_frsize


# each source: (HF repo, config, data_dir, text column, split, is_code)
SOURCES = {
    "fineweb":   ("HuggingFaceFW/fineweb-edu", "sample-350BT", None, "text", "train", False),
    "finemath":  ("HuggingFaceTB/finemath", "finemath-4plus", None, "text", "train", False),
    "finemath3": ("HuggingFaceTB/finemath", "finemath-3plus", None, "text", "train", False),
    "cosmopedia":  ("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", None, "text", "train", False),
    "infiwebmath": ("HuggingFaceTB/finemath", "infiwebmath-4plus", None, "text", "train", False),
    "fineweb_hq":  ("HuggingFaceFW/fineweb-edu", "sample-350BT", None, "text", "train", False),
    "stack":     ("bigcode/the-stack-dedup", None, None, "content", "train", True),
}
# one source per language: "stack_python" reads only the data/python folder
CODE_LANGS = ["python", "javascript", "java", "cpp", "typescript",
              "go", "rust", "shell", "sql"]
for _lang in CODE_LANGS:
    SOURCES[f"stack_{_lang}"] = ("bigcode/the-stack-dedup", None, f"data/{_lang}",
                                 "content", "train", True)
# keep only documents with at least this quality score (fineweb-edu's 0-5 "int_score")
MIN_INT_SCORE = {"fineweb_hq": 4}

# --- code quality filter (StarCoder-style rules) --------------------------
_HEX = re.compile(r"0x[0-9a-fA-F]{2,}")


def _looks_like_good_code(text: str) -> bool:
    lines = text.split("\n")
    if max(len(l) for l in lines) > 1000:                  # minified / one giant line
        return False
    if len(text) / len(lines) > 100:                        # very long lines on average
        return False
    if sum(c.isalnum() for c in text) / len(text) < 0.25:   # mostly symbols
        return False
    if sum(len(m) for m in _HEX.findall(text)) / len(text) > 0.2:   # hex data tables
        return False
    if sum(c.isdigit() for c in text) / len(text) > 0.4:    # numeric data tables
        return False
    return True


# --- worker side ---------------------------------------------------------
# each worker builds its own tokenizer once (tokenizers aren't shareable
# across processes), then reuses it for every document it receives
_worker_tok = None
_worker_is_code = False


def _init_worker(is_code):
    global _worker_tok, _worker_is_code
    _worker_tok = Tokenizer()
    _worker_is_code = is_code


def _encode_one(text: str) -> np.ndarray:
    # a rejected document comes back empty; the main process skips it
    if not text or (_worker_is_code and not _looks_like_good_code(text)):
        return np.empty(0, dtype=DTYPE)
    ids = _worker_tok.encode(text)
    ids.append(_worker_tok.eot_id)
    return np.array(ids, dtype=DTYPE)


# --- main side -----------------------------------------------------------
SCORE_STATS = {"seen": 0, "low": 0}   # docs streamed / docs below min score


def _doc_texts(ds, text_col, min_score=0):
    # pull the text out of each document, skipping low-quality ones if asked
    for doc in ds:
        SCORE_STATS["seen"] += 1
        if min_score and (doc.get("int_score") or 0) < min_score:
            SCORE_STATS["low"] += 1
            continue
        yield doc[text_col] or ""

def prepare(source: str, out_dir: str, target_tokens: int,
            num_workers: int, chunk_docs: int = 10_000, skip_docs: int = 0,
            data_files=None):
    repo, config, data_dir, text_col, split, is_code = SOURCES[source]
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"{source}.bin")
    # fresh data appended to an old file would mix two different runs
    if skip_docs == 0 and os.path.exists(out_path) and os.path.getsize(out_path) > 0:
        raise SystemExit(f"[refuse] {out_path} already has data: delete it first")

    if data_files:
        # read one specific file directly (e.g. a late shard for a held-out set)
        ds = load_dataset(repo, data_files=data_files, split=split, streaming=True)
    else:
        ds = load_dataset(repo, config, data_dir=data_dir, split=split, streaming=True)
    print(f"[read] {repo} files={data_files} config={config} data_dir={data_dir}", flush=True)
    texts = _doc_texts(ds, text_col, MIN_INT_SCORE.get(source, 0))
    
    
    # skip documents an earlier run already consumed, so batches never overlap
    if skip_docs > 0:
        print(f"  [skip] skipping first {skip_docs:,} documents", flush=True)
        texts = itertools.islice(texts, skip_docs, None)

    written = n_docs = n_kept = 0
    stopped_early = False

    # imap keeps results IN ORDER and pulls lazily, so memory stays small
    pool = mp.Pool(num_workers, initializer=_init_worker, initargs=(is_code,))
    try:
        with open(out_path, "ab") as f:
            for arr in pool.imap(_encode_one, texts, chunksize=64):
                n_docs += 1                    # counts every streamed doc (for --skip_docs)
                if arr.size:
                    arr.tofile(f)              # main process is the only writer
                    written += arr.size
                    n_kept += 1
                if written >= target_tokens:
                    break
                if n_docs % chunk_docs == 0:
                    if free_bytes(out_dir) < MIN_FREE_BYTES:
                        print(f"\n[stop] free space below {MIN_FREE_BYTES/1e9:.0f} GB",
                              flush=True)
                        stopped_early = True
                        break
    finally:
        pool.terminate()
        pool.join()

    if MIN_INT_SCORE.get(source):
        low = 100 * SCORE_STATS["low"] / max(SCORE_STATS["seen"], 1)
        print(f"\n[score] {low:.0f}% of streamed docs had int_score < {MIN_INT_SCORE[source]}")
    status = "STOPPED EARLY (disk)" if stopped_early else "done"
    dropped = 100 * (1 - n_kept / max(n_docs, 1))
    print(f"\n[{status}] {source}: {written/1e6:.1f}M tokens, "
          f"{n_docs:,} docs read, {dropped:.0f}% dropped by code filter")
    print(f"[next] for the following run, pass --skip_docs {skip_docs + n_docs}",
          flush=True)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--source", required=True, choices=list(SOURCES))
    p.add_argument("--out_dir", default="/fsx/ajay-stu/ar-data/bin")
    p.add_argument("--target_tokens", type=int, default=100_000_000)
    p.add_argument("--skip_docs", type=int, default=0)
    p.add_argument("--num_workers", type=int, default=max(1, mp.cpu_count() - 4))
    p.add_argument("--data_files", default=None)
    
    args = p.parse_args()
    prepare(args.source, args.out_dir, args.target_tokens, args.num_workers,
            skip_docs=args.skip_docs, data_files=args.data_files or None)

    # HF's fast tokenizer leaves native threads that crash normal interpreter
    # shutdown; the file is closed above, so exit immediately
    os._exit(0)