import gradio as gr
from transformers import pipeline

# Load AI translation models
translator_models = {
    "french": pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr"),
    "spanish": pipeline("translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es"),
    "german": pipeline("translation_en_to_de", model="Helsinki-NLP/opus-mt-en-de"),
    "arabic": pipeline("translation_en_to_ar", model="Helsinki-NLP/opus-mt-en-ar"),
}

# Function to translate messages
def chatbot_translate(message, history):
    try:
        # Normalize input (convert to lowercase and remove leading/trailing spaces)
        message = message.strip()
        lower_message = message.lower()

        # Check if the user specified a language
        for lang in translator_models.keys():
            if lower_message.startswith(f"{lang}:"):
                text_to_translate = message[len(lang) + 1 :].strip()  # Extract text after "Language:"
                translation = translator_models[lang](text_to_translate)[0]["translation_text"]
                return translation

        # If no valid language is specified, prompt the user
        return "Please specify a language (e.g., 'French: Hello' or 'Spanish: Hello')."

    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio chatbot UI
iface = gr.ChatInterface(
    chatbot_translate,
    title="🌍 AI Translator Bot",
    description="Type 'French: [your text]' or 'Spanish: [your text]' to translate!\nSupported: French, Spanish, German, Arabic",
)

iface.launch()