# app.py (Translation Space)
import gradio as gr
from transformers import pipeline
import asyncio

# Initialize models
try:
    translator_ar_en = pipeline("translation_ar_to_en", model="Helsinki-NLP/opus-mt-ar-en")
    translator_en_ar = pipeline("translation_en_to_ar", model="Helsinki-NLP/opus-mt-en-ar")
    print("Translation models loaded successfully!")
except Exception as e:
    print(f"Error loading translation models: {e}")
    raise  # Stop execution if models fail to load

# Translation function (Batched)
async def translate(texts, source_lang, target_lang):
    try:
        if not texts:
            return []  # Handle empty input

        if source_lang == "ar" and target_lang == "en":
            translations = translator_ar_en(texts)
            return [translation['translation_text'] for translation in translations]
        elif source_lang == "en" and target_lang == "ar":
            translations = translator_en_ar(texts)
            return [translation['translation_text'] for translation in translations]
        else:
            return ["Invalid language combination"] * len(texts)  # Error for each input
    except Exception as e:
        print(f"Error in translation: {e}")
        return ["Translation Error"] * len(texts)  # Indicate error for each input

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌍 Language Translation")

    with gr.Row():
        source_lang = gr.Dropdown(["ar", "en"], label="Source Language")
        target_lang = gr.Dropdown(["ar", "en"], label="Target Language")

    text_input = gr.Textbox(label="Input Text (Separate multiple texts with newlines)", lines=5)

    translation_output = gr.Textbox(label="Translated Text", lines=5)

    submit = gr.Button("Translate", variant="primary")

    async def process_translation(text_input, source_lang, target_lang):
        texts = text_input.split("\n")  # Split input into multiple texts
        translations = await translate(texts, source_lang, target_lang)
        return "\n".join(translations)  # Join translations with newlines

    submit.click(process_translation,
                 inputs=[text_input, source_lang, target_lang],
                 outputs=[translation_output])

demo.launch()