#!/usr/bin/env python3
"""
Exact replication of CustomWebDatasetMultiModalDataset initialization on single CPU.
This uses the ACTUAL CustomWebDatasetMultiModalDataset class to reproduce the training issue.

Usage:
    # Simulate rank 0 out of 8
    python3 scripts/test_exact_dataset_init.py \
        --wdsmeta-path /lustre/fsw/portfolios/nvr/users/ligengz/datasets/vila-webds-meta/--home--ligengz--nvr_elm_llm--dataset--coyo-25m-vila.wdsmeta.json \
        --rank 0 --world-size 8
    
    # Simulate other ranks
    python3 scripts/test_exact_dataset_init.py \
        --wdsmeta-path /lustre/fsw/portfolios/nvr/users/ligengz/datasets/vila-webds-meta/--home--ligengz--nvr_elm_llm--dataset--coyo-25m-vila.wdsmeta.json \
        --rank 3 --world-size 8
"""

import argparse
import os
import sys
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description='Test exact dataset initialization')
    parser.add_argument('--wdsmeta-path', required=True, help='Path to .wdsmeta.json file')
    parser.add_argument('--rank', type=int, default=0, help='Simulate this rank')
    parser.add_argument('--world-size', type=int, default=8, help='Simulate this world size')
    parser.add_argument('--test-getitem', action='store_true', help='Also test __getitem__ for first few samples')
    parser.add_argument('--num-samples', type=int, default=3, help='Number of samples to test with __getitem__')
    
    args = parser.parse_args()
    
    print("=" * 80)
    print(f"EXACT DATASET INITIALIZATION TEST")
    print("=" * 80)
    print(f"Simulating: Rank {args.rank}/{args.world_size}")
    print(f"Metadata: {args.wdsmeta_path}")
    print("=" * 80)
    print()
    
    # Set environment variables BEFORE importing LMFlow
    os.environ['RANK'] = str(args.rank)
    os.environ['LOCAL_RANK'] = str(args.rank)
    os.environ['WORLD_SIZE'] = str(args.world_size)
    os.environ['NNODES'] = '1'
    os.environ['NODE_RANK'] = '0'
    
    print(f"[Test] Environment variables set:")
    print(f"  RANK={os.environ['RANK']}")
    print(f"  LOCAL_RANK={os.environ['LOCAL_RANK']}")
    print(f"  WORLD_SIZE={os.environ['WORLD_SIZE']}")
    print()
    
    # Import AFTER setting env vars (important!)
    sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
    from lmflow.args import DatasetArguments
    from lmflow.datasets.webdataset_multi_modal_dataset import CustomWebDatasetMultiModalDataset
    
    print("[Test] Creating DatasetArguments...")
    data_args = DatasetArguments(
        dataset_path=args.wdsmeta_path,
    )
    print(f"[Test] ✅ DatasetArguments created")
    print()
    
    print("[Test] ========================================")
    print("[Test] Creating CustomWebDatasetMultiModalDataset...")
    print("[Test] Watch for [WebDataset] debug output below:")
    print("[Test] ========================================")
    print()
    
    try:
        dataset = CustomWebDatasetMultiModalDataset(args.wdsmeta_path, data_args)
        
        print()
        print("[Test] ========================================")
        print("[Test] Dataset Creation Result:")
        print("[Test] ========================================")
        print(f"[Test] ✅ Dataset object created successfully")
        print(f"[Test] Dataset type: {type(dataset)}")
        print(f"[Test] Dataset length: {len(dataset)}")
        print()
        
        if len(dataset) == 0:
            print("[Test] ❌❌❌ CRITICAL: Dataset is EMPTY! ❌❌❌")
            print("[Test] This REPRODUCES the training failure!")
            print("[Test] Expected non-zero samples but got 0.")
            print()
            print("[Test] Check the [WebDataset] debug output above to see:")
            print("[Test]   - How many shards were assigned to this rank")
            print("[Test]   - How many samples were loaded from each shard")
            print("[Test]   - Whether any errors occurred during indexing")
            sys.exit(1)
        else:
            print(f"[Test] ✅ Dataset has {len(dataset)} samples (NON-EMPTY)")
            print()
            
            if args.test_getitem:
                print(f"[Test] Testing __getitem__ for first {args.num_samples} samples...")
                print()
                for i in range(min(args.num_samples, len(dataset))):
                    try:
                        print(f"[Test] Loading sample {i}...")
                        sample = dataset[i]
                        
                        if 'messages' not in sample:
                            print(f"[Test] ❌ Sample {i}: Missing 'messages' key")
                            print(f"[Test]   Available keys: {list(sample.keys())}")
                            continue
                        
                        messages = sample['messages']
                        print(f"[Test] ✅ Sample {i}: {len(messages)} messages")
                        
                        # Check message structure
                        for msg_idx, msg in enumerate(messages):
                            role = msg.get('role', '?')
                            content = msg.get('content', [])
                            if isinstance(content, list):
                                types = [c.get('type', '?') for c in content]
                                print(f"[Test]     Message {msg_idx} ({role}): content types = {types}")
                            else:
                                print(f"[Test]     Message {msg_idx} ({role}): content type = {type(content)}")
                        print()
                        
                    except Exception as e:
                        print(f"[Test] ❌ Error loading sample {i}: {e}")
                        import traceback
                        traceback.print_exc()
                        break
            
            print()
            print("[Test] ========================================")
            print("[Test] ✅✅✅ ALL TESTS PASSED ✅✅✅")
            print("[Test] ========================================")
            print(f"[Test] Dataset is NOT empty (has {len(dataset)} samples)")
            print("[Test] This means the issue is NOT in dataset initialization")
            print("[Test] but might be in how the training script uses the dataset.")
    
    except Exception as e:
        print()
        print("[Test] ========================================")
        print("[Test] ❌❌❌ DATASET CREATION FAILED ❌❌❌")
        print("[Test] ========================================")
        print(f"[Test] Error: {e}")
        print()
        import traceback
        traceback.print_exc()
        print()
        print("[Test] This error might explain the training failure.")
        sys.exit(1)


if __name__ == '__main__':
    main()

