#!/usr/bin/env python3
"""
Simulate the exact dataset loading process from finetune_qwenvl.py on single CPU.
This reproduces the training initialization to debug empty dataset issue.

Usage:
    python3 scripts/test_dataset_load_single_cpu.py \
        --wdsmeta-path /path/to/metadata.wdsmeta.json \
        --simulate-rank 0 \
        --simulate-world-size 8
"""

import argparse
import os
import sys
from pathlib import Path

# Add LMFlow to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

def main():
    parser = argparse.ArgumentParser(description='Test dataset loading on single CPU')
    parser.add_argument('--wdsmeta-path', required=True, help='Path to .wdsmeta.json file')
    parser.add_argument('--simulate-rank', type=int, default=0, help='Simulate this rank')
    parser.add_argument('--simulate-world-size', type=int, default=8, help='Simulate this world size')
    parser.add_argument('--max-samples', type=int, default=10, help='Max samples to iterate')
    
    args = parser.parse_args()
    
    # Set environment variables to simulate distributed training
    os.environ['RANK'] = str(args.simulate_rank)
    os.environ['LOCAL_RANK'] = str(args.simulate_rank)
    os.environ['WORLD_SIZE'] = str(args.simulate_world_size)
    os.environ['NNODES'] = '1'
    os.environ['NODE_RANK'] = '0'
    
    print(f"=" * 80)
    print(f"Simulating Rank {args.simulate_rank}/{args.simulate_world_size} on single CPU")
    print(f"Metadata path: {args.wdsmeta_path}")
    print(f"=" * 80)
    print()
    
    # Import after setting env vars
    from lmflow.args import DatasetArguments
    from lmflow.datasets.webdataset_multi_modal_dataset import CustomWebDatasetMultiModalDataset
    
    # Create minimal DatasetArguments (same as training would do)
    print("[Test] Creating DatasetArguments...")
    data_args = DatasetArguments(
        dataset_path=args.wdsmeta_path,
    )
    print(f"[Test] DatasetArguments created: dataset_path={data_args.dataset_path}")
    print()
    
    # Create the dataset (this is where the issue happens in training)
    print("[Test] Creating CustomWebDatasetMultiModalDataset...")
    print("[Test] This should trigger all the [WebDataset] debug prints...")
    print()
    
    try:
        dataset = CustomWebDatasetMultiModalDataset(args.wdsmeta_path, data_args)
        print()
        print("[Test] ✅ Dataset created successfully!")
        print(f"[Test] Dataset length: {len(dataset)}")
        print()
        
        if len(dataset) == 0:
            print("[Test] ❌ ERROR: Dataset is EMPTY!")
            print("[Test] This reproduces the training failure.")
        else:
            print(f"[Test] ✅ Dataset has {len(dataset)} samples")
            print()
            
            # Try to iterate a few samples
            print(f"[Test] Attempting to load first {min(args.max_samples, len(dataset))} samples...")
            for i in range(min(args.max_samples, len(dataset))):
                try:
                    sample = dataset[i]
                    if 'messages' in sample:
                        num_messages = len(sample['messages'])
                        print(f"[Test] Sample {i}: {num_messages} messages")
                        # Print first message structure
                        if num_messages > 0:
                            first_msg = sample['messages'][0]
                            if 'content' in first_msg:
                                content_types = [c.get('type', '?') for c in first_msg['content']]
                                print(f"[Test]   First message content types: {content_types}")
                    else:
                        print(f"[Test] Sample {i}: No 'messages' key found")
                except Exception as e:
                    print(f"[Test] ❌ Error loading sample {i}: {e}")
                    import traceback
                    traceback.print_exc()
                    break
            
            print()
            print("[Test] ✅ Successfully loaded samples!")
    
    except Exception as e:
        print()
        print(f"[Test] ❌ ERROR: Failed to create dataset!")
        print(f"[Test] Error: {e}")
        import traceback
        traceback.print_exc()
        print()
        print("[Test] This might explain why training fails with empty dataset.")
        sys.exit(1)
    
    print()
    print("=" * 80)
    print("Test completed successfully!")
    print("=" * 80)


if __name__ == '__main__':
    main()

