"""
Push pretrain dataset (40:30:30) to HF Datasets.
Supports: local data/raw/*.txt -> HF repo ViuAI/viu-mini-pretrain-40-30-30

Usage:
  python push_dataset_to_hf.py --repo ViuAI/viu-mini-pretrain-40-30-30 --dry-run
  python push_dataset_to_hf.py --repo ViuAI/viu-mini-pretrain-40-30-30
"""
import argparse
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
RAW = ROOT / "data" / "raw"
EVAL = ROOT / "data" / "eval"

def lang_tag(name: str) -> str:
    n = name.lower()
    if "hinglish" in n:
        return "hinglish"
    if "hindi" in n:
        return "hindi"
    if "english" in n:
        return "english"
    return "other"


def collect():
    files = []
    if RAW.exists():
        files += sorted(RAW.glob("*.txt"))
    return files

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", required=True, help="HF dataset repo id")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    files = collect()
    print(f"[data] found {len(files)} raw files:")
    for f in files:
        print(" ", f.name, f.stat().st_size, "bytes")
    if EVAL.exists():
        for f in sorted(EVAL.glob("*.txt")):
            print(" [eval]", f.name, f.stat().st_size, "bytes")

    if args.dry_run:
        print(f"[dry-run] would push to {args.repo} (no upload)")
        # also show how train.py will load it:
        print("[info] load via: datasets.load_dataset('{}', streaming=True)".format(args.repo))
        return

    try:
        from huggingface_hub import HfApi
    except ImportError as e:
        print(f"[error] missing dep: {e}. pip install huggingface_hub")
        return

    api = HfApi()
    try:
        api.create_repo(repo_id=args.repo, repo_type="dataset", exist_ok=True)
        print(f"[hf] dataset repo ensured: {args.repo}")
    except Exception as e:
        print(f"[warn] create_repo: {e}")

    # Build HF dataset from raw txt (Hinglish/Hindi/English tags) — RAM-safe sharded jsonl
    import json
    import tempfile
    shard_rows = 100_000
    tmpdir = Path(tempfile.mkdtemp(prefix="viu_ds_"))
    shard_id, n_rows, f = 0, 0, None
    def new_shard():
        nonlocal shard_id, f
        p = tmpdir / f"train-{shard_id:05d}.jsonl"
        f = open(p, "w", encoding="utf-8")
        return p
    cur = new_shard()
    print(f"[data] writing shards to {tmpdir} ...")
    total = 0
    for fp in files:
        tag = lang_tag(fp.name)
        for line in open(fp, encoding="utf-8"):
            line = line.strip()
            if line:
                if f is None:
                    cur = new_shard()
                f.write(json.dumps({"text": line, "lang": tag, "source": fp.name}, ensure_ascii=False) + "\n")
                n_rows += 1
                total += 1
                if n_rows >= shard_rows:
                    f.close()
                    f = None
                    n_rows = 0
                    shard_id += 1
    if f is not None:
        f.close()
    shards = sorted(tmpdir.glob("*.jsonl"))

    if total == 0:
        print("[error] no rows to push — data/raw khaali hai. Real 4-5B data Phase-1 me aayega.")
        return

    print(f"[hf] pushing {total} rows ({len(shards)} shards) to {args.repo} ...")
    api.upload_folder(
        folder_path=str(tmpdir),
        repo_id=args.repo,
        repo_type="dataset",
        commit_message="upload pretrain 40:30:30 sharded jsonl",
        path_in_repo="data",
    )
    print(f"[ok] dataset pushed: https://huggingface.co/datasets/{args.repo}")
    print("[info] train.py me: --hf_dataset", args.repo)

if __name__ == "__main__":
    main()