"""
Kaggle Cloud Data Collector for ViuMini-MoE-242M.
Streams massive open datasets, cleans with stable blake2b hash dedup,
chunks into ~150MB Parquet shards (configurable via --shard_mb, default 150),
uploads to Hugging Face dataset repo `ViuAI/viu-mini-raw-pretrain`, and immediately deletes
local temporary files to keep Kaggle disk usage < 2GB.

Features:
  - Smart Auto-Resume: Automatically queries HF dataset repo, finds existing shards,
    starts numbering at (max_shard_id + 1), and fast-skips already processed rows.
  - ~150MB Shards: keeps HF viewer + DataLoader stable (old 1GB shards caused 1.6GB
    english files + viewer HTTP 500). Row cap 100k also enforced.
  - Zero Sub-Sampling: Ingests full corpora without artificial truncation.

Usage on Kaggle:
  python data/scripts/kaggle_collector.py --stream all --repo ViuAI/viu-mini-raw-pretrain --shard_mb 150
  python data/scripts/kaggle_collector.py --stream hindi --repo ViuAI/viu-mini-raw-pretrain --shard_mb 150
"""
import argparse
import gc
import gzip
import hashlib
import itertools
import json
import os
from pathlib import Path
import re
import sys
import tempfile
import time

DEV = re.compile(r'[\u0900-\u097F]')
LAT = re.compile(r'[A-Za-z]')


def get_hf_token():
    """Retrieve HF token from Kaggle secrets, Colab secrets, env var, or user input."""
    # 1. Check Kaggle secrets
    try:
        from kaggle_secrets import UserSecretsClient
        tok = UserSecretsClient().get_secret("HF_TOKEN")
        if tok:
            return tok
    except Exception:
        pass
    # 2. Check Google Colab secrets
    try:
        from google.colab import userdata
        tok = userdata.get("HF_TOKEN")
        if tok:
            return tok
    except Exception:
        pass
    # 3. Check environment variable
    tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if tok:
        return tok
    return None


def stable_hash(s: str) -> int:
    """Stable 64-bit hash (PYTHONHASHSEED-independent). hash() is randomized per run,
    so dedup was unstable across Kaggle restarts -> duplicates. Use blake2b."""
    return int.from_bytes(hashlib.blake2b(s.encode("utf-8"), digest_size=8).digest(), "big")


def dev_ratio(s: str) -> float:
    dev = len(DEV.findall(s))
    lat = len(LAT.findall(s))
    tot = dev + lat
    return dev / max(tot, 1)


def is_valid_line(s: str, lang: str) -> bool:
    if len(s) < 10 or len(s.split()) < 3:
        return False
    if lang == "bilingual":
        return len(s) >= 15
    r = dev_ratio(s)
    if lang == "hinglish":
        # Must have Latin letters, Devanagari ratio <= 30%
        return bool(LAT.search(s)) and r <= 0.3
    elif lang == "hindi":
        # Must have Devanagari script
        return bool(DEV.search(s))
    elif lang == "english":
        # Must have Latin, zero Devanagari
        return bool(LAT.search(s)) and not bool(DEV.search(s))
    return True


def apply_skip(ds, skip_rows: int, source_name: str):
    """Applies fast skipping or falls back to itertools.islice."""
    if skip_rows <= 0:
        return ds
    try:
        ds = ds.skip(skip_rows)
        print(f"[resume] Fast-skipped {skip_rows:,} rows from {source_name}.")
        return ds
    except Exception as e:
        print(f"[resume] Fast-skip failed ({e}), using iterator slice for {source_name}...")
        return itertools.islice(ds, skip_rows, None)


def stream_hindi_data(skip_rows: int = 0):
    """Streams full IndicCorpV2 hin_Deva (152M rows) and Sangraha."""
    try:
        from datasets import load_dataset
        print("[source] Loading ai4bharat/IndicCorpV2 (indiccorp_v2, hin_Deva) streaming...")
        ds = load_dataset("ai4bharat/IndicCorpV2", "indiccorp_v2", split="hin_Deva", streaming=True)
        ds = apply_skip(ds, skip_rows, "IndicCorpV2")
        for r in ds:
            txt = (r.get("text") or "").strip()
            if txt:
                yield txt, "indiccorp_v2"
    except Exception as e:
        print(f"[warn] IndicCorpV2 stream error: {e}, falling back to sangraha...")
        try:
            from datasets import load_dataset
            ds = load_dataset("ai4bharat/sangraha", "verified", split="hin", streaming=True)
            ds = apply_skip(ds, skip_rows, "sangraha")
            for r in ds:
                txt = (r.get("text") or "").strip()
                if txt:
                    yield txt, "sangraha"
        except Exception as e2:
            print(f"[error] Sangraha stream error: {e2}")


def stream_english_data(skip_rows: int = 0):
    """Streams Fineweb-edu / TinyStories educational slices."""
    from datasets import load_dataset
    print("[source] Loading HuggingFaceFW/fineweb-edu (sample-10BT or train) streaming...")
    try:
        ds = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "fineweb-edu")
        for r in ds:
            txt = (r.get("text") or "").strip()
            if txt:
                yield txt, "fineweb-edu"
    except Exception as e:
        print(f"[warn] fineweb-edu fallback to TinyStories: {e}")
        ds = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "tinystories")
        for r in ds:
            txt = (r.get("text") or "").strip()
            if txt:
                yield txt, "tinystories"


def stream_hinglish_data(skip_rows: int = 0):
    """Streams sk-community/romanized_hindi and conversational corpora."""
    from datasets import load_dataset
    print("[source] Loading sk-community/romanized_hindi streaming...")
    try:
        ds = load_dataset("sk-community/romanized_hindi", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "romanized_hindi")
        for r in ds:
            txt = (r.get("Transliterated Hindi") or r.get("hi_rom") or r.get("text") or "").strip()
            if txt:
                yield txt, "romanized_hindi"
    except Exception as e:
        print(f"[warn] sk-community/romanized_hindi stream error: {e}")


def stream_domain_data():
    """Streams curated 26-pillar knowledge from collect_domains."""
    try:
        from data.scripts.collect_domains import generate_domain_samples
        for rec in generate_domain_samples():
            yield rec["text"], rec["source"], rec["lang"], rec["domain"]
    except ImportError:
        try:
            from collect_domains import generate_domain_samples
            for rec in generate_domain_samples():
                yield rec["text"], rec["source"], rec["lang"], rec["domain"]
        except Exception as e:
            print(f"[warn] domain generator import failed: {e}")


def stream_translation_data(skip_rows: int = 0):
    """Streams parallel translation data from cfilt/iitb-english-hindi & Romanized Hindi."""
    from datasets import load_dataset
    print("[source] Loading cfilt/iitb-english-hindi parallel translation streaming...")
    try:
        ds = load_dataset("cfilt/iitb-english-hindi", split="train", streaming=True)
        ds = apply_skip(ds, skip_rows, "iitb-english-hindi")
        for r in ds:
            trans = r.get("translation") or {}
            en = (trans.get("en") or "").strip()
            hi = (trans.get("hi") or "").strip()
            if en and hi and len(en) >= 10 and len(hi) >= 10:
                text = f"English: {en}\nHindi: {hi}"
                yield text, "iitb-english-hindi", "bilingual", "translation"
    except Exception as e:
        print(f"[warn] iitb translation stream fallback: {e}")


def get_existing_shards_info(api, repo_id: str, stream_name: str):
    """
    Checks HF repo for existing shards in {stream_name}/.
    Returns (next_shard_id, max_shard_id, existing_count).
    """
    if not api:
        return 0, -1, 0
    try:
        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
        prefix = f"{stream_name}/train-"
        shard_ids = []
        for f in files:
            if f.startswith(prefix) and (f.endswith(".parquet") or f.endswith(".jsonl.gz")):
                base = f[len(prefix):]
                num_str = base.split(".")[0]
                if num_str.isdigit():
                    shard_ids.append(int(num_str))
        if shard_ids:
            max_id = max(shard_ids)
            return max_id + 1, max_id, len(shard_ids)
    except Exception as e:
        print(f"[warn] Could not inspect existing shards for {stream_name}: {e}")
    return 0, -1, 0


def write_and_upload_shard(records, shard_id, stream_name, repo_id, api, dry_run=False, tmp_dir=None):
    """
    Writes records to a compressed parquet shard, uploads to HF,
    and immediately removes the local file and frees memory.
    """
    if tmp_dir is None:
        tmp_dir = Path(tempfile.gettempdir()) / "viu_shards"
    tmp_path = Path(tmp_dir) / stream_name
    tmp_path.mkdir(parents=True, exist_ok=True)

    file_name = f"train-{shard_id:05d}.parquet"
    out_file = tmp_path / file_name
    remote_path = f"{stream_name}/{file_name}"

    try:
        import pandas as pd
        df = pd.DataFrame(records)
        df.to_parquet(out_file, index=False, compression="snappy")
        del df
        gc.collect()
    except Exception as e:
        # Fallback to jsonl.gz if parquet/pandas is unavailable
        file_name = f"train-{shard_id:05d}.jsonl.gz"
        out_file = tmp_path / file_name
        remote_path = f"{stream_name}/{file_name}"
        with gzip.open(out_file, "wt", encoding="utf-8") as f:
            for r in records:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        gc.collect()

    size_mb = out_file.stat().st_size / (1024 * 1024)
    print(f"[shard] Serialized {file_name} with {len(records):,} rows ({size_mb:.2f} MB)")

    if dry_run:
        print(f"[dry-run] Skipped upload of {out_file} to {repo_id}/{remote_path}")
        out_file.unlink(missing_ok=True)
        return True

    # Upload to HF with retry and exponential backoff + 429 rate limit handler (Overnight resilient)
    retries = 999
    success = False
    for attempt in range(1, retries + 1):
        try:
            api.upload_file(
                path_or_fileobj=str(out_file),
                path_in_repo=remote_path,
                repo_id=repo_id,
                repo_type="dataset",
                commit_message=f"Upload {remote_path} ({len(records):,} rows, {size_mb:.1f} MB)"
            )
            print(f"[hf] Uploaded successfully: {remote_path} ({size_mb:.2f} MB)")
            success = True
            break
        except Exception as err:
            err_str = str(err)
            if "429" in err_str or "rate limit" in err_str.lower():
                wait_time = 120
                m = re.search(r"Retry after (\d+) seconds", err_str)
                if m:
                    wait_time = int(m.group(1)) + 5
                print(f"[rate-limit] HF 128 commits/hr limit hit. Waiting {wait_time}s before retry ({attempt}/{retries})...")
                time.sleep(wait_time)
            else:
                print(f"[warn] Upload attempt {attempt}/{retries} failed: {err}")
                time.sleep(4 * attempt)

    # Immediately delete local temporary shard to keep Kaggle disk clean
    out_file.unlink(missing_ok=True)
    if not success:
        print(f"[error] Failed to upload {remote_path} after {retries} retries!")
    return success


def run_collector(stream_name, repo_id, shard_mb=150, shard_rows=100_000,
                  resume=True, skip_rows=None, limit=None, dry_run=False, token=None):
    from huggingface_hub import HfApi

    api = None
    if not dry_run:
        if not token:
            token = get_hf_token()
        if not token:
            print("[error] HF token missing! Set HF_TOKEN environment variable or Kaggle secret.")
            sys.exit(1)
        api = HfApi(token=token)
        try:
            api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
            print(f"[hf] Verified dataset repo: {repo_id}")
        except Exception as e:
            print(f"[warn] create_repo warning: {e}")

    # Determine starting shard_id and rows to skip if resuming
    # FIX: old code assumed 100k rows/shard even for 1.6GB english shards (~3-4M rows),
    # causing massive duplication on resume. Now shard_rows is enforced, so
    # skip = (max_id+1) * shard_rows is correct. For legacy 1GB shards, require
    # explicit --skip_rows or manual backfill (see docs/HF_REPAIR_RUNBOOK.md).
    shard_id = 0
    actual_skip = 0
    if resume and not dry_run:
        next_id, max_id, count = get_existing_shards_info(api, repo_id, stream_name)
        if count > 0:
            shard_id = next_id
            if skip_rows is None:
                actual_skip = (max_id + 1) * shard_rows
                print(f"[resume] Found {count} existing shards in '{stream_name}' (up to train-{max_id:05d}.parquet).")
                print(f"[resume] Resuming at shard_id={shard_id} (train-{shard_id:05d}.parquet), auto-skipping ~{actual_skip:,} rows.")
                print(f"[resume] NOTE: legacy 1GB shards (english/hindi-big) used variable rows. "
                      f"If resuming those, pass explicit --skip_rows. See docs/HF_REPAIR_RUNBOOK.md")
            else:
                actual_skip = skip_rows
                print(f"[resume] Resuming at shard_id={shard_id}, skipping specified {actual_skip:,} rows.")
        else:
            actual_skip = skip_rows or 0
    else:
        actual_skip = skip_rows or 0

    # Target byte calculation: Snappy parquet ~2.7x for hindi text, ~1.6x for english.
    # FIX: old default 1000MB caused 1.6GB english shards -> HF viewer 500 + OOM.
    # New default 150MB + hard row cap 100k keeps shards 30-150MB, viewer-stable.
    # Conservative ratio 2.0 to avoid overshoot on english.
    target_raw_bytes = int(shard_mb * 1024 * 1024 * 2.0)
    if stream_name == "domains":
        # Domain samples are curated; use smaller chunk threshold to flush promptly
        target_raw_bytes = int(min(shard_mb, 150) * 1024 * 1024 * 2.7)

    print(f"\n=== Starting Collector: stream={stream_name}, repo={repo_id} ===")
    print(f"    Target shard size: ~{shard_mb} MB compressed | Start shard: train-{shard_id:05d}.parquet | Skip: {actual_skip:,} rows")

    seen_hashes = set()
    current_records = []
    current_raw_bytes = 0
    total_kept = 0
    start_time = time.time()
    last_log_time = start_time

    # Determine generator
    if stream_name == "hindi":
        gen = ((txt, src, "hindi", "general_hindi") for txt, src in stream_hindi_data(skip_rows=actual_skip))
    elif stream_name == "english":
        gen = ((txt, src, "english", "general_english") for txt, src in stream_english_data(skip_rows=actual_skip))
    elif stream_name == "hinglish":
        gen = ((txt, src, "hinglish", "general_hinglish") for txt, src in stream_hinglish_data(skip_rows=actual_skip))
    elif stream_name == "domains":
        gen = stream_domain_data()
    elif stream_name == "translation":
        gen = stream_translation_data(skip_rows=actual_skip)
    else:
        raise ValueError(f"Unknown stream: {stream_name}")

    for item in gen:
        txt, src, lang, dom = item[0], item[1], item[2], item[3]
        s = txt.strip()
        if not is_valid_line(s, lang):
            continue

        # In-flight stable-hash deduplication (FIX: hash() was seed-randomized)
        # NOTE: seen_hashes grows with total rows. For >10M-row runs on Kaggle,
        # memory can blow up. We cap at 5M entries then clear oldest via sampling
        # reset to keep RAM <2GB (small duplicate leak acceptable vs OOM).
        h = stable_hash(s)
        if h in seen_hashes:
            continue
        seen_hashes.add(h)
        if len(seen_hashes) > 5_000_000:
            print("[dedup] seen_hashes >5M, clearing to bound RAM (minor dup leak accepted)")
            seen_hashes.clear()

        row_bytes = len(s.encode('utf-8'))
        current_records.append({
            "text": s,
            "lang": lang,
            "source": src,
            "domain": dom
        })
        current_raw_bytes += row_bytes
        total_kept += 1

        # Periodic status log every 100,000 rows or every 30 seconds
        now = time.time()
        if total_kept % 100_000 == 0 or (now - last_log_time) >= 30.0:
            elapsed = max(now - start_time, 1)
            rate = total_kept / elapsed
            est_parquet_mb = (current_raw_bytes / 2.7) / (1024 * 1024)
            print(f"[progress] {total_kept:,} rows | Shard buffer: {len(current_records):,} rows (~{est_parquet_mb:.1f} MB est. parquet) | {rate:.1f} rows/s")
            last_log_time = now

        # Flush condition: Target MB reached OR explicit row cap reached
        hit_bytes = current_raw_bytes >= target_raw_bytes
        hit_rows = (shard_rows is not None and len(current_records) >= shard_rows)
        if hit_bytes or hit_rows:
            write_and_upload_shard(current_records, shard_id, stream_name, repo_id, api, dry_run=dry_run)
            shard_id += 1
            current_records = []
            current_raw_bytes = 0
            gc.collect()

        if limit and total_kept >= limit:
            print(f"[limit] Reached limit of {limit:,} rows.")
            break

    # Flush remaining records
    if current_records:
        write_and_upload_shard(current_records, shard_id, stream_name, repo_id, api, dry_run=dry_run)
        shard_id += 1
        current_records = []
        gc.collect()

    elapsed = max(time.time() - start_time, 1)
    print(f"=== Collector Finished: {total_kept:,} total rows in this session ({elapsed:.1f}s, {total_kept/elapsed:.1f} rows/s) ===")


def main():
    parser = argparse.ArgumentParser(description="Kaggle High-Throughput Cloud Stream Collector")
    parser.add_argument("--stream", type=str, required=True,
                        choices=["hindi", "english", "hinglish", "domains", "translation", "all"],
                        help="Data stream to collect (or 'all' for sequential run)")
    parser.add_argument("--repo", type=str, default="ViuAI/viu-mini-raw-pretrain", help="HF dataset repo ID")
    parser.add_argument("--shard_mb", type=int, default=150, help="Target compressed Parquet shard size in MB (default 150 MB, viewer-stable; old 1000 MB caused 1.6GB shards + viewer 500)")
    parser.add_argument("--shard_rows", type=int, default=100000, help="Row cap per shard (default 100k; whichever hits first: bytes or rows)")
    parser.add_argument("--no-resume", action="store_true", help="Do not auto-resume; start from shard 0 and row 0")
    parser.add_argument("--skip_rows", type=int, default=None, help="Explicitly skip N rows before streaming")
    parser.add_argument("--limit", type=int, default=None, help="Optional max rows to collect")
    parser.add_argument("--dry-run", action="store_true", help="Run without uploading to HF")
    parser.add_argument("--token", type=str, default=None, help="HF API token")
    args = parser.parse_args()

    resume = not args.no_resume

    if args.stream == "all":
        # Prioritize Hinglish and English so they accumulate immediately alongside Hindi
        streams = ["domains", "translation", "hinglish", "english", "hindi"]
        for s in streams:
            print(f"\n==========================================")
            print(f"      STARTING STREAM: {s.upper()}")
            print(f"==========================================")
            # Smart check: if domains or translation are already fully uploaded, proceed immediately
            if resume and not args.dry_run:
                try:
                    from huggingface_hub import HfApi
                    check_api = HfApi(token=args.token or get_hf_token())
                    _, _, existing_cnt = get_existing_shards_info(check_api, args.repo, s)
                    if s == "domains" and existing_cnt >= 2:
                        print(f"[all] Stream 'domains' already fully uploaded ({existing_cnt} shards). Skipping to next...")
                        continue
                    if s == "translation" and existing_cnt >= 13:
                        print(f"[all] Stream 'translation' already fully uploaded ({existing_cnt} shards). Skipping to next...")
                        continue
                except Exception as e_chk:
                    print(f"[warn] Existing check error for {s}: {e_chk}")

            run_collector(
                stream_name=s,
                repo_id=args.repo,
                shard_mb=args.shard_mb,
                shard_rows=args.shard_rows,
                resume=resume,
                skip_rows=args.skip_rows,
                limit=args.limit,
                dry_run=args.dry_run,
                token=args.token
            )
    else:
        run_collector(
            stream_name=args.stream,
            repo_id=args.repo,
            shard_mb=args.shard_mb,
            shard_rows=args.shard_rows,
            resume=resume,
            skip_rows=args.skip_rows,
            limit=args.limit,
            dry_run=args.dry_run,
            token=args.token
        )


if __name__ == "__main__":
    main()