#!/usr/bin/env python3
"""
Clean Chunking Test - NO EMOTION TAGS

This test isolates the chunking bug from the emotion tag issue by using ONLY plain text.

Goal: Verify all chunks are processed and no text is truncated.
"""

import os
import requests
import subprocess
from openai import OpenAI
import re

# Configuration
API_KEY = "vn3_cdd6d45f2045d03d5adac56eda6af9a9b781211038972807f35d52dfb6400144"
API_URL = "http://localhost:8000/v1/tts/generate"
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')

client = OpenAI(api_key=OPENAI_API_KEY)

# Test cases - NO emotion tags, just plain text
TEST_CASES = {
    "short": {
        "text": "This is a short test with no chunking needed.",
        "expected_duration": "<5s",
        "output_file": "test_chunking_short.wav"
    },
    
    "medium": {
        "text": (
            "Once upon a time in a distant kingdom, there lived a young princess named Aurora. "
            "She was known throughout the land for her kindness and bravery. "
            "One day, a mysterious stranger arrived at the castle gates. "
            "He brought news of a terrible dragon that was terrorizing nearby villages. "
            "Without hesitation, Aurora decided to face the dragon herself. "
            "She gathered her sword and shield, preparing for the journey ahead. "
            "The kingdom watched in awe as their princess rode off into the sunset. "
            "Little did they know, this would be the beginning of a legendary adventure that would change the fate of the entire kingdom forever."
        ),
        "expected_duration": "10-15s",
        "output_file": "test_chunking_medium.wav"
    },
    
    "long": {
        "text": (
            "In the ancient kingdom of Solara, there lived a wise queen named Aria. "
            "She ruled with justice and compassion, earning the love of all her subjects. "
            "One fateful morning, a dark shadow fell across the land. "
            "The royal astronomer rushed to the throne room with urgent news. "
            "A comet of unprecedented size was heading directly toward their world. "
            "The queen immediately convened her council of advisors. "
            "Together, they devised a plan to protect the kingdom. "
            "The royal mages would create a magical shield, while the engineers built reinforced shelters. "
            "Days turned into weeks as the preparations continued. "
            "The entire kingdom worked together, united by a common purpose. "
            "Finally, on the night of the comet's arrival, everyone gathered in the great plaza. "
            "The queen stood before her people, radiating confidence and strength. "
            "Tonight, we face our greatest challenge, she proclaimed. "
            "But together, we are invincible! "
            "As the comet blazed across the sky, the magical shield shimmered to life. "
            "The ground shook, and brilliant light filled the heavens. "
            "But when dawn broke, the kingdom still stood. "
            "The people erupted in celebration, their faith vindicated. "
            "Queen Aria had saved her kingdom once again. "
            "And from that day forward, the people remembered the night they defied the stars. "
            "The legend of Queen Aria's Shield would be told for generations to come, "
            "inspiring hope in the hearts of all who heard it. "
            "Children would grow up hearing tales of the brave queen who never gave up. "
            "Scholars would study her wisdom and leadership for centuries. "
            "And whenever darkness threatened, the people would remember her words: "
            "Together, we are invincible! "
            "CRITICAL TEST: This sentence is at the END and must NOT be cut off. "
            "If you hear this sentence in the audio, the chunking system is working correctly. "
            "If this sentence is missing, we have confirmed the last chunk truncation bug. "
            "End of test text. Thank you for listening!"
        ),
        "expected_duration": "30-40s",
        "output_file": "test_chunking_long.wav"
    },
}


def generate_audio(text, output_file):
    """Generate audio via API."""
    print(f"🎙️  Generating audio...")
    print(f"   Text length: {len(text)} chars")
    print(f"   Output: {output_file}")
    
    response = requests.post(
        API_URL,
        headers={"Authorization": f"Bearer {API_KEY}"},
        json={
            "text": text,
            "speaker": "lipakshi",
            "temperature": 0.4,
            "top_p": 0.9,
            "max_tokens": 4096,
            "seed": 42,
        },
        timeout=120
    )
    
    if response.status_code != 200:
        print(f"❌ API error: {response.status_code}")
        print(response.text)
        return None
    
    # Save audio
    with open(output_file, 'wb') as f:
        f.write(response.content)
    
    # Get audio info
    audio_bytes = len(response.content)
    duration = (audio_bytes - 44) / (24000 * 2)  # 16-bit PCM, 24kHz
    
    print(f"✅ Audio generated:")
    print(f"   Size: {audio_bytes:,} bytes")
    print(f"   Duration: {duration:.2f}s")
    
    # Check response headers
    if 'X-Text-Chunked' in response.headers:
        print(f"   Chunked: {response.headers['X-Text-Chunked']}")
    if 'X-RTF' in response.headers:
        print(f"   RTF: {response.headers['X-RTF']}")
    
    print()
    return output_file


def transcribe_audio(audio_file):
    """Transcribe audio with OpenAI Whisper."""
    print(f"🎧 Transcribing with Whisper...")
    
    with open(audio_file, 'rb') as f:
        transcription = client.audio.transcriptions.create(
            model="whisper-1",
            file=f,
            response_format="verbose_json",
            timestamp_granularity=["word"]
        )
    
    transcript_text = transcription.text
    words = transcription.words if hasattr(transcription, 'words') else []
    
    print(f"✅ Transcription complete:")
    print(f"   Words transcribed: {len(words)}")
    print(f"   Duration: {transcription.duration:.2f}s")
    print()
    
    return transcript_text, words


def compare_texts(reference, transcribed):
    """Compare reference and transcribed text."""
    # Clean texts
    ref_words = re.findall(r'\b\w+\b', reference.lower())
    trans_words = re.findall(r'\b\w+\b', transcribed.lower())
    
    print(f"📊 Comparison:")
    print(f"   Reference words: {len(ref_words)}")
    print(f"   Transcribed words: {len(trans_words)}")
    print(f"   Accuracy: {len(trans_words)/len(ref_words)*100:.1f}%")
    print()
    
    # Find missing words
    trans_set = set(trans_words)
    missing = [w for w in ref_words if w not in trans_set]
    
    if missing:
        print(f"⚠️  Missing words: {len(missing)}")
        print(f"   First 20: {missing[:20]}")
        print()
        
        # Check if ending is missing
        last_30_ref = ref_words[-30:]
        last_30_in_trans = sum(1 for w in last_30_ref if w in trans_set)
        
        print(f"🎯 CRITICAL: Last 30 words check:")
        print(f"   Last 30 reference words: {' '.join(last_30_ref)}")
        print(f"   Found in transcript: {last_30_in_trans}/30")
        
        if last_30_in_trans < 20:
            print(f"   ❌ ENDING IS MISSING! Last chunk truncation confirmed!")
        else:
            print(f"   ✅ Ending is present. Missing words are scattered.")
        print()
    else:
        print(f"✅ All words present!")
        print()
    
    return len(missing), missing


def test_case(test_name, test_data):
    """Run a single test case."""
    print("=" * 80)
    print(f"🧪 TEST: {test_name}")
    print("=" * 80)
    print()
    
    text = test_data['text']
    output_file = test_data['output_file']
    
    # Step 1: Generate audio
    audio_file = generate_audio(text, output_file)
    if not audio_file:
        print(f"❌ Test FAILED: Audio generation failed")
        return False
    
    # Step 2: Transcribe
    transcribed, words = transcribe_audio(audio_file)
    
    # Step 3: Compare
    missing_count, missing_words = compare_texts(text, transcribed)
    
    # Step 4: Result
    if missing_count == 0:
        print(f"✅ TEST PASSED: All text present")
        result = True
    elif missing_count < 10:
        print(f"⚠️  TEST WARNING: {missing_count} words missing (likely ASR errors)")
        result = True
    else:
        print(f"❌ TEST FAILED: {missing_count} words missing (likely chunking bug)")
        result = False
    
    print()
    return result


def main():
    """Run all test cases."""
    print("\n" + "=" * 80)
    print("🚀 CLEAN CHUNKING TEST SUITE (NO EMOTION TAGS)")
    print("=" * 80)
    print()
    print("This test isolates the chunking bug by using ONLY plain text.")
    print("No emotion tags = No extra words in ASR = Clean comparison")
    print()
    
    results = {}
    for test_name, test_data in TEST_CASES.items():
        results[test_name] = test_case(test_name, test_data)
    
    # Summary
    print("=" * 80)
    print("📊 TEST SUMMARY")
    print("=" * 80)
    print()
    
    for test_name, passed in results.items():
        status = "✅ PASS" if passed else "❌ FAIL"
        print(f"{status} - {test_name}")
    
    total = len(results)
    passed = sum(1 for p in results.values() if p)
    
    print()
    print(f"Total: {total}, Passed: {passed}, Failed: {total - passed}")
    print()
    
    if all(results.values()):
        print("✅ ALL TESTS PASSED - Chunking system works correctly!")
    else:
        print("❌ SOME TESTS FAILED - Chunking bug confirmed!")
    
    print()


if __name__ == "__main__":
    main()