#!/usr/bin/env python3
"""
Script para gerar audio com StyleTTS2 e clonagem de voz.
Uso: python generate_audio.py --text "Seu texto aqui" --voice voice_ref.wav --output audio.wav
"""

import argparse
import os

# Fix para PyTorch 2.6+ com weights_only
import torch
original_load = torch.load
def patched_load(*args, **kwargs):
    kwargs['weights_only'] = False
    return original_load(*args, **kwargs)
torch.load = patched_load

def generate_audio(text: str, output_path: str, voice_ref: str = None, diffusion_steps: int = 10):
    """Gera audio usando StyleTTS2."""

    print(f"Carregando StyleTTS2...")
    from styletts2 import tts

    my_tts = tts.StyleTTS2()

    print(f"Gerando audio...")
    print(f"  Texto: {text[:50]}...")
    print(f"  Voz de referencia: {voice_ref or 'padrao'}")

    if voice_ref and os.path.exists(voice_ref):
        wav = my_tts.inference(
            text,
            target_voice_path=voice_ref,
            diffusion_steps=diffusion_steps
        )
    else:
        wav = my_tts.inference(
            text,
            diffusion_steps=diffusion_steps
        )

    # Salvar audio
    import scipy.io.wavfile as wavfile
    wavfile.write(output_path, 24000, wav)

    print(f"Audio salvo em: {output_path}")
    return output_path


def main():
    parser = argparse.ArgumentParser(description='Gerar audio com StyleTTS2')
    parser.add_argument('--text', '-t', required=True, help='Texto para converter em audio')
    parser.add_argument('--output', '-o', default='output.wav', help='Arquivo de saida')
    parser.add_argument('--voice', '-v', help='Arquivo WAV de referencia para clonagem de voz')
    parser.add_argument('--steps', '-s', type=int, default=10, help='Passos de difusao (mais = melhor qualidade)')

    args = parser.parse_args()

    generate_audio(
        text=args.text,
        output_path=args.output,
        voice_ref=args.voice,
        diffusion_steps=args.steps
    )


if __name__ == '__main__':
    main()