#!/usr/bin/env python
# coding=utf-8
"""
Test script for WebDataset Multi Modal Dataset loading.
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from lmflow.datasets.webdataset_multi_modal_dataset import CustomWebDatasetMultiModalDataset
from lmflow.args import DatasetArguments

wdsmeta_path = '/lustre/fsw/portfolios/nvr/users/ligengz/datasets/vila-webds-meta/--home--ligengz--nvr_elm_llm--dataset--coyo-25m-vila.wdsmeta.json'

print("Creating dataset...")
data_args = DatasetArguments(dataset_path=wdsmeta_path)
dataset = CustomWebDatasetMultiModalDataset(wdsmeta_path, data_args)

print(f"\nDataset length: {len(dataset)}")

# Test loading a few samples
print("\nTesting sample loading...")
for i in range(min(3, len(dataset))):
    print(f"\n--- Sample {i} ---")
    try:
        sample = dataset[i]
        messages = sample['messages']
        print(f"Number of messages: {len(messages)}")
        for j, msg in enumerate(messages):
            print(f"  Message {j}: role={msg['role']}, content_types={[c.get('type') for c in msg['content']]}")
            # Check if image is present
            for content in msg['content']:
                if content.get('type') == 'image':
                    img = content.get('image')
                    if img is not None:
                        print(f"    Image: {type(img)}, size={img.size if hasattr(img, 'size') else 'N/A'}")
    except Exception as e:
        print(f"Error loading sample {i}: {e}")
        import traceback
        traceback.print_exc()

print("\nTest completed!")
