"""
Quick test script to verify data loading works correctly
Run this before starting the full training pipeline
"""

import sys
from pathlib import Path
from transformers import WhisperProcessor

# Add current directory to path
sys.path.insert(0, str(Path(__file__).parent))

from train import DataPreparer
from config import DataConfig
from dataset import IndicMultilingualDataset, AudioAugmentor
from config import AugmentationConfig


def test_data_loading():
    """Test that data loading works correctly"""
    print("="*80)
    print("TESTING DATA LOADING")
    print("="*80)
    
    # Initialize config
    data_config = DataConfig()
    aug_config = AugmentationConfig()
    
    # Test data preparer
    print("\n1. Testing DataPreparer...")
    preparer = DataPreparer(data_config, base_path=".")
    
    try:
        train_samples = preparer.prepare_training_data()
        print(f"✅ Successfully loaded {len(train_samples)} training samples")
        
        # Show sample structure
        if train_samples:
            print("\nSample structure:")
            sample = train_samples[0]
            for key, value in sample.items():
                if isinstance(value, str) and len(value) > 50:
                    print(f"  {key}: {value[:50]}...")
                else:
                    print(f"  {key}: {value}")
    except Exception as e:
        print(f"❌ Error loading training data: {e}")
        return False
    
    # Test dataset class
    print("\n2. Testing IndicMultilingualDataset...")
    try:
        processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
        print("✅ Loaded Whisper processor")
        
        augmentor = AudioAugmentor(aug_config)
        print("✅ Created audio augmentor")
        
        # Create dataset with data list
        dataset = IndicMultilingualDataset(
            processor=processor,
            data=train_samples[:5],  # Just test with 5 samples
            config=data_config,
            augmentor=None,  # Skip augmentation for test
            is_training=False,
        )
        print(f"✅ Created dataset with {len(dataset)} samples")
        
        # Try to load one sample
        print("\n3. Testing sample loading...")
        sample = dataset[0]
        print(f"✅ Successfully loaded sample")
        print(f"  Input features shape: {sample['input_features'].shape}")
        print(f"  Labels length: {len(sample['labels'])}")
        print(f"  Text: {sample['text'][:50]}...")
        print(f"  Language: {sample['language']}")
        print(f"  Code-switched: {sample['is_code_switched']}")
        
    except Exception as e:
        print(f"❌ Error creating dataset: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    print("\n" + "="*80)
    print("✅ ALL TESTS PASSED!")
    print("="*80)
    print("\nYou can now run the full training with:")
    print("  python train.py")
    print("\nor:")
    print("  ./run_training.sh")
    
    return True


if __name__ == "__main__":
    success = test_data_loading()
    sys.exit(0 if success else 1)