import gradio as gr
import asyncio
import time
import os
from tts import synthesize_and_play_audio
from pydub import AudioSegment
import shutil

async def generate_tts(input_text, reference_audio_path, output_path="cloned.wav"):
    print(reference_audio_path)
    await synthesize_and_play_audio(
        input_text=input_text,
        reference_audio_path=reference_audio_path,
        model="voxtral-mini-tts-260213",
        api_key=os.getenv("MISTRAL_API_KEY"),
        output_path=output_path,
        no_play=True
    )
    return output_path

def gradio_tts(input_text, audio_choice, uploaded_audio=None):
    # If an audio file is uploaded, save it to a fixed path
    if uploaded_audio is not None:
        reference_audio = "uploadedreference.wav"
        shutil.copy(uploaded_audio, reference_audio)
    else:
        reference_audio = audio_choice

    output_path = "cloned.wav"
    try:
        generated_audio = asyncio.run(generate_tts(input_text, reference_audio, output_path))
        return generated_audio
    except Exception as e:
        print(f"Error: {e}")
        return None



with open("styles.css", "r") as f:
    css = f.read()

examples = [
    ["Frontier AI in your hands.", "sample.mp3"],
    ["Hello, world!", "voice.wav"],
    ["This is a test.", "4languages.mp3"],
]

with gr.Blocks() as demo:
    gr.Markdown("## Voxtral TTS Demo", elem_classes="markdown")
    gr.Markdown("Voxtral TTS is a text-to-speech model that can clone voice from any reference audio. Learn more about [Voxtral TTS](https://huggingface.co/mistralai/Voxtral-3B-TTS-2603).", elem_classes="markdown")
    with gr.Tabs():
        with gr.TabItem("Predefined Voices"):
            gr.Markdown("# Predefined Voices TTS", elem_classes="markdown")
            gr.Markdown("Enter text to synthesize and select a predefined voice.", elem_classes="markdown")
            with gr.Row():
                with gr.Column(elem_classes="gradio-box"):
                    input_text_predefined = gr.Textbox(
                        label="Enter text to synthesize",
                        placeholder="Frontier AI in your hands.",
                        elem_classes="gradio-textbox"
                    )
                    audio_choice = gr.Dropdown(
                        label="Select a predefined voice",
                        choices=["sample.mp3", "voice.wav", "4languages.mp3"],
                        value="4languages.mp3",
                    )
                    submit_btn_predefined = gr.Button("Generate Audio", elem_classes="gradio-button")
                with gr.Column(elem_classes="gradio-box"):
                    output_audio_predefined = gr.Audio(label="Generated audio", elem_classes="gradio-audio")

            submit_btn_predefined.click(
                fn=lambda text, choice: gradio_tts(text, choice, None),
                inputs=[input_text_predefined, audio_choice],
                outputs=[output_audio_predefined],
            )
            gr.Examples(
                examples=examples,
                inputs=[input_text_predefined, audio_choice],
                outputs=[output_audio_predefined],
                fn=lambda text, choice: gradio_tts(text, choice, None),
                cache_examples=True,
            )
        with gr.TabItem("Voice Cloning"):
            gr.Markdown("# Voice Cloning TTS", elem_classes="markdown")
            gr.Markdown("Enter text to synthesize and upload your reference audio.", elem_classes="markdown")
            with gr.Row():
                with gr.Column(elem_classes="gradio-box"):
                    input_text_cloning = gr.Textbox(
                        label="Enter text to synthesize",
                        placeholder="Frontier AI in your hands.",
                        elem_classes="gradio-textbox"
                    )
                    uploaded_audio = gr.Audio(
                        label="Upload your reference audio",
                        type="filepath",
                        sources=["upload"],
                        elem_classes="gradio-audio"
                    )
                    submit_btn_cloning = gr.Button("Generate Audio", elem_classes="gradio-button")
                with gr.Column(elem_classes="gradio-box"):
                    output_audio_cloning = gr.Audio(label="Generated audio", elem_classes="gradio-audio")

            submit_btn_cloning.click(
                fn=lambda text, audio: gradio_tts(text, None, audio),
                inputs=[input_text_cloning, uploaded_audio],
                outputs=[output_audio_cloning],
            )
            gr.Examples(
                examples=[
                    ["Frontier AI in your hands.", "sample.mp3"],
                    ["Hello, world!", "voice.wav"],
                ],
                inputs=[input_text_cloning, uploaded_audio],
                outputs=[output_audio_cloning],
                fn=lambda text, audio: gradio_tts(text, None, audio),
                cache_examples=True,
            )

if __name__ == "__main__":
    demo.launch(share=False, css=css)