import gradio as gr
from transformers import pipeline

# Language codes and names (only ISO 639-1 supported by Helsinki models)
languages = {
    "en": "English",
    "fr": "French",
    "de": "German",
    "es": "Spanish",
    "it": "Italian",
    "ru": "Russian",
    "zh": "Chinese",
    "ja": "Japanese",
    "ko": "Korean",
    "ar": "Arabic",
    "hi": "Hindi",
    "te": "Telugu",
    "ta": "Tamil",
    "bn": "Bengali",
    "ur": "Urdu",
    "ml": "Malayalam",
    "gu": "Gujarati",
    "mr": "Marathi",
    "pa": "Punjabi",
    "sw": "Swahili",
    "tr": "Turkish",
    "fa": "Persian",
    "pt": "Portuguese",
    "pl": "Polish",
    "uk": "Ukrainian",
    "nl": "Dutch",
    "sv": "Swedish",
    "no": "Norwegian",
    "fi": "Finnish",
    "cs": "Czech",
    "el": "Greek",
    "ro": "Romanian",
    "hu": "Hungarian",
    "vi": "Vietnamese",
    "th": "Thai",
    "id": "Indonesian",
    "ms": "Malay",
    "he": "Hebrew",
    "bg": "Bulgarian",
    "sr": "Serbian",
    "hr": "Croatian",
    "sk": "Slovak",
    "da": "Danish",
    "et": "Estonian",
    "lt": "Lithuanian",
    "lv": "Latvian",
    "sl": "Slovenian",
    "mt": "Maltese"
}

# Reverse mapping from full name to code
name_to_code = {v: k for k, v in languages.items()}

# Define the function to translate

def translate_text(text, source_lang_name, target_lang_name):
    if source_lang_name == target_lang_name:
        return text

    source_lang = name_to_code[source_lang_name]
    target_lang = name_to_code[target_lang_name]

    if source_lang == "en" and target_lang == "te":
        model_name = "Helsinki-NLP/opus-mt-en-dra"
    else:
        model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"

    try:
        trans_pipeline = pipeline("translation", model=model_name)
        translation = trans_pipeline(text, max_length=400)[0]['translation_text']
        return translation
    except Exception as e:
        return f"Error loading model for {source_lang} to {target_lang}: {str(e)}"

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🌐 Multi-language Translation App")
    input_text = gr.Textbox(label="Enter Text", lines=4)
    source_lang = gr.Dropdown(label="Source Language", choices=list(languages.values()), value="English")
    target_lang = gr.Dropdown(label="Target Language", choices=list(languages.values()), value="Telugu")
    output_text = gr.Textbox(label="Translated Text")

    translate_button = gr.Button("Translate")
    translate_button.click(fn=translate_text, inputs=[input_text, source_lang, target_lang], outputs=output_text)

if __name__ == "__main__":
    demo.launch()