"""
LUNA — Dataset Fetcher
======================
Downloads the tokenized litdata dataset from either:
  - HuggingFace Hub (recommended, free, fast)
  - Google Drive (direct link, requires gdown)

Usage:
    python fetch_data.py --source huggingface --hf_repo YourName/LUNA-pretrain-data --out_dir /workspace/data
    python fetch_data.py --source gdrive     --gdrive_id <FILE_OR_FOLDER_ID>       --out_dir /workspace/data
    python fetch_data.py --source local      --local_path Base/data/litdata_pretrain_final  --out_dir /workspace/data

After running, pass --data_path /workspace/data/litdata_pretrain_final to train.py
"""

import os
import sys
import json
import shutil
import argparse
from pathlib import Path


# ─── HuggingFace Download ─────────────────────────────────────────────────────

def download_huggingface(repo_id: str, out_dir: Path, hf_token: str = None):
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("  Installing huggingface_hub...")
        os.system(f"{sys.executable} -m pip install -q huggingface_hub")
        from huggingface_hub import snapshot_download

    print(f"  Downloading from HuggingFace: {repo_id}")
    out_dir.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        local_dir=str(out_dir),
        token=hf_token,
        ignore_patterns=["*.md", ".gitattributes"],
    )
    print(f"  Downloaded to: {out_dir}")

    # Auto-extract any zip files found in the download
    _extract_zips(out_dir)

    # If index.json landed in a subdirectory, move contents up
    _flatten_to_root(out_dir)

    _verify(out_dir)


# ─── Google Drive Download ────────────────────────────────────────────────────

def download_gdrive(gdrive_id: str, out_dir: Path):
    try:
        import gdown
    except ImportError:
        print("  Installing gdown...")
        os.system(f"{sys.executable} -m pip install -q gdown")
        import gdown

    out_dir.mkdir(parents=True, exist_ok=True)
    # Try as folder first, then single file
    url = f"https://drive.google.com/drive/folders/{gdrive_id}"
    print(f"  Attempting GDrive folder download: {gdrive_id}")
    try:
        gdown.download_folder(url=url, output=str(out_dir), quiet=False, use_cookies=False)
    except Exception as e:
        print(f"  Folder download failed ({e}), trying single file...")
        url = f"https://drive.google.com/uc?id={gdrive_id}"
        dest = out_dir / "data.zip"
        gdown.download(url, str(dest), quiet=False)
        if dest.suffix == ".zip":
            print("  Extracting zip...")
            import zipfile
            with zipfile.ZipFile(dest) as z:
                z.extractall(out_dir)
            dest.unlink()
    print(f"  Downloaded to: {out_dir}")
    _verify(out_dir)


# ─── Local Copy ───────────────────────────────────────────────────────────────

def copy_local(local_path: str, out_dir: Path):
    src = Path(local_path)
    if not src.exists():
        raise FileNotFoundError(f"Local path not found: {src}")
    if out_dir.resolve() == src.resolve():
        print(f"  Source == destination, no copy needed.")
        _verify(out_dir)
        return
    print(f"  Copying {src} → {out_dir}")
    if out_dir.exists():
        shutil.rmtree(out_dir)
    shutil.copytree(src, out_dir)
    print(f"  Copied to: {out_dir}")
    _verify(out_dir)


# ─── Zip Extraction & Flattening ──────────────────────────────────────────────

def _extract_zips(data_dir: Path):
    """Find and extract all .zip files in data_dir, then delete the zips."""
    import zipfile
    zips = list(data_dir.glob("*.zip"))
    if not zips:
        return
    for zf in zips:
        print(f"  Extracting {zf.name} ...")
        with zipfile.ZipFile(zf) as z:
            z.extractall(data_dir)
        zf.unlink()
        print(f"  Removed {zf.name}")


def _flatten_to_root(data_dir: Path):
    """If index.json is nested (e.g. data_dir/a/b/index.json),
    move everything from that subfolder up to data_dir."""
    if (data_dir / "index.json").exists():
        return  # already at root
    candidates = list(data_dir.glob("**/index.json"))
    if len(candidates) != 1:
        return  # ambiguous or not found
    sub = candidates[0].parent
    print(f"  Moving contents from {sub.relative_to(data_dir)}/ up to {data_dir.name}/ ...")
    for item in sub.iterdir():
        dest = data_dir / item.name
        if dest.exists():
            if dest.is_dir():
                shutil.rmtree(dest)
            else:
                dest.unlink()
        shutil.move(str(item), str(dest))
    # Remove the now-empty nested directories
    # Walk up from sub to data_dir, removing empty dirs
    while sub != data_dir:
        try:
            sub.rmdir()
        except OSError:
            break
        sub = sub.parent


# ─── Verify ───────────────────────────────────────────────────────────────────

def _verify(data_dir: Path):
    idx_path = data_dir / "index.json"
    if not idx_path.exists():
        # Search one level deeper
        found = list(data_dir.glob("**/index.json"))
        if found:
            print(f"  Note: index.json found at {found[0]}, not root. Check your --out_dir.")
        else:
            print(f"  WARNING: index.json NOT found in {data_dir}")
        return

    with open(idx_path) as f:
        idx = json.load(f)
    chunks = idx.get("chunks", [])
    total_tokens = sum(c.get("dim", 0) for c in chunks)
    present = sum(1 for c in chunks if (data_dir / c["filename"]).exists())
    missing = len(chunks) - present

    print(f"\n  Dataset verified:")
    print(f"  Chunks declared : {len(chunks)}")
    print(f"  Chunks on disk  : {present}")
    print(f"  Missing chunks  : {missing}")
    print(f"  Total tokens    : {total_tokens:,}")
    if missing > 0:
        print(f"  WARNING: {missing} chunk(s) missing — training will error on those blocks!")
    else:
        print(f"  All chunks present. Ready to train.")


# ─── Args ─────────────────────────────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="LUNA dataset fetcher")
    p.add_argument("--source", choices=["huggingface", "gdrive", "local"], required=True)
    p.add_argument("--out_dir", type=str, default="/workspace/data/litdata_pretrain_final",
                   help="Where to save the dataset")
    p.add_argument("--hf_repo", type=str, default="",
                   help="HuggingFace dataset repo ID (e.g. YourName/LUNA-pretrain-data)")
    p.add_argument("--hf_token", type=str, default=os.environ.get("HF_TOKEN", ""),
                   help="HuggingFace token (or set HF_TOKEN env var)")
    p.add_argument("--gdrive_id", type=str, default="",
                   help="Google Drive file/folder ID")
    p.add_argument("--local_path", type=str, default="Base/data/litdata_pretrain_final",
                   help="Local path to the dataset (for local source)")
    return p.parse_args()


if __name__ == "__main__":
    args = parse_args()
    out = Path(args.out_dir)

    if args.source == "huggingface":
        if not args.hf_repo:
            print("ERROR: --hf_repo required for HuggingFace source")
            sys.exit(1)
        download_huggingface(args.hf_repo, out, hf_token=args.hf_token or None)

    elif args.source == "gdrive":
        if not args.gdrive_id:
            print("ERROR: --gdrive_id required for GDrive source")
            sys.exit(1)
        download_gdrive(args.gdrive_id, out)

    elif args.source == "local":
        copy_local(args.local_path, out)