import gradio as gr
from transformers import pipeline
from gtts import gTTS
import os

# Define translation models for each language pair
translation_models = {
    'Chinese': "Helsinki-NLP/opus-mt-en-zh",
    'French': "Helsinki-NLP/opus-mt-en-fr",
    'Spanish': "Helsinki-NLP/opus-mt-en-es",
    'Russian': "Helsinki-NLP/opus-mt-en-ru",
    'Arabic': "Helsinki-NLP/opus-mt-en-ar",
    'Urdu': "Helsinki-NLP/opus-mt-en-ur",
    'German': "Helsinki-NLP/opus-mt-en-de"
}

# Initialize translators
translators = {lang: pipeline("translation", model=model) for lang, model in translation_models.items() if model}

# Supported Languages
language_options = {
    'Arabic': 'ar',
    'Chinese': 'zh',
    'French': 'fr',
    'German': 'de',
    'Russian': 'ru',
    'Spanish': 'es',
    'Urdu': 'ur',
  }

def translate_text(input_text, language):
    if language == 'English':
        return input_text
    model_name = translation_models.get(language)
    if model_name:
        try:
            translator = translators[language]
            translation = translator(input_text)
            return translation[0]['translation_text']
        except Exception as e:
            return f"Error during translation: {str(e)}"
    else:
        return "Translation model not available."

def text_to_speech(output_text, language):
    tts = gTTS(output_text, lang=language_options[language])
    tts.save("output.mp3")
    return "output.mp3"

def input_to_speech(input_text):
    tts = gTTS(input_text, lang='en')
    tts.save("input_speech.mp3")
    return "input_speech.mp3"

def generate(input_text, language, translate, text_to_speech_opt, input_to_speech_opt):
    translated_text = ""
    translated_audio = None
    input_audio = None
    
    if translate:
        translated_text = translate_text(input_text, language)
    
    if text_to_speech_opt and translated_text:
        translated_audio = text_to_speech(translated_text, language)
    
    if input_to_speech_opt:
        input_audio = input_to_speech(input_text)
    
    return translated_text, translated_audio, input_audio

def clear_fields(input_text, output_text, translated_audio_output, input_audio_output, translate_checkbox, text_to_speech_checkbox, input_to_speech_checkbox):
    """Clear input and output fields and uncheck option boxes."""
    return "", "", None, None, False, False, False  # Clear and set checkboxes to False


# Gradio Interface
with gr.Blocks() as demo:
  
    # Add a caption at the top
    gr.Markdown("<h1 style='text-align: center;'>Multilingual T2T/Speech Translator</h1>")
    with gr.Row():
        input_text = gr.Textbox(label="Input Text", placeholder="Enter text here...")
        output_text = gr.Textbox(label="Translation", interactive=False)
    
    with gr.Row():
        language_dropdown = gr.Dropdown(label="Select Output Language", choices=list(language_options.keys()))
    
    with gr.Row():
        translate_checkbox = gr.Checkbox(label="Translate")
        text_to_speech_checkbox = gr.Checkbox(label="T2S - Translation")
        input_to_speech_checkbox = gr.Checkbox(label="T2S - Actual-Text")
    
    with gr.Row():
        generate_btn = gr.Button("Generate")
        clear_btn = gr.Button("Clear")

        # Define Audio components outside the click events
    translated_audio_output = gr.Audio(label="T2S - Translation", interactive=False)  
    input_audio_output = gr.Audio(label="T2S - Actual Text", interactive=False)
    
    # Generate button action
    generate_btn.click(generate, 
                       inputs=[input_text, language_dropdown, translate_checkbox, text_to_speech_checkbox, input_to_speech_checkbox], 
                       outputs=[output_text, translated_audio_output, input_audio_output])  # Update existing components
    
    # Clear button action (Modified)
    def clear_audio_output(audio_component):
        return None  # Return None to clear the audio output
    
    clear_btn.click(
        fn=clear_fields,
        inputs=None,
        outputs=[input_text, output_text, translate_checkbox, text_to_speech_checkbox, input_to_speech_checkbox])  # Excluded audio components
    clear_btn.click(
        fn=clear_audio_output,
        inputs=translated_audio_output,
        outputs=translated_audio_output)
    clear_btn.click(
        fn=clear_audio_output,
        inputs=input_audio_output,
        outputs=input_audio_output)
    
    # Add footnote
    gr.Markdown("<p style='text-align: center; font-size: 12px;'>By Khalid Baig - Thanks to Aspire Pakistan</p>")

    demo.launch()
    share=True