"""MicroSupra-10k — preparação de dados (Fineweb-Edu + DCLM-Edu, streaming)."""
import os
import time
import numpy as np
from datasets import load_dataset, interleave_datasets
from transformers import PreTrainedTokenizerFast
from huggingface_hub import hf_hub_download

TOKENS_TRAIN = int(os.environ.get("TOKENS_TRAIN", 20_000_000))
TOKENS_VAL = int(os.environ.get("TOKENS_VAL", 1_000_000))
VOCAB = 1024
BOS, EOS = 0, 2  # <s>=0, </s>=2 (pad do original é <pad>=1; não usamos padding)
OUT = os.path.dirname(os.path.abspath(__file__)) + "/data"
os.makedirs(OUT, exist_ok=True)

print("[*] Baixando tokenizador do original (SupraLabs/MicroSupra-1k)...", flush=True)
tok_path = hf_hub_download("SupraLabs/MicroSupra-1k", "tokenizer.json")
tokenizer = PreTrainedTokenizerFast(
    tokenizer_file=tok_path,
    bos_token="<s>", eos_token="</s>", pad_token="<pad>", unk_token="<unk>",
)

print("[*] Abrindo streams: fineweb-edu sample-10BT + dclm-edu (edu_int_score>=3)...", flush=True)
fw = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)
dclm = load_dataset("HuggingFaceTB/dclm-edu", split="train", streaming=True)
dclm = dclm.filter(lambda score: score >= 3, input_columns=["edu_int_score"])
mix = interleave_datasets([fw, dclm])  # 50/50

total_target = TOKENS_TRAIN + TOKENS_VAL
buf = np.empty(total_target, dtype=np.uint16)
pos, docs, t0 = 0, 0, time.time()

for ex in mix:
    text = ex.get("text") or ""
    if not text:
        continue
    ids = tokenizer.encode(text)
    if not ids:
        continue
    ids = np.array([BOS] + ids + [EOS], dtype=np.int64)
    ids = ids[ids < VOCAB]  # segurança (BPE 1024 não deve estourar)
    k = min(len(ids), total_target - pos)
    buf[pos:pos + k] = ids[:k].astype(np.uint16)
    pos += k
    docs += 1
    if pos % 2_000_000 < len(ids):
        print(f"    {pos/1e6:.1f}M tokens ({docs} docs, {time.time()-t0:.0f}s)", flush=True)
    if pos >= total_target:
        break

print(f"[*] Total tokenizado: {pos:,} tokens em {time.time()-t0:.0f}s", flush=True)

# separação: últimos TOKENS_VAL tokens viram validação
train = buf[:pos - TOKENS_VAL]
val = buf[pos - TOKENS_VAL:pos]
train.tofile(f"{OUT}/train.bin")
val.tofile(f"{OUT}/val.bin")
print(f"[*] train.bin: {len(train):,} tokens | val.bin: {len(val):,} tokens", flush=True)
print("[*] DONE", flush=True)