#!/usr/bin/env python3
"""
Simple test script for the Coqui TTS API
Run this to test the API locally
"""

import requests
import time
import os

# Configuration
API_BASE_URL = "http://localhost:7860"
TEST_TEXTS = [
    "Hello world, this is a test of the Coqui TTS API using the VITS model.",
    "The quick brown fox jumps over the lazy dog.",
    "Welcome to our minimal text-to-speech service!"
]

def test_health_check():
    """Test the health check endpoints"""
    print("🔍 Testing health check endpoints...")
    
    try:
        # Test root endpoint
        response = requests.get(f"{API_BASE_URL}/")
        print(f"GET / - Status: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"Model: {data.get('model')}")
            print(f"Engine: {data.get('engine')}")
        
        # Test health endpoint
        response = requests.get(f"{API_BASE_URL}/health")
        print(f"GET /health - Status: {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"Model loaded: {data.get('model_loaded')}")
        
    except requests.exceptions.ConnectionError:
        print("❌ Could not connect to the API. Make sure it's running on localhost:7860")
        return False
    
    return True

def test_get_endpoint():
    """Test the GET TTS endpoint"""
    print("\n🎤 Testing GET TTS endpoint...")
    
    for i, text in enumerate(TEST_TEXTS):
        try:
            params = {"text": text}
            
            print(f"Testing text {i+1}: '{text[:50]}...'")
            response = requests.get(f"{API_BASE_URL}/tts", params=params)
            
            if response.status_code == 200:
                # Save the audio file
                filename = f"test_output_get_{i+1}.wav"
                with open(filename, "wb") as f:
                    f.write(response.content)
                print(f"✅ Audio saved as {filename} ({len(response.content)} bytes)")
            else:
                print(f"❌ Error: {response.status_code} - {response.text}")
                
        except Exception as e:
            print(f"❌ Exception: {str(e)}")

def test_post_endpoint():
    """Test the POST TTS endpoint"""
    print("\n🎵 Testing POST TTS endpoint...")
    
    for i, text in enumerate(TEST_TEXTS):
        try:
            data = {"text": text}
            
            print(f"Testing text {i+1}: '{text[:50]}...'")
            response = requests.post(f"{API_BASE_URL}/tts", json=data)
            
            if response.status_code == 200:
                # Save the audio file
                filename = f"test_output_post_{i+1}.wav"
                with open(filename, "wb") as f:
                    f.write(response.content)
                print(f"✅ Audio saved as {filename} ({len(response.content)} bytes)")
            else:
                print(f"❌ Error: {response.status_code} - {response.text}")
                
        except Exception as e:
            print(f"❌ Exception: {str(e)}")

def test_form_endpoint():
    """Test the POST TTS endpoint with form data"""
    print("\n📋 Testing POST TTS endpoint with form data...")
    
    try:
        data = {"text": "This is a test using form data submission with Coqui TTS VITS model."}
        
        response = requests.post(f"{API_BASE_URL}/tts", data=data)
        
        if response.status_code == 200:
            filename = "test_output_form_vits.wav"
            with open(filename, "wb") as f:
                f.write(response.content)
            print(f"✅ Audio saved as {filename} ({len(response.content)} bytes)")
        else:
            print(f"❌ Error: {response.status_code} - {response.text}")
            
    except Exception as e:
        print(f"❌ Exception: {str(e)}")

def test_long_text():
    """Test with longer text"""
    print("\n📖 Testing with longer text...")
    
    long_text = """
    This is a longer text passage to test the VITS model's ability to handle 
    extended speech synthesis. The model should maintain good quality and 
    natural prosody throughout this longer utterance.
    """
    
    try:
        data = {"text": long_text.strip()}
        
        print("Testing with longer text passage...")
        response = requests.post(f"{API_BASE_URL}/tts", json=data)
        
        if response.status_code == 200:
            filename = "test_output_long_text.wav"
            with open(filename, "wb") as f:
                f.write(response.content)
            print(f"✅ Long text audio saved as {filename} ({len(response.content)} bytes)")
        else:
            print(f"❌ Error: {response.status_code} - {response.text}")
            
    except Exception as e:
        print(f"❌ Exception: {str(e)}")

def cleanup_test_files():
    """Clean up generated test files"""
    print("\n🧹 Cleaning up test files...")
    
    test_files = [f for f in os.listdir(".") if f.startswith("test_output_") and f.endswith(".wav")]
    
    for file in test_files:
        try:
            os.remove(file)
            print(f"Removed {file}")
        except Exception as e:
            print(f"Could not remove {file}: {str(e)}")

if __name__ == "__main__":
    print("🚀 Starting Coqui TTS VITS API Test Suite")
    print("=" * 50)
    
    # Test health check first
    if not test_health_check():
        print("\n❌ Health check failed. Exiting.")
        exit(1)
    
    # Wait for model to load
    print("\n⏳ Waiting for VITS model to load...")
    time.sleep(5)
    
    # Run tests
    test_get_endpoint()
    test_post_endpoint()
    test_form_endpoint()
    test_long_text()
    
    print("\n" + "=" * 50)
    print("✅ Test suite completed!")
    print("\nGenerated files demonstrate:")
    print("- VITS model voice quality")
    print("- Various input methods (GET, POST JSON, POST form)")
    print("- Short and long text handling")
    
    print("\nTo clean up test files, run:")
    print("python test_api.py --cleanup")
    
    # Check if cleanup flag is provided
    import sys
    if "--cleanup" in sys.argv:
        cleanup_test_files()