#!/usr/bin/env python3
"""39_coletar_dados.py — Streaming CONTROLADO de dados reais (Wikipédia PT-BR).

Tetos duros (não satura o ambiente): max_artigos e max_bytes do config.
Salva shards JSONL locais (treino/validação 90/10) + treina o BPE do
KHTST_RUST (vocab-alvo do config) e salva em dados/bpe.json.
Nada além dos tetos é baixado; 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:  # descarta esboços mínimos
            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")

    # normalização leve: primeiros parágrafos relevantes (descarta infobox ruidosa)
    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 do KHTST_RUST
    import khtst_rust as K
    vocab_alvo = int(cfg["modelo"]["vocab"])
    t0 = time.time()
    bpe = K.Bpe.treinar(treino, vocab_alvo)
    bpe.salvar(os.path.join(pasta, "dados", "bpe.json"))
    print(f"BPE: vocab={bpe.n_vocab} (alvo {vocab_alvo}) em {time.time()-t0:.1f}s")
    amostra = treino[0][:120]
    ids = bpe.codificar(amostra)
    rt = bpe.decodificar(ids)
    print("roundtrip OK:", rt[:80] == amostra[:80].strip()[:80] or len(rt) > 50)

if __name__ == "__main__":
    main()