#!/usr/bin/env python3
import sys
sys.path.append('third_party/Matcha-TTS')

print("=== Teste de TTS em Português Melhorado ===\n")

try:
    from cosyvoice.cli.cosyvoice import CosyVoice
    from cosyvoice.utils.file_utils import load_wav
    import torchaudio
    import os
    
    # Teste 1: Usar o modelo SFT que tem vozes pré-treinadas
    print("1. Testando com modelo SFT (vozes pré-treinadas)")
    print("-" * 50)
    
    sft_path = 'pretrained_models/CosyVoice-300M-SFT'
    
    # Verificar se o modelo SFT está completo
    if os.path.exists(f'{sft_path}/llm.pt'):
        print("Modelo SFT encontrado!")
    else:
        # Copiar arquivos necessários
        print("Preparando modelo SFT...")
        os.system(f'cp pretrained_models/CosyVoice-300M-direct/*.pt {sft_path}/')
        os.system(f'cp pretrained_models/CosyVoice-300M-direct/*.onnx {sft_path}/')
    
    cosyvoice_sft = CosyVoice(sft_path, load_jit=False, load_trt=False, fp16=False)
    
    # Listar vozes disponíveis
    print("\nVozes disponíveis no modelo SFT:")
    speakers = cosyvoice_sft.list_available_spks()
    for i, spk in enumerate(speakers[:10]):  # Mostrar primeiras 10
        print(f"  {i}: {spk}")
    
    # Testar com diferentes vozes
    test_text = "Olá! Este é um teste de síntese de voz em português brasileiro. Espero que a qualidade esteja melhor agora."
    
    # Tentar vozes que podem ser mais adequadas para português
    voice_tests = []
    
    # Procurar vozes femininas/masculinas em chinês que podem soar melhor
    for spk in speakers:
        if '中文女' in spk or '中文男' in spk or 'Chinese' in spk.lower():
            voice_tests.append(spk)
            if len(voice_tests) >= 3:
                break
    
    for i, voice in enumerate(voice_tests):
        print(f"\nTestando com voz: {voice}")
        for j, audio in enumerate(cosyvoice_sft.inference_sft(test_text, voice, stream=False)):
            output_file = f'output_pt_sft_{i}.wav'
            torchaudio.save(output_file, audio['tts_speech'], cosyvoice_sft.sample_rate)
            duration = audio['tts_speech'].shape[1] / cosyvoice_sft.sample_rate
            print(f"✓ Gerado: {output_file} ({duration:.2f}s)")
    
    # Teste 2: Usar áudio de referência diferente
    print("\n\n2. Testando com áudio de referência cross-lingual")
    print("-" * 50)
    
    # Usar o modelo base com cross-lingual prompt
    cosyvoice_base = CosyVoice('pretrained_models/CosyVoice-300M-direct', load_jit=False, load_trt=False, fp16=False)
    
    # Usar o prompt cross-lingual que pode ter melhor qualidade
    if os.path.exists('asset/cross_lingual_prompt.wav'):
        prompt_cross = load_wav('asset/cross_lingual_prompt.wav', 16000)
        
        # Teste com prompt em inglês para gerar português
        for i, audio in enumerate(cosyvoice_base.inference_cross_lingual(
            '<|pt|>' + test_text,  # Indicar idioma português
            prompt_cross,
            stream=False
        )):
            output_file = 'output_pt_crosslingual.wav'
            torchaudio.save(output_file, audio['tts_speech'], cosyvoice_base.sample_rate)
            duration = audio['tts_speech'].shape[1] / cosyvoice_base.sample_rate
            print(f"✓ Gerado: {output_file} ({duration:.2f}s)")
    
    # Teste 3: Ajustar o texto de prompt
    print("\n\n3. Testando com prompts otimizados")
    print("-" * 50)
    
    # Usar prompts mais curtos e simples
    prompt_speech = load_wav('asset/zero_shot_prompt.wav', 16000)
    
    prompts_test = [
        "Olá.",  # Prompt muito curto
        "Teste de voz.",  # Prompt médio
        "Esta é uma voz brasileira."  # Prompt com contexto
    ]
    
    for i, prompt_text in enumerate(prompts_test):
        print(f"\nPrompt: '{prompt_text}'")
        for j, audio in enumerate(cosyvoice_base.inference_zero_shot(
            test_text,
            prompt_text,
            prompt_speech,
            stream=False
        )):
            output_file = f'output_pt_prompt_{i}.wav'
            torchaudio.save(output_file, audio['tts_speech'], cosyvoice_base.sample_rate)
            duration = audio['tts_speech'].shape[1] / cosyvoice_base.sample_rate
            print(f"✓ Gerado: {output_file} ({duration:.2f}s)")
    
    print("\n✅ Testes concluídos!")
    print("\nArquivos gerados:")
    for f in sorted(os.listdir('.')):
        if f.startswith('output_pt_') and f.endswith('.wav'):
            size = os.path.getsize(f) / 1024
            print(f"  - {f} ({size:.1f} KB)")
    
    print("\n💡 Dicas para melhorar a qualidade:")
    print("1. Use um áudio de referência de um falante nativo de português")
    print("2. O modelo CosyVoice2-0.5B pode ter melhor suporte multilíngue")
    print("3. Fine-tuning com dados em português melhoraria significativamente")
    
except Exception as e:
    print(f"\n❌ Erro: {e}")
    import traceback
    traceback.print_exc()