import os
import gradio as gr
from googletrans import Translator, LANGUAGES
from TTS.api import TTS

# Set environment variable to automatically accept the license terms
os.environ["COQUI_TOS_AGREED"] = "1"

# Load the TTS model once
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)
translator = Translator()

# Define the list of supported languages
SUPPORTED_LANGUAGES = {
    "English (en)": "en",
    "Spanish (es)": "es",
    "French (fr)": "fr",
    "German (de)": "de",
    "Italian (it)": "it",
    "Portuguese (pt)": "pt",
    "Polish (pl)": "pl",
    "Turkish (tr)": "tr",
    "Russian (ru)": "ru",
    "Dutch (nl)": "nl",
    "Czech (cs)": "cs",
    "Arabic (ar)": "ar",
    "Chinese (zh-cn)": "zh-cn",
    "Japanese (ja)": "ja",
    "Hungarian (hu)": "hu",
    "Korean (ko)": "ko",
    "Hindi (hi)": "hi"
}

def translate_text(text, input_lang, output_lang):
    try:
        translated = translator.translate(text, src=input_lang, dest=output_lang)
        return translated.text
    except Exception as e:
        return str(e)

def generate_speech(input_text, voice_path, input_lang, output_lang, translation_enabled):
    try:
        # Check if translation is enabled
        if translation_enabled:
            # Translate text
            translated_text = translate_text(input_text, input_lang, output_lang)
            if 'ERROR' in translated_text:
                return translated_text, None
        else:
            translated_text = input_text
            output_lang = input_lang

        # Generate speech
        output_path = "output.wav"
        tts.tts_to_file(
            text=translated_text,
            file_path=output_path,
            speaker_wav=voice_path,
            language=output_lang,
            split_sentences=True
        )
        return translated_text, output_path
    except Exception as e:
        return str(e), None

def update_output_lang_options(translation_enabled, input_lang):
    if not translation_enabled:
        return gr.update(choices=[input_lang], value=input_lang, interactive=False)
    return gr.update(choices=list(SUPPORTED_LANGUAGES.values()), value=input_lang, interactive=True)

with gr.Blocks() as iface:
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(label="Input Text")
            voice_path = gr.Audio(type="filepath", label="Upload Voice")
            input_lang = gr.Dropdown(choices=list(SUPPORTED_LANGUAGES.values()), label="Input Language")
            output_lang = gr.Dropdown(choices=list(SUPPORTED_LANGUAGES.values()), label="Output Language")
            translation_enabled = gr.Checkbox(label="Enable Translation", value=True)
            generate_button = gr.Button("Generate Speech")

        with gr.Column(scale=1):
            translated_text = gr.Textbox(label="Translated Text")
            generated_speech = gr.Audio(label="Generated Speech")

    translation_enabled.change(
        update_output_lang_options,
        inputs=[translation_enabled, input_lang],
        outputs=output_lang
    )

    generate_button.click(
        generate_speech,
        inputs=[input_text, voice_path, input_lang, output_lang, translation_enabled],
        outputs=[translated_text, generated_speech]
    )

iface.launch()