import gradio as gr
import torch
from TTS.api import TTS
import os
import subprocess

import nltk
nltk.download('punkt')

def run_all_models(input_text, progress=gr.Progress()):
    # Initialize output file paths
    xtts_output = "xtts_output.wav"
    styletts2_output = "styletts2_output.wav"
    piper_output = "piper_output.wav"

    # 1. Running XttsV2
    progress(0, desc="Running XttsV2...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    os.environ["COQUI_TOS_AGREED"] = "1"
    os.system('curl -L -o default_voice.wav https://github.com/DrewThomasson/ebook2audiobookXTTS/raw/main/default_voice.wav')
    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
    tts.tts_to_file(text=input_text, speaker_wav="default_voice.wav", language="en", file_path=xtts_output)
    progress(1/3, desc="XttsV2 completed")

    # 2. Running StyleTTS2
    progress(1/3, desc="Running StyleTTS2...")
    from styletts2 import tts as style_tts
    my_tts = style_tts.StyleTTS2()
    my_tts.inference(input_text, output_wav_file=styletts2_output)
    progress(2/3, desc="StyleTTS2 completed")

    # 3. Running Piper-TTS
    progress(2/3, desc="Running Piper-TTS...")
    command = f'echo "{input_text}" | piper --model en_US-lessac-medium --output_file {piper_output}'
    subprocess.run(command, shell=True)
    progress(3/3, desc="Piper-TTS completed")

    # Return the file paths for downloading
    return xtts_output, styletts2_output, piper_output

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Test TTS Models with Progress Bar")
    
    # Input text box for user input
    input_text = gr.Textbox(label="Input Text", placeholder="Enter the text you want all models to use", value="Hello world!")

    # Download buttons for the outputs
    xtts_output = gr.File(label="Download XttsV2 Output")
    styletts2_output = gr.File(label="Download StyleTTS2 Output")
    piper_output = gr.File(label="Download Piper Output")

    # Button to run all models and show progress
    run_button = gr.Button("Run All Models")

    # Link button click to the run_all_models function
    run_button.click(fn=run_all_models, inputs=input_text, outputs=[xtts_output, styletts2_output, piper_output])

# Launch the app
demo.launch()