"""
Repair script 2: Re-shard oversized files (english 37x1.6GB, hindi-big 13x1.15GB)
Problem: 1GB shard_mb setting created viewer-breaking 1.6GB parquets (HTTP 500 + OOM).
Fix (Kaggle, streaming, delete-after-upload):
  1. Download one big shard -> read in batches -> re-write as ~150MB / 100k-row shards
     under staging folder (english_fixed/, hindi_fixed/).
  2. Verify row counts, then delete originals and rename staging -> live in 2 commits.
  3. NEVER delete before verify ( HF has no trash for datasets).

Run:
  python data/scripts/repair_reshard_large.py --repo ViuAI/viu-mini-raw-pretrain --stream english --dry-run
  python data/scripts/repair_reshard_large.py --repo ViuAI/viu-mini-raw-pretrain --stream english
  python data/scripts/repair_reshard_large.py --repo ViuAI/viu-mini-raw-pretrain --stream hindi_big

Note: 1.6GB download per file needs ~8GB disk+RAM; process ONE file at a time,
upload splits immediately, delete local. Total time ~2-4 hrs on Kaggle.
"""
import argparse
import gc
import tempfile
from pathlib import Path

TARGET_MB = 150
TARGET_ROWS = 100_000

def shard_plan(stream: str):
    if stream == "english":
        return [f"english/train-{i:05d}.parquet" for i in range(37)]
    if stream == "hindi_big":
        return [f"hindi/train-{i:05d}.parquet" for i in list(range(434, 447))]
    raise ValueError("stream must be english or hindi_big")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="ViuAI/viu-mini-raw-pretrain")
    ap.add_argument("--stream", required=True, choices=["english", "hindi_big"])
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--token", default=None)
    ap.add_argument("--limit-files", type=int, default=None)
    args = ap.parse_args()

    from huggingface_hub import HfApi, hf_hub_download
    from data.scripts.kaggle_collector import get_hf_token
    import pandas as pd

    token = args.token or get_hf_token()
    api = HfApi(token=token) if not args.dry_run else None
    files = shard_plan(args.stream)
    if args.limit_files:
        files = files[:args.limit_files]
    stage = "english_fixed" if args.stream == "english" else "hindi_fixed"
    print(f"[reshard] {args.stream}: {len(files)} big files -> staging {stage}/ (~{TARGET_MB}MB / {TARGET_ROWS} rows each)")

    if args.dry_run:
        for f in files:
            print(f"  would split {f}")
        print("[dry-run] no download/upload. Run without --dry-run on Kaggle.")
        return

    tmp = Path(tempfile.mkdtemp(prefix="viu_reshard_"))
    out_id = 0
    # resume staging counter
    existing = set(api.list_repo_files(repo_id=args.repo, repo_type="dataset"))
    staged = sorted([f for f in existing if f.startswith(stage + "/")])
    if staged:
        nums = [int(f.split("-")[-1].split(".")[0]) for f in staged]
        out_id = max(nums) + 1
        print(f"[resume] staging has {len(staged)} shards, next out_id={out_id}")

    for big in files:
        print(f"\n[dl] {big} ...")
        local = hf_hub_download(repo_id=args.repo, filename=big, repo_type="dataset")
        import pyarrow.parquet as pq
        pf = pq.ParquetFile(local)
        print(f"  rows={pf.metadata.num_rows:,}, row_groups={pf.metadata.num_row_groups}")
        buf = []
        def flush():
            nonlocal out_id, buf
            if not buf:
                return
            df = pd.DataFrame(buf, columns=["text", "lang", "source", "domain"])
            out_local = tmp / f"train-{out_id:05d}.parquet"
            df.to_parquet(out_local, index=False, compression="snappy")
            mb = out_local.stat().st_size / 1024 / 1024
            remote = f"{stage}/train-{out_id:05d}.parquet"
            print(f"  [up] {remote} {len(df):,} rows {mb:.1f}MB")
            api.upload_file(path_or_fileobj=str(out_local), path_in_repo=remote,
                            repo_id=args.repo, repo_type="dataset",
                            commit_message=f"reshard {big} -> {remote} ({len(df):,} rows)")
            out_local.unlink(missing_ok=True)
            out_id += 1
            buf = []
            gc.collect()
        for batch in pf.iter_batches(batch_size=20000):
            for r in batch.to_pylist():
                # keep only unified cols (drop stray cols from legacy files)
                buf.append({k: r.get(k, "") for k in ("text", "lang", "source", "domain")})
                if len(buf) >= TARGET_ROWS:
                    flush()
        flush()
        print(f"[ok] split {big} done")
    print(f"\n[done] staged {out_id} shards in {stage}/")
    print("NEXT (manual, after verifying row counts in viewer):")
    print(f"  1. Delete originals: {files[0]} ... ({len(files)} files)")
    print(f"  2. Rename {stage}/* -> original folder via 2 commits (copy+delete).")
    print("  See docs/HF_REPAIR_RUNBOOK.md for exact delete/rename commands.")

if __name__ == "__main__":
    main()