#!/usr/bin/env python3
"""
Test TTS API with Urdu text and save audio outputs
"""
import requests
import os
from pathlib import Path
from datetime import datetime

# API configuration
API_URL = "https://ebitlogix-parler-tts-api.hf.space"
# For local testing, change to: API_URL = "http://localhost:7860"

OUTPUT_DIR = Path("tts_output")
OUTPUT_DIR.mkdir(exist_ok=True)

# Test cases with Urdu text
TEST_CASES = [
    {
        "text": "سلام دنیا",
        "speaker": "Divya",
        "description": "Hello World"
    },
    {
        "text": "مرحبا بك في تطبيق تحويل النص إلى كلام",
        "speaker": "Rani",
        "description": "Welcome to TTS App"
    },
    {
        "text": "یہ ایک اردو متن ہے جو آواز میں تبدیل ہوگا",
        "speaker": "Rohit",
        "description": "This is Urdu text converted to speech"
    },
    {
        "text": "کتاب پڑھنا میرا پسندیدہ کام ہے",
        "speaker": "Aman",
        "description": "Reading books is my favorite"
    },
    {
        "text": "خوشامدید، یہ ایک خوبصورت دن ہے",
        "speaker": "Generic Female",
        "description": "Welcome, it's a beautiful day"
    },
    {
        "text": "ہمیں اپنے لیے بہتر مستقبل بنانی ہے",
        "speaker": "Generic Male",
        "description": "We must build a better future"
    },
]

def test_health():
    """Test API health check."""
    print("\n" + "="*60)
    print("Testing API Health Check")
    print("="*60)

    try:
        response = requests.get(f"{API_URL}/")
        if response.status_code == 200:
            data = response.json()
            print(f"✅ API is running")
            print(f"   Model: {data.get('model')}")
            print(f"   Speakers: {', '.join(data.get('speakers', []))}")
            print(f"   Sample Rate: {data.get('sample_rate')} Hz")
            return True
        else:
            print(f"❌ API health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Cannot connect to API: {e}")
        return False

def test_tts_generation():
    """Generate TTS audio for all test cases."""
    print("\n" + "="*60)
    print("Testing TTS Generation with Urdu Text")
    print("="*60 + "\n")

    successful = 0
    failed = 0

    for i, test_case in enumerate(TEST_CASES, 1):
        text = test_case["text"]
        speaker = test_case["speaker"]
        description = test_case["description"]

        print(f"[{i}/{len(TEST_CASES)}] Testing: {description}")
        print(f"     Text: {text}")
        print(f"     Speaker: {speaker}")

        try:
            payload = {
                "text": text,
                "speaker": speaker,
                "pitch": "Moderate",
                "rate": "Moderate",
                "temperature": 0.8,
                "do_sample": True
            }

            response = requests.post(f"{API_URL}/tts", json=payload, timeout=60)

            if response.status_code != 200:
                print(f"     ❌ Failed: {response.status_code} - {response.text[:100]}")
                failed += 1
                continue

            # Save audio file
            audio_data = response.content
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{timestamp}_{speaker.replace(' ', '_')}_{i}.wav"
            filepath = OUTPUT_DIR / filename

            with open(filepath, "wb") as f:
                f.write(audio_data)

            file_size_kb = len(audio_data) / 1024
            print(f"     ✅ Success - Saved to: {filepath} ({file_size_kb:.1f} KB)")
            successful += 1

        except requests.exceptions.Timeout:
            print(f"     ❌ Failed: Request timeout (60s)")
            failed += 1
        except Exception as e:
            print(f"     ❌ Failed: {str(e)[:100]}")
            failed += 1

        print()

    # Summary
    print("="*60)
    print(f"Results: {successful} successful, {failed} failed")
    print(f"Audio files saved to: {OUTPUT_DIR.absolute()}")
    print("="*60)

    return successful, failed

def main():
    """Run all tests."""
    print("\n" + "█"*60)
    print("█  TTS API Test Suite - Urdu Text Generation")
    print("█"*60)

    # Health check
    if not test_health():
        print("\n❌ API is not available. Exiting...")
        return

    # TTS generation tests
    successful, failed = test_tts_generation()

    # Final status
    if failed == 0:
        print("\n✅ All tests passed!")
    else:
        print(f"\n⚠️  {failed} test(s) failed.")

    print(f"\n📁 Output folder: {OUTPUT_DIR.absolute()}")
    print(f"📊 Generated {successful} audio files\n")

if __name__ == "__main__":
    main()