"""
Faz 1 / Adım 1 — Pre-tokenize + parquet shard (kaynak bazlı).

Akış: HF streaming -> sc_tokenizer (SP) ile encode + EOS -> 2048'lik dizilere paketle
(doc sınırları arası carry-over) -> ~256MB parquet shard'lar + manifest.json.
Mixture EĞİTİMDE uygulanır; shard'lar source/lang etiketli tutulur.

Kaynaklar (--source):
  en_fineweb_edu : HuggingFaceFW/fineweb-edu (sample-10BT)   lang=en
  tr_fineweb2_hq : epfml/FineWeb2-HQ (tur_Latn)              lang=tr
  tr_tc100b      : lumees/turkish-corpus-100b (data_dir=pretrain, ~103B, Apache) lang=tr  [v1.5b ANA TR]

Kullanım:
  python kod/faz1_01_tokenize_shard.py --source en_fineweb_edu --target_tokens 5_000_000 --seqs_per_shard 1500
"""
import os, sys, json, time, argparse
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pyarrow as pa
import pyarrow.parquet as pq
from datasets import load_dataset
from sc_tokenizer import SCTokenizer
from decontam import Decontaminator

SOURCES = {
    "en_fineweb_edu": ("HuggingFaceFW/fineweb-edu", "sample-10BT", "en"),
    "tr_fineweb2_hq": ("epfml/FineWeb2-HQ", "tur_Latn", "tr"),
    # v1.5b ANA TR: TC-100B (~103B, Apache, FineWeb-2-TR+Cosmos-sentetik+TR-Wiki+haber+math). 4. eleman = data_dir.
    "tr_tc100b": ("lumees/turkish-corpus-100b", None, "tr", "pretrain"),
    # NOT: StarCoder2/the-stack gate'li, datasets 4.x script kaynaklarını desteklemiyor.
    # codeparrot-clean (ham Python, content kolonu) gate'siz ikamesi.
    "code_codeparrot": ("codeparrot/codeparrot-clean", None, "code"),
    "math_openwebmath": ("open-web-math/open-web-math", None, "math"),
}

# içerik kolonu önceliği (codeparrot=content, owm/fineweb=text)
TEXT_KEYS = ("text", "content", "code")


def text_of(rec):
    for k in TEXT_KEYS:
        v = rec.get(k)
        if isinstance(v, str) and v:
            return v
    for v in rec.values():
        if isinstance(v, str) and len(v) > 0:
            return v
    return ""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", required=True, choices=list(SOURCES))
    ap.add_argument("--target_tokens", type=lambda x: int(float(x)), default=5_000_000)  # 5e7/7.5e9 kabul
    ap.add_argument("--seq_len", type=int, default=2048)
    ap.add_argument("--seqs_per_shard", type=int, default=32768)  # ~256MB @2048 int32
    ap.add_argument("--out_dir", default="kod/data/shards")
    ap.add_argument("--decontam", default="kod/data/decontam_13gram.pkl.gz",
                    help="13-gram index; yoksa filtre atlanır")
    ap.add_argument("--n_shards", type=int, default=1, help="paralel worker sayısı (ds.shard ile parça böl)")
    ap.add_argument("--shard_index", type=int, default=0, help="bu worker'ın parça indeksi (0..n_shards-1)")
    args = ap.parse_args()

    src = SOURCES[args.source]
    repo, cfg, lang = src[0], src[1], src[2]
    data_dir = src[3] if len(src) > 3 else None   # dizin-yapılı dataset (TC-100B: pretrain/)
    outdir = os.path.join(args.out_dir, args.source)
    if args.n_shards > 1:                          # paralel: her worker kendi alt-dizinine yazar (çakışma yok)
        outdir = os.path.join(outdir, f"w{args.shard_index:02d}")
    os.makedirs(outdir, exist_ok=True)
    per_target = args.target_tokens // max(1, args.n_shards)   # worker başına hedef
    tok = SCTokenizer()
    dec = None
    if args.decontam and os.path.exists(args.decontam):
        dec = Decontaminator.load(args.decontam)
        print(f"[decontam] {len(dec.grams):,} adet 13-gram yüklendi → kontamine doc'lar atlanacak")
    else:
        print("[decontam] index yok → filtre KAPALI (uyarı: eval sızıntısı kontrol edilmiyor)")
    print(f"[{args.source} w{args.shard_index}/{args.n_shards}] repo={repo} dir={data_dir} lang={lang} "
          f"| seq_len={args.seq_len} | hedef ~{per_target/1e6:.1f}M token | vocab={tok.vocab_size}")

    ds = load_dataset(repo, name=cfg, data_dir=data_dir, split="train", streaming=True)
    if args.n_shards > 1:
        ds = ds.shard(num_shards=args.n_shards, index=args.shard_index)   # bu worker'a düşen parçalar
    buf, seqs, shard_idx = [], [], 0
    tot_tokens = tot_seqs = ndocs = n_skip = 0
    manifest = {"source": args.source, "repo": repo, "config": cfg, "data_dir": data_dir, "lang": lang,
                "seq_len": args.seq_len, "shards": []}
    t0 = time.perf_counter()

    def flush():
        nonlocal seqs, shard_idx, tot_seqs
        if not seqs:
            return
        fn = f"shard_{shard_idx:05d}.parquet"
        table = pa.table({
            # uint16: token id < 48000 < 65535 → diski/upload'ı yarıya indirir (int32 yerine)
            "input_ids": pa.array(seqs, type=pa.list_(pa.uint16())),
            "source": pa.array([args.source] * len(seqs)),
            "lang": pa.array([lang] * len(seqs)),
        })
        pq.write_table(table, os.path.join(outdir, fn), compression="zstd")
        manifest["shards"].append({"file": fn, "n_seqs": len(seqs)})
        tot_seqs += len(seqs)
        shard_idx += 1
        seqs = []

    for rec in ds:
        ndocs += 1
        txt = text_of(rec)
        if dec is not None and dec.is_contaminated(txt):
            n_skip += 1
            continue                       # eval ile örtüşen doc'u at (dekontaminasyon)
        buf.extend(tok.encode(txt, add_eos=True))
        while len(buf) >= args.seq_len:
            seqs.append(buf[:args.seq_len]); buf = buf[args.seq_len:]
            tot_tokens += args.seq_len
            if len(seqs) >= args.seqs_per_shard:
                flush()
        if tot_tokens >= per_target:
            break
        if ndocs % 5000 == 0:
            dt = time.perf_counter() - t0
            print(f"  {tot_tokens/1e6:6.2f}M tok | {ndocs} doc | {shard_idx} shard | "
                  f"{tot_tokens/1e6/max(1e-9,dt):.2f} M tok/s", flush=True)

    flush()  # kalan
    manifest["n_seqs"] = tot_seqs
    manifest["n_tokens"] = tot_seqs * args.seq_len
    manifest["n_docs"] = ndocs
    manifest["n_decontam_skipped"] = n_skip
    manifest["decontam"] = bool(dec)
    with open(os.path.join(outdir, "manifest.json"), "w", encoding="utf-8") as f:
        json.dump(manifest, f, ensure_ascii=False, indent=2)
    dt = time.perf_counter() - t0
    skip_pct = 100 * n_skip / max(1, ndocs)
    print(f"[{args.source}] BİTTİ: {tot_seqs} dizi × {args.seq_len} = {tot_seqs*args.seq_len/1e6:.2f}M token "
          f"| {shard_idx} shard | {ndocs} doc | {dt:.0f}s ({tot_tokens/1e6/max(1e-9,dt):.2f} M tok/s)")
    print(f"  dekontaminasyon: {n_skip} doc atlandı ({skip_pct:.2f}%) | -> {outdir}/")


if __name__ == "__main__":
    main()