"""
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 the intended dataset up
    _flatten_to_root(out_dir, preferred_name=out_dir.name)

    _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}")
    _flatten_to_root(out_dir, preferred_name=out_dir.name)
    _verify(out_dir)


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

def _nearest_existing_ancestor(path: Path):
    current = path
    while True:
        if current.exists():
            return current
        if current.parent == current:
            return None
        current = current.parent


def _discover_local_dataset(local_path: str):
    requested = Path(local_path)
    preferred_name = requested.name.lower()
    roots = []

    for candidate in [requested, _nearest_existing_ancestor(requested), Path.cwd(), Path("/workspace")]:
        if candidate is None:
            continue
        if candidate.exists():
            candidate = candidate.resolve()
        if candidate not in roots:
            roots.append(candidate)

    found = []
    seen = set()
    for root in roots:
        if not root.exists() or not root.is_dir():
            continue
        for index_path in root.glob("**/index.json"):
            parent = index_path.parent.resolve()
            if parent in seen:
                continue
            seen.add(parent)
            summary = _read_index_summary(index_path)
            if summary["chunks"] <= 0:
                continue
            name = parent.name.lower()
            exact = int(bool(preferred_name and name == preferred_name))
            contains = int(bool(preferred_name and preferred_name in name))
            found.append((exact, contains, summary["tokens"], summary["chunks"], parent, summary))

    if not found:
        return None, None

    found.sort(reverse=True, key=lambda item: (item[0], item[1], item[2], item[3]))
    best = found[0]
    return best[4], best[5]


def copy_local(local_path: str, out_dir: Path):
    src = Path(local_path)
    if not src.exists():
        discovered, summary = _discover_local_dataset(local_path)
        if discovered is None:
            raise FileNotFoundError(f"Local path not found: {src}")
        print(
            f"  Requested local path not found: {src}\n"
            f"  Auto-selected local dataset: {discovered} "
            f"({summary['chunks']} chunks, {summary['tokens']:,} tokens)"
        )
        src = discovered
    if out_dir.resolve() == src.resolve():
        print(f"  Source == destination, no copy needed.")
        _flatten_to_root(out_dir, preferred_name=out_dir.name)
        _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}")
    _flatten_to_root(out_dir, preferred_name=out_dir.name)
    _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 _read_index_summary(index_path: Path):
    try:
        with open(index_path, encoding="utf-8") as f:
            idx = json.load(f)
    except Exception:
        return {"chunks": 0, "tokens": 0}
    chunks = idx.get("chunks", [])
    total_tokens = sum(c.get("dim", 0) for c in chunks)
    return {"chunks": len(chunks), "tokens": total_tokens}


def _pick_dataset_subdir(candidates, preferred_name: str | None = None):
    """Pick the most likely dataset subdir from multiple nested index.json files.

    Priority:
    1. Parent folder name exactly matches preferred_name
    2. Parent folder name contains preferred_name
    3. Highest token count
    """
    ranked = []
    preferred_name = (preferred_name or "").lower()

    for index_path in candidates:
        parent = index_path.parent
        summary = _read_index_summary(index_path)
        name = parent.name.lower()
        exact = int(bool(preferred_name and name == preferred_name))
        contains = int(bool(preferred_name and preferred_name in name))
        ranked.append((exact, contains, summary["tokens"], summary["chunks"], parent, summary))

    ranked.sort(reverse=True, key=lambda item: (item[0], item[1], item[2], item[3]))
    return ranked[0][4], ranked[0][5]


def _flatten_to_root(data_dir: Path, preferred_name: str | None = None):
    """If index.json is nested, move the intended dataset subfolder up to data_dir."""
    if (data_dir / "index.json").exists():
        return

    candidates = list(data_dir.glob("**/index.json"))
    if not candidates:
        return

    if len(candidates) == 1:
        sub = candidates[0].parent
        summary = _read_index_summary(candidates[0])
    else:
        sub, summary = _pick_dataset_subdir(candidates, preferred_name=preferred_name)
        print(
            f"  Multiple nested datasets found; selected {sub.relative_to(data_dir)}/ "
            f"({summary['chunks']} chunks, {summary['tokens']:,} tokens)"
        )

    if sub == data_dir:
        return

    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)