#!/usr/bin/env python3
"""
Test script for KittenTTS functionality
"""

import soundfile as sf
import numpy as np
from kittentts import KittenTTS

def test_kittentts():
    """Test the KittenTTS model with a simple example"""
    
    print("🚀 Testing KittenTTS...")
    
    try:
        # Initialize the model
        print("📦 Loading KittenTTS model...")
        model = KittenTTS("KittenML/kitten-tts-nano-0.1")
        print("✅ Model loaded successfully!")
        
        # Test text
        test_text = "Hello! This is a test of the KittenTTS model."
        test_voice = 'expr-voice-2-f'
        
        print(f"🎤 Generating speech for: '{test_text}'")
        print(f"🎵 Using voice: {test_voice}")
        
        # Generate audio
        audio = model.generate(test_text, voice=test_voice)
        
        print(f"✅ Audio generated successfully!")
        print(f"📊 Audio shape: {audio.shape}")
        print(f"📊 Audio dtype: {audio.dtype}")
        print(f"📊 Audio range: {np.min(audio):.4f} to {np.max(audio):.4f}")
        
        # Save test audio
        output_file = 'test_output.wav'
        sf.write(output_file, audio, 24000)
        print(f"💾 Audio saved to: {output_file}")
        
        # Test all voices
        print("\n🎭 Testing all available voices...")
        available_voices = [
            'expr-voice-2-m', 'expr-voice-2-f', 'expr-voice-3-m', 'expr-voice-3-f',
            'expr-voice-4-m', 'expr-voice-4-f', 'expr-voice-5-m', 'expr-voice-5-f'
        ]
        
        for i, voice in enumerate(available_voices, 1):
            print(f"  {i}/8: Testing {voice}...")
            try:
                test_audio = model.generate("Test voice.", voice=voice)
                print(f"    ✅ {voice} works!")
            except Exception as e:
                print(f"    ❌ {voice} failed: {e}")
        
        print("\n🎉 All tests completed successfully!")
        return True
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        return False

if __name__ == "__main__":
    success = test_kittentts()
    exit(0 if success else 1)