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

# Agree to Coqui TOS
os.environ["COQUI_TOS_AGREED"] = "1"

# Initialize TTS model
device = "cuda" if torch.cuda.is_available() else "cpu"
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", device=device)

def clone(text, audio):
    output_path = "./output.wav"
    tts.tts_to_file(text=text, speaker_wav=audio, language="en", file_path=output_path)
    return output_path

# Define the UI using Blocks
with gr.Blocks(title="Advanced Voice Clone", theme=gr.themes.Soft(primary_hue="teal")) as demo:
    gr.Markdown(
        """
        # 🎤 Voice Clone

        **by [Tony Assi](https://www.tonyassi.com/)**

        This application uses the **xtts_v2** model for voice cloning.  
        *Non-commercial use only.*

        [Coqui Public Model License](https://coqui.ai/cpml) | 
        📧 [Contact Me](mailto:tony.assi.media@gmail.com)

        ---
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            text_input = gr.Textbox(
                label="Enter Text",
                placeholder="Type the text you want to clone...",
                lines=4
            )
            audio_input = gr.Audio(
                type="filepath",
                label="Upload Reference Voice",
                elem_id="audio_upload"
            )
            clone_button = gr.Button("Clone Voice", variant="primary")
        
        with gr.Column(scale=1):
            output_audio = gr.Audio(
                type="filepath",
                label="Cloned Voice Output",
                interactive=False
            )
            with gr.Accordion("Example Voices", open=False):
                gr.Markdown(
                    """
                    ### Preloaded Examples

                    - **Dorothy from Wizard of Oz**
                      - *Sample Audio:* [Download](./audio/Wizard-of-Oz-Dorthy.wav)
                    - **Vito Corleone from The Godfather**
                      - *Sample Audio:* [Download](./audio/Godfather.wav)
                    - **Paris Hilton**
                      - *Sample Audio:* [Download](./audio/Paris-Hilton.mp3)
                    - **Megan Fox from Transformers**
                      - *Sample Audio:* [Download](./audio/Megan-Fox.mp3)
                    - **Jeff Goldblum**
                      - *Sample Audio:* [Download](./audio/Jeff-Goldblum.mp3)
                    - **Heath Ledger as the Joker**
                      - *Sample Audio:* [Download](./audio/Heath-Ledger.mp3)
                    """
                )

    gr.Markdown(
        """
        ---
        ❤️ If you find this tool useful, please consider giving it a thumbs up!
        """
    )

    # Connect the button to the function
    clone_button.click(
        clone,
        inputs=[text_input, audio_input],
        outputs=output_audio,
        queue=True
    )

    # Add custom CSS for enhanced styling
    demo.add_css("""
    #audio_upload > label {
        background-color: #14b8a6;
        color: white;
        padding: 10px;
        border-radius: 5px;
        cursor: pointer;
    }
    #audio_upload > label:hover {
        background-color: #0d9488;
    }
    """)

# Launch the app
demo.launch(server_name="0.0.0.0", server_port=7860)