#!/usr/bin/env python3
"""70_coleta_shards_v090.py — re-coleta os shards REAIS (Wikipédia PT-BR)
para a validação v0.9.0 em dados reais. NÃO re-treina o BPE (dados/bpe.json
já existe — só os shards tinham sido apagados na limpeza de estados).

Tetos duros idênticos ao 39_coletar_dados.py (max_artigos/max_bytes do
config): streaming controlado, sem cache global do datasets."""
import json, os, sys, time

TETO_PADRAO_ARTIGOS = 1200
TETO_PADRAO_BYTES = 6_000_000

def main():
    pasta = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    cfg = json.load(open(os.path.join(pasta, "configs", "base_rust.json")))
    dados_cfg = cfg.get("dados", {})
    max_artigos = int(dados_cfg.get("max_artigos", TETO_PADRAO_ARTIGOS))
    max_bytes = int(dados_cfg.get("max_bytes", TETO_PADRAO_BYTES))
    os.makedirs(os.path.join(pasta, "dados"), exist_ok=True)

    from datasets import load_dataset
    t0 = time.time()
    ds = load_dataset("wikimedia/wikipedia", "20231101.pt", split="train", streaming=True)
    textos, total_bytes = [], 0
    for ex in ds:
        t = (ex.get("text") or "").strip()
        if len(t) < 200:
            continue
        textos.append(t)
        total_bytes += len(t.encode("utf-8"))
        if len(textos) % 200 == 0:
            print(f"  ... {len(textos)} artigos, {total_bytes/1e6:.2f} MB", flush=True)
        if len(textos) >= max_artigos or total_bytes >= max_bytes:
            break
    print(f"coletado: {len(textos)} artigos, {total_bytes/1e6:.2f} MB em {time.time()-t0:.1f}s")

    def limpa(t: str) -> str:
        partes = [p.strip() for p in t.split("\n") if len(p.strip()) > 80]
        return " ".join(partes[:6])[:1800]

    textos = [limpa(t) for t in textos]
    n_val = max(1, len(textos) // 10)
    val, treino = textos[:n_val], textos[n_val:]
    with open(os.path.join(pasta, "dados", "shard_treino.jsonl"), "w") as f:
        for t in treino:
            f.write(json.dumps({"texto": t}, ensure_ascii=False) + "\n")
    with open(os.path.join(pasta, "dados", "shard_validacao.jsonl"), "w") as f:
        for t in val:
            f.write(json.dumps({"texto": t}, ensure_ascii=False) + "\n")
    print(f"shards: treino={len(treino)} val={len(val)} (bpe.json preservado)")

if __name__ == "__main__":
    main()