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

# Mapping user-selected languages to available translation models
TRANSLATION_MODELS = {
    "fr": "Helsinki-NLP/opus-mt-en-fr",
    "es": "Helsinki-NLP/opus-mt-en-es",
    "de": "Helsinki-NLP/opus-mt-en-de",
    "it": "Helsinki-NLP/opus-mt-en-it",
    "zh": "Helsinki-NLP/opus-mt-en-zh",
    "ar": "Helsinki-NLP/opus-mt-en-ar",
    "hi": "Helsinki-NLP/opus-mt-en-hi",
    "ur": "Helsinki-NLP/opus-mt-en-ur"
}

# Language mapping for gTTS
TTS_LANGUAGES = {
    "fr": "fr",
    "es": "es",
    "de": "de",
    "it": "it",
    "zh": "zh",
    "ar": "ar",
    "hi": "hi",
    "ur": "ur"
}

def translate_and_speak(text, target_language):
    """ Translates input text to the selected language and converts it to speech. """
    
    if not text.strip():
        return "Please enter some text.", None

    if target_language not in TRANSLATION_MODELS:
        return "Translation model not available for this language.", None

    try:
        # Load translation model dynamically based on user choice
        translator = pipeline("translation", model=TRANSLATION_MODELS[target_language])
        translated_text = translator(text)[0]['translation_text']

        # Convert text to speech using gTTS
        output_path = "output.mp3"
        tts = gTTS(text=translated_text, lang=TTS_LANGUAGES[target_language])
        tts.save(output_path)

        return translated_text, output_path
    
    except Exception as e:
        return f"Error: {str(e)}", None

# Gradio Interface
demo = gr.Interface(
    fn=translate_and_speak,
    inputs=[
        gr.Textbox(label="Enter Text"),
        gr.Dropdown(["fr", "es", "de", "it", "zh", "ar", "hi", "ur"], label="Target Language")
    ],
    outputs=[
        gr.Textbox(label="Translated Text"),
        gr.Audio(label="Generated Speech")
    ]
)

# Launch Gradio App
demo.launch(share=True)