#!/usr/bin/env python3
"""
Test script for C-3PO TTS model integration
"""

import os
import requests
import json
import tempfile
from pathlib import Path

# Test configuration
API_BASE_URL = "http://localhost:7860"
TEST_TEXTS = [
    "I am C-3PO, human-cyborg relations.",
    "The odds of successfully navigating an asteroid field are approximately 3,720 to 1.",
    "R2-D2, you know better than to trust a strange computer!",
    "Oh my! How interesting!"
]

def test_health_check():
    """Test the health check endpoint"""
    print("🔍 Testing health check...")
    try:
        response = requests.get(f"{API_BASE_URL}/health")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Health check passed")
            print(f"   Model: {data.get('model', 'Unknown')}")
            print(f"   Device: {data.get('device', 'Unknown')}")
            print(f"   C-3PO voice available: {data.get('c3po_voice_available', False)}")
            print(f"   Model path: {data.get('model_path', 'Not specified')}")
            return True
        else:
            print(f"❌ Health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Health check error: {e}")
        return False

def test_c3po_endpoint():
    """Test the dedicated C-3PO endpoint"""
    print("\n🎭 Testing C-3PO endpoint...")
    
    test_text = "I am C-3PO, human-cyborg relations."
    
    try:
        data = {
            'text': test_text,
            'language': 'en'
        }
        
        response = requests.post(f"{API_BASE_URL}/tts-c3po", data=data)
        
        if response.status_code == 200:
            # Save the audio file
            output_path = Path(tempfile.gettempdir()) / "c3po_test_output.wav"
            with open(output_path, 'wb') as f:
                f.write(response.content)
            
            print(f"✅ C-3PO endpoint test passed")
            print(f"   Audio saved to: {output_path}")
            print(f"   File size: {os.path.getsize(output_path)} bytes")
            return True
        else:
            print(f"❌ C-3PO endpoint failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ C-3PO endpoint error: {e}")
        return False

def test_general_tts_with_c3po():
    """Test the general TTS endpoint with C-3PO voice enabled"""
    print("\n🎤 Testing general TTS with C-3PO voice...")
    
    test_text = "The odds of successfully navigating an asteroid field are approximately 3,720 to 1."
    
    try:
        data = {
            'text': test_text,
            'language': 'en',
            'use_c3po_voice': 'true'
        }
        
        response = requests.post(f"{API_BASE_URL}/tts", data=data)
        
        if response.status_code == 200:
            # Save the audio file
            output_path = Path(tempfile.gettempdir()) / "general_c3po_test_output.wav"
            with open(output_path, 'wb') as f:
                f.write(response.content)
            
            print(f"✅ General TTS with C-3PO test passed")
            print(f"   Audio saved to: {output_path}")
            print(f"   File size: {os.path.getsize(output_path)} bytes")
            return True
        else:
            print(f"❌ General TTS with C-3PO failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ General TTS with C-3PO error: {e}")
        return False

def test_json_endpoint():
    """Test the JSON endpoint"""
    print("\n📄 Testing JSON endpoint...")
    
    test_text = "R2-D2, you know better than to trust a strange computer!"
    
    try:
        data = {
            'text': test_text,
            'language': 'en'
        }
        
        headers = {'Content-Type': 'application/json'}
        response = requests.post(f"{API_BASE_URL}/tts-json", json=data, headers=headers)
        
        if response.status_code == 200:
            # Save the audio file
            output_path = Path(tempfile.gettempdir()) / "json_c3po_test_output.wav"
            with open(output_path, 'wb') as f:
                f.write(response.content)
            
            print(f"✅ JSON endpoint test passed")
            print(f"   Audio saved to: {output_path}")
            print(f"   File size: {os.path.getsize(output_path)} bytes")
            return True
        else:
            print(f"❌ JSON endpoint failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ JSON endpoint error: {e}")
        return False

def test_multilingual_support():
    """Test multilingual support with C-3PO voice"""
    print("\n🌍 Testing multilingual support...")
    
    test_cases = [
        ("Hello, I am C-3PO", "en"),
        ("Hola, soy C-3PO", "es"),
        ("Bonjour, je suis C-3PO", "fr"),
        ("Guten Tag, ich bin C-3PO", "de")
    ]
    
    success_count = 0
    
    for text, language in test_cases:
        try:
            data = {
                'text': text,
                'language': language
            }
            
            response = requests.post(f"{API_BASE_URL}/tts-c3po", data=data)
            
            if response.status_code == 200:
                output_path = Path(tempfile.gettempdir()) / f"c3po_test_{language}.wav"
                with open(output_path, 'wb') as f:
                    f.write(response.content)
                
                print(f"   ✅ {language}: {text} -> {output_path}")
                success_count += 1
            else:
                print(f"   ❌ {language}: Failed ({response.status_code})")
                
        except Exception as e:
            print(f"   ❌ {language}: Error - {e}")
    
    print(f"\n   Multilingual test: {success_count}/{len(test_cases)} languages successful")
    return success_count == len(test_cases)

def main():
    """Run all tests"""
    print("🚀 Starting C-3PO TTS Model Tests")
    print("=" * 50)
    
    tests = [
        test_health_check,
        test_c3po_endpoint,
        test_general_tts_with_c3po,
        test_json_endpoint,
        test_multilingual_support
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        if test():
            passed += 1
    
    print("\n" + "=" * 50)
    print(f"🎯 Test Results: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! C-3PO model integration is working correctly.")
    else:
        print("⚠️  Some tests failed. Check the API logs for more details.")
    
    print("\n💡 Tips:")
    print("   - Make sure the API server is running on http://localhost:7860")
    print("   - Check that the C-3PO model downloaded successfully")
    print("   - Generated audio files are saved in the system temp directory")

if __name__ == "__main__":
    main()