#!/usr/bin/env python3
"""
TTS Synchronization Tool
Compares word timestamps between espeak and ElevenLabs
"""

import subprocess
import asyncio
import json
import os
import sys
import wave

WHISPER_CLI = "/workspace/MuseTalk1.5/vendor/whisper/build/bin/whisper-cli"
WHISPER_MODEL = "/workspace/MuseTalk1.5/vendor/whisper/models/ggml-large-v3-turbo-q5_0.bin"

def get_audio_duration(filepath):
    """Get audio duration in seconds"""
    if filepath.endswith('.wav'):
        with wave.open(filepath, 'r') as f:
            return f.getnframes() / float(f.getframerate())
    else:
        result = subprocess.run([
            'ffprobe', '-v', 'error', '-show_entries', 'format=duration',
            '-of', 'default=noprint_wrappers=1:nokey=1', filepath
        ], capture_output=True, text=True)
        return float(result.stdout.strip())

def extract_word_timestamps(audio_path):
    """Extract word timestamps using local Whisper"""
    # Convert to wav if needed
    if not audio_path.endswith('.wav'):
        wav_path = audio_path.replace('.mp3', '.wav')
        subprocess.run([
            'ffmpeg', '-y', '-v', 'quiet', '-i', audio_path, '-ar', '16000', '-ac', '1', wav_path
        ], capture_output=True)
        audio_path = wav_path
    
    # Run whisper-cli
    result = subprocess.run([
        WHISPER_CLI, '-m', WHISPER_MODEL, '-f', audio_path,
        '--output-json', '--max-len', '1', '-l', 'pt', '--no-prints'
    ], capture_output=True, text=True)
    
    # Read JSON output
    json_path = audio_path + '.json'
    if os.path.exists(json_path):
        with open(json_path) as f:
            data = json.load(f)
        os.remove(json_path)
        
        words = []
        for item in data.get('transcription', []):
            text = item['text'].strip()
            if text and text not in [',', '.', '?', '!']:
                words.append({
                    'text': text,
                    'start': item['offsets']['from'],
                    'end': item['offsets']['to']
                })
        return words
    return []

def generate_espeak(text, output_path, speed=175):
    """Generate espeak audio"""
    subprocess.run([
        "espeak-ng", "-v", "pt-br", "-s", str(speed), "-w", output_path, text
    ], capture_output=True)
    return get_audio_duration(output_path)

async def generate_elevenlabs(text, output_path):
    """Generate ElevenLabs audio"""
    import httpx
    
    ELEVENLABS_API_KEY = "sk_857e9e6f2412ddf3ff5334b736e4b571641d26225c0d8d62"
    ELEVENLABS_VOICE_ID = "21m00Tcm4TlvDq8ikWAM"
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(
            f"https://api.elevenlabs.io/v1/text-to-speech/{ELEVENLABS_VOICE_ID}/stream",
            headers={"xi-api-key": ELEVENLABS_API_KEY, "Content-Type": "application/json"},
            json={"text": text, "model_id": "eleven_flash_v2_5", "voice_settings": {"stability": 0.5, "similarity_boost": 0.75}}
        )
        with open(output_path, "wb") as f:
            f.write(response.content)
    return get_audio_duration(output_path)

def generate_espeak_synced(text, output_path, target_duration, base_speed=175):
    """Generate espeak with adjusted speed to match target duration"""
    # Get base duration
    temp_path = "/tmp/espeak_temp.wav"
    subprocess.run([
        "espeak-ng", "-v", "pt-br", "-s", str(base_speed), "-w", temp_path, text
    ], capture_output=True)
    base_duration = get_audio_duration(temp_path)
    
    # Calculate required speed
    required_speed = int(base_speed * (base_duration / target_duration))
    required_speed = max(80, min(450, required_speed))
    
    # Generate with adjusted speed
    subprocess.run([
        "espeak-ng", "-v", "pt-br", "-s", str(required_speed), "-w", output_path, text
    ], capture_output=True)
    
    return get_audio_duration(output_path), required_speed

async def compare_timestamps(text):
    """Compare word timestamps between espeak and ElevenLabs"""
    print(f"\n{'='*70}")
    print(f"COMPARAÇÃO DE TIMESTAMPS: \"{text}\"")
    print('='*70)
    
    # Generate ElevenLabs
    eleven_path = "/tmp/compare_eleven.mp3"
    print("\n1. Gerando ElevenLabs...")
    eleven_dur = await generate_elevenlabs(text, eleven_path)
    print(f"   Duração: {eleven_dur:.2f}s")
    
    # Generate synced espeak
    espeak_path = "/tmp/compare_espeak.wav"
    print("\n2. Gerando espeak sincronizado...")
    espeak_dur, speed = generate_espeak_synced(text, espeak_path, eleven_dur)
    print(f"   Duração: {espeak_dur:.2f}s (speed={speed})")
    
    # Extract timestamps
    print("\n3. Extraindo timestamps com Whisper...")
    eleven_words = extract_word_timestamps(eleven_path)
    espeak_words = extract_word_timestamps(espeak_path)
    
    # Compare
    print("\n" + "-"*70)
    print(f"{'PALAVRA':<15} {'ELEVENLABS':<20} {'ESPEAK':<20} {'DIFF':<10}")
    print("-"*70)
    
    total_diff = 0
    count = 0
    
    for i, ew in enumerate(eleven_words):
        if i < len(espeak_words):
            sw = espeak_words[i]
            diff = abs(ew['start'] - sw['start'])
            total_diff += diff
            count += 1
            
            print(f"{ew['text']:<15} {ew['start']:>5}-{ew['end']:<5} ms    {sw['start']:>5}-{sw['end']:<5} ms    {diff:>5} ms")
    
    print("-"*70)
    
    if count > 0:
        avg_diff = total_diff / count
        print(f"\nMÉDIA DE DIFERENÇA: {avg_diff:.0f} ms")
        
        if avg_diff < 50:
            print("✓ EXCELENTE SINCRONIZAÇÃO!")
        elif avg_diff < 100:
            print("✓ BOA SINCRONIZAÇÃO")
        elif avg_diff < 200:
            print("⚠ SINCRONIZAÇÃO ACEITÁVEL")
        else:
            print("✗ SINCRONIZAÇÃO RUIM")
    
    return eleven_words, espeak_words

if __name__ == "__main__":
    text = sys.argv[1] if len(sys.argv) > 1 else "Olá, tudo bem? Como posso te ajudar?"
    asyncio.run(compare_timestamps(text))