"""Download training datasets from HuggingFace Hub.

Usage:
    python3 scripts/download_datasets.py --config configs/scale_1.3b.yaml --stage all
    python3 scripts/download_datasets.py --config configs/scale_1.3b.yaml --stage stage1
    python3 scripts/download_datasets.py --config configs/scale_1.3b.yaml --stage stage2 --dummy
"""

import argparse
import os
import sys
import yaml
import json

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def download_hf_dataset(name: str, source: str, num_samples: int,
                        output_dir: str, streaming: bool = True):
    """Download a dataset from HuggingFace Hub.

    Args:
        name: Dataset identifier
        source: HuggingFace dataset path
        num_samples: Max samples to download
        output_dir: Where to save
        streaming: Use streaming mode (recommended for large datasets)
    """
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f"{name}.jsonl")

    if os.path.exists(output_file):
        # Count existing lines
        with open(output_file) as f:
            existing = sum(1 for _ in f)
        if existing >= num_samples:
            print(f"  {name}: Already downloaded ({existing} samples)")
            return existing

    try:
        from datasets import load_dataset
        print(f"  {name}: Downloading from {source}...")

        if streaming:
            ds = load_dataset(source, split="train", streaming=True, trust_remote_code=True)
        else:
            ds = load_dataset(source, split="train", trust_remote_code=True)

        count = 0
        with open(output_file, "w") as f:
            for item in ds:
                if count >= num_samples:
                    break

                # Extract text fields (skip image data for now — download separately)
                record = {}
                for key in ["question", "text", "caption", "answer", "multiple_choice_answer",
                           "answers", "conversations", "image_id", "id"]:
                    if key in item:
                        val = item[key]
                        if isinstance(val, (str, int, float, bool, list)):
                            record[key] = val
                        elif hasattr(val, 'tolist'):
                            record[key] = val.tolist()

                f.write(json.dumps(record) + "\n")
                count += 1

                if count % 10000 == 0:
                    print(f"    {name}: {count}/{num_samples} samples...")

        print(f"  {name}: Downloaded {count} samples to {output_file}")
        return count

    except Exception as e:
        print(f"  {name}: Download failed ({e}), will use dummy data")
        return 0


def generate_dummy_data(name: str, num_samples: int, output_dir: str):
    """Generate dummy training data for testing."""
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f"{name}.jsonl")

    with open(output_file, "w") as f:
        for i in range(min(num_samples, 500)):
            record = {
                "question": f"Sample question {i} for {name}",
                "answer": f"Sample answer {i}",
                "id": i,
            }
            f.write(json.dumps(record) + "\n")

    print(f"  {name}: Generated {min(num_samples, 500)} dummy samples")


def main():
    parser = argparse.ArgumentParser(description="Download ArcisVLM training datasets")
    parser.add_argument("--config", default="configs/scale_1.3b.yaml")
    parser.add_argument("--stage", default="all", choices=["all", "stage1", "stage2", "stage3"])
    parser.add_argument("--output_dir", default="data/downloads")
    parser.add_argument("--dummy", action="store_true", help="Generate dummy data only")
    args = parser.parse_args()

    with open(args.config) as f:
        config = yaml.safe_load(f)

    stages = []
    if args.stage in ["all", "stage1"]:
        stages.append(("stage1", config.get("data", {}).get("train_datasets", {}).get("stage1", [])))
    if args.stage in ["all", "stage2"]:
        stages.append(("stage2", config.get("data", {}).get("train_datasets", {}).get("stage2", [])))
    if args.stage in ["all", "stage3"]:
        stages.append(("stage3", config.get("data", {}).get("train_datasets", {}).get("stage3", [])))

    for stage_name, datasets in stages:
        print(f"\n{'='*50}")
        print(f"Downloading {stage_name} datasets")
        print(f"{'='*50}")

        stage_dir = os.path.join(args.output_dir, stage_name)

        for ds_config in datasets:
            name = ds_config["name"]
            source = ds_config.get("source", name)
            num_samples = ds_config.get("samples", 1000)

            if args.dummy:
                generate_dummy_data(name, num_samples, stage_dir)
            else:
                count = download_hf_dataset(name, source, num_samples, stage_dir)
                if count == 0:
                    print(f"    Falling back to dummy data for {name}")
                    generate_dummy_data(name, num_samples, stage_dir)

    print(f"\n{'='*50}")
    print("Download complete!")
    # Show disk usage
    total = 0
    for root, dirs, files in os.walk(args.output_dir):
        for f in files:
            total += os.path.getsize(os.path.join(root, f))
    print(f"Total disk usage: {total / 1e9:.2f} GB")
    print(f"{'='*50}")


if __name__ == "__main__":
    main()