from transformers import pipeline
import gradio as gr

# Initialize translation pipeline
translator = pipeline("translation", model="facebook/m2m100_418M")

def translate_text(text, src_lang_code, tgt_lang_code):
    if not text.strip():
        return ""
    
    # Set source language
    translator.tokenizer.src_lang = src_lang_code
    
    # Translate
    result = translator(
        text,
        src_lang=src_lang_code,
        tgt_lang=tgt_lang_code,
        clean_up_tokenization_spaces=True
    )
    return result[0]['translation_text']

# Language dictionary with native text examples
LANGUAGES = {
    "Afrikaans - Afrikaans": "af",
    "Albanian - Shqip": "sq",
    "Amharic - አማርኛ": "am",
    "Arabic - العربية": "ar",
    "Armenian - Հայերեն": "hy",
    "Azerbaijani - Azərbaycanca": "az",
    "Basque - Euskara": "eu",
    "Belarusian - Беларуская": "be",
    "Bengali - বাংলা": "bn",
    "Bosnian - Bosanski": "bs",
    "Bulgarian - Български": "bg",
    "Catalan - Català": "ca",
    "Cebuano - Cebuano": "ceb",
    "Chichewa - Nyanja": "ny",
    "Chinese Simplified - 中文": "zh",
    "Chinese Traditional - 中文": "zh",
    "Corsican - Corsu": "co",
    "Croatian - Hrvatski": "hr",
    "Czech - Čeština": "cs",
    "Danish - Dansk": "da",
    "Dutch - Nederlands": "nl",
    "English - English": "en",
    "Esperanto - Esperanto": "eo",
    "Estonian - Eesti": "et",
    "Filipino - Filipino": "tl",
    "Finnish - Suomi": "fi",
    "French - Français": "fr",
    "Frisian - Frysk": "fy",
    "Galician - Galego": "gl",
    "Georgian - ქართული": "ka",
    "German - Deutsch": "de",
    "Greek - Ελληνικά": "el",
    "Gujarati - ગુજરાતી": "gu",
    "Haitian Creole - Kreyòl ayisyen": "ht",
    "Hausa - Hausa": "ha",
    "Hawaiian - ʻŌlelo Hawaiʻi": "haw",
    "Hebrew - עברית": "he",
    "Hindi - हिन्दी": "hi",
    "Hmong - Hmong": "hmn",
    "Hungarian - Magyar": "hu",
    "Icelandic - Íslenska": "is",
    "Igbo - Asụsụ Igbo": "ig",
    "Indonesian - Bahasa Indonesia": "id",
    "Irish - Gaeilge": "ga",
    "Italian - Italiano": "it",
    "Japanese - 日本語": "ja",
    "Javanese - Basa Jawa": "jv",
    "Kannada - ಕನ್ನಡ": "kn",
    "Kazakh - Қазақша": "kk",
    "Khmer - ភាសាខ្មែរ": "km",
    "Kinyarwanda - Kinyarwanda": "rw",
    "Korean - 한국어": "ko",
    "Kurdish - Kurdî": "ku",
    "Kyrgyz - Кыргызча": "ky",
    "Lao - ພາສາລາວ": "lo",
    "Latin - Latina": "la",
    "Latvian - Latviešu": "lv",
    "Lithuanian - Lietuvių": "lt",
    "Luxembourgish - Lëtzebuergesch": "lb",
    "Macedonian - Македонски": "mk",
    "Malagasy - Malagasy": "mg",
    "Malay - Bahasa Melayu": "ms",
    "Malayalam - മലയാളം": "ml",
    "Maltese - Malti": "mt",
    "Maori - Te Reo Māori": "mi",
    "Marathi - मराठी": "mr",
    "Mongolian - Монгол": "mn",
    "Myanmar - မြန်မာဘာသာ": "my",
    "Nepali - नेपाली": "ne",
    "Norwegian - Norsk": "no",
    "Odia - ଓଡ଼ିଆ": "or",
    "Pashto - پښتو": "ps",
    "Persian - فارسی": "fa",
    "Polish - Polski": "pl",
    "Portuguese - Português": "pt",
    "Punjabi - ਪੰਜਾਬੀ": "pa",
    "Romanian - Română": "ro",
    "Russian - Русский": "ru",
    "Samoan - Gagana Samoa": "sm",
    "Scots Gaelic - Gàidhlig": "gd",
    "Serbian - Српски": "sr",
    "Sesotho - Sesotho": "st",
    "Shona - ChiShona": "sn",
    "Sindhi - سنڌي": "sd",
    "Sinhala - සිංහල": "si",
    "Slovak - Slovenčina": "sk",
    "Slovenian - Slovenščina": "sl",
    "Somali - Soomaali": "so",
    "Spanish - Español": "es",
    "Sundanese - Basa Sunda": "su",
    "Swahili - Kiswahili": "sw",
    "Swedish - Svenska": "sv",
    "Tajik - Тоҷикӣ": "tg",
    "Tamil - தமிழ்": "ta",
    "Tatar - Татарча": "tt",
    "Telugu - తెలుగు": "te",
    "Thai - ไทย": "th",
    "Turkish - Türkçe": "tr",
    "Turkmen - Türkmençe": "tk",
    "Ukrainian - Українська": "uk",
    "Urdu - اردو": "ur",
    "Uyghur - ئۇيغۇرچە": "ug",
    "Uzbek - O'zbekcha": "uz",
    "Vietnamese - Tiếng Việt": "vi",
    "Welsh - Cymraeg": "cy",
    "Xhosa - IsiXhosa": "xh",
    "Yiddish - יידיש": "yi",
    "Yoruba - Èdè Yorùbá": "yo",
    "Zulu - IsiZulu": "zu"
}

# Create sorted lists for dropdowns
language_names = sorted(LANGUAGES.keys())
default_src = "English - English"
default_tgt = "Spanish - Español"

# Gradio interface
with gr.Blocks(title="Universal Translator") as demo:
    gr.Markdown("# 🌍 Universal Language Translator")
    gr.Markdown("Translate between 100+ languages using Hugging Face's M2M100 model")
    
    with gr.Row():
        with gr.Column():
            src_lang = gr.Dropdown(
                choices=language_names,
                value=default_src,
                label="Source Language"
            )
            input_text = gr.Textbox(
                lines=6,
                placeholder="Enter text to translate...",
                label="Input Text"
            )
            
        with gr.Column():
            tgt_lang = gr.Dropdown(
                choices=language_names,
                value=default_tgt,
                label="Target Language"
            )
            output_text = gr.Textbox(
                lines=6,
                interactive=False,
                label="Translation"
            )
    
    translate_btn = gr.Button("Translate", variant="primary")
    
    def translate_wrapper(text, src_lang_name, tgt_lang_name):
        src_code = LANGUAGES[src_lang_name]
        tgt_code = LANGUAGES[tgt_lang_name]
        return translate_text(text, src_code, tgt_code)
    
    translate_btn.click(
        translate_wrapper,
        inputs=[input_text, src_lang, tgt_lang],
        outputs=output_text
    )
    
    gr.Examples(
        [
            ["Hello, how are you today?", "English - English", "French - Français"],
            ["The weather is beautiful", "English - English", "German - Deutsch"],
            ["I love programming with Python", "English - English", "Spanish - Español"],
            ["Machine learning is fascinating", "English - English", "Chinese Simplified - 中文"],
            ["こんにちは、元気ですか？", "Japanese - 日本語", "English - English"],
            ["Bonjour, comment allez-vous?", "French - Français", "Russian - Русский"]
        ],
        inputs=[input_text, src_lang, tgt_lang]
    )
    
    gr.Markdown("### Supported Languages")
    gr.Markdown(f"- **Total Languages**: {len(language_names)}")
    gr.Markdown("- **Model**: [facebook/m2m100_418M](https://huggingface.co/facebook/m2m100_418M)")
    gr.Markdown("- **Tip**: Each language shows how it's written in its own script (e.g., Russian - Русский)")

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