# Prompt: 给我生成一些test script，能单独在纯cpu的环境下 iterate这个数据集，不跑训练
#!/usr/bin/env python3
"""CPU-only sanity checker for CustomWebDatasetMultiModalDataset.

This script instantiates the dataset, iterates over a limited number of
samples, and prints basic diagnostics so you can verify that the metadata
points to accessible shards before launching a full training job.
"""

import argparse
import os
from pathlib import Path
from typing import Any, Dict

import torch
from torch.utils.data import DataLoader

from lmflow.args import DatasetArguments
from lmflow.datasets.webdataset_multi_modal_dataset import (
    CustomWebDatasetMultiModalDataset,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Iterate WebDataset multimodal dataset on CPU."
    )
    parser.add_argument(
        "--wdsmeta-path",
        required=True,
        type=Path,
        help="Path to the .wdsmeta.json file.",
    )
    parser.add_argument(
        "--max-batches",
        type=int,
        default=2,
        help="Number of DataLoader batches to iterate (default: 2).",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=2,
        help="Batch size for the sanity-check DataLoader (default: 2).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=0,
        help="Number of worker processes for DataLoader (default: 0 for CPU-only).",
    )
    parser.add_argument(
        "--env-path-map",
        type=str,
        default=None,
        help=(
            "Optional LMFLOW_WDS_PATH_MAP override, e.g. "
            "'/home/ligengz/nvr_elm_llm/dataset=/lustre/fsw/portfolios/nvr/users/ligengz/datasets'"
        ),
    )
    parser.add_argument(
        "--print-images",
        action="store_true",
        help="Print extra diagnostics about image tensors (may be verbose).",
    )
    return parser.parse_args()


def make_dataset_args(meta_path: Path) -> DatasetArguments:
    return DatasetArguments(
        dataset_path=str(meta_path),
        dataset_name="customized",
        is_custom_dataset=True,
    )


def main() -> None:
    args = parse_args()
    meta_path = args.wdsmeta_path.expanduser().resolve()
    if args.env_path_map:
        os.environ["LMFLOW_WDS_PATH_MAP"] = args.env_path_map
    data_args = make_dataset_args(meta_path)

    dataset = CustomWebDatasetMultiModalDataset(str(meta_path), data_args)
    print(f"[sanity] dataset length: {len(dataset)} samples")

    if len(dataset) == 0:
        print("[sanity] dataset reported zero samples; aborting.")
        return

    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        collate_fn=lambda batch: batch,
        pin_memory=False,
    )

    total_items = 0
    for batch_idx, batch in enumerate(loader):
        print(f"[sanity] batch {batch_idx}: size={len(batch)}")
        for item_idx, item in enumerate(batch):
            show_sample(item, args.print_images)
            total_items += 1
        if args.max_batches > 0 and batch_idx + 1 >= args.max_batches:
            break

    print(f"[sanity] inspected {total_items} samples across {batch_idx + 1} batches.")


def show_sample(sample: Dict[str, Any], include_images: bool) -> None:
    messages = sample.get("messages", [])
    print(f"  - sample with {len(messages)} messages")
    for msg in messages:
        role = msg.get("role")
        content = msg.get("content", [])
        text_snippets = [
            c.get("text", "") for c in content if c.get("type") == "text"
        ]
        has_image = any(c.get("type") == "image" for c in content)
        preview = text_snippets[0][:80] + ("..." if len(text_snippets[0]) > 80 else "")
        print(
            f"    role={role:>8} text='{preview}' "
            f"{'[image]' if has_image else ''}"
        )
        if include_images and has_image:
            for c in content:
                if c.get("type") == "image":
                    img = c.get("image")
                    if hasattr(img, "size"):
                        print(f"      image size={img.size} mode={img.mode}")


if __name__ == "__main__":
    main()

