"""
File: module_translation.py
Description: Translating text with a multimodal interface.
Author: Didier Guillevic
Date: 2025-03-16
"""

import gradio as gr
import vlm
from deep_translator import GoogleTranslator

tgt_language_codes = {
    'English': 'en',
    'French': 'fr'
}
code_to_languages = {v: k for k, v in tgt_language_codes.items()}

translation_prompt = (
    "Please translate the following text into {}. "
    "Do not include any other information. "
    "Output only the translation of the given text. "
    "Text:\n\n{}"
)

def translate_text(text, tgt_lang):
    """Translate the given text into the given target language.
    """
    # Build messages
    messages = [
        {
            'role': 'user',
            'content': [
                {
                    "type": "text",
                    "text": translation_prompt.format(
                        code_to_languages[tgt_lang], text)
                }
            ]
        }
    ]

    # Get the translation
    yield from vlm.stream_response(messages)

def translate_text_with_google(text, tgt_lang):
    """Translate the given text into the given target language using Google Translate.
    """
    translated_text_google = GoogleTranslator(
        source='auto', target=tgt_lang).translate(text=text)
    
    return translated_text_google


#
# User interface
#
with gr.Blocks() as demo:
    with gr.Row():
        input_text = gr.Textbox(
            lines=5,
            placeholder="Enter text to translate",
            label="Text to translate",
            render=True
        )
    
    with gr.Row():   
        output_text = gr.Textbox(
            lines=5,
            label=vlm.model_id,
            render=True
        )
        output_text_google = gr.Textbox(
            lines=5,
            label="Google Translate",
            render=True
        )
    
    with gr.Row():
        tgt_lang = gr.Dropdown(
            choices=tgt_language_codes.items(),
            value="en",
            label="Target language",
            render=True
        )
        translate_btn = gr.Button(value="Translate (Mistral)", variant="primary")
        translate_google_btn = gr.Button(value="Translate (Google)", variant="primary")
        clear_btn = gr.Button("Clear", variant="secondary")
    
    # Examples
    with gr.Accordion("Examples", open=False):
        examples = gr.Examples(
            [
                ["""بی بی سی ویریفائی کا تجزیہ

بی بی سی ویریفائی نے تین ایسی ویڈیوز کی تصدیق کی ہے جن کے بارے میں دعوی کیا گیا تھا کہ ان میں نظر آنے والا ملبہ ایک فرانسیسی ساختہ رفال طیارے کا ہے جو انڈیا کی فضائیہ کے زیراستعمال ہیں۔
ان میں سے ایک ویڈیو کی جیو لوکیشن سے بی بی سی ویریفائی کو علم ہوا کہ یہ انڈین ریاست پنجاب میں بٹھنڈہ کے مقام کی ہے۔ اس ویڈیو میں یونیفارم میں ملبوس اہلکار لڑاکا طیارے کا ملبہ اکھٹا کرتے ہوئے دیکھے جا سکتے ہیں۔
بی بی سی ویریفائی نے رات کے وقت بنائی جانے والی دو مذید ویڈیوز کو بھی چیک کیا اور یہ بھی اسی مقام پر بنی ہیں۔ ان میں سے ایک میں کھیتوں میں موجود ملبہ دیکھا جا سکتا ہے جبکہ ایک اور ویڈیو میں دیکھا جا سکتا ہے کہ ایک اڑنے والی چیز کو پہلے آسمان میں آگ لگتی ہے اور پھر وہ ایک کھلے کھیت میں گرتی ہے۔
جسٹن کرمپ برطانوی فوج کے سابق افسر ہیں جو اب رسک انٹیلیجنس کمپنی سبیلائن چلاتے ہیں۔ انھوں نے بی بی سی ویرفائی سے بات کرتے ہوئے کہا کہ ویڈیو میں نظر آنے والا یہ ملبہ فرانسیسی میزائل کا لگتا ہے جو میراج دو ہزار اور رفال لڑاکا طیاروں، دونوں میں ہی استعمال ہوتا ہے۔
سوشل میڈیا پر گردش کرنے والی ایک اور تصویر میں ایک لڑاکا طیارے کی ٹیل فن، جس پر بی ایس زیرو زیرو ون اور رفال لکھا ہوا ہے، دیکھی جا سکتی ہے۔
بی بی سی ویریفائی نے اس تصویر کو جانچنے کے لیے گوگل ریورس امیج کا استعمال کیا تو معلوم ہوا کہ اس تصویر کا اس سے پہلے کوئی اور ورژن نہیں ہے۔""", "en"],
                ["ریچارد مور، رئیس سازمان مخفی اطلاعاتی بریتانیا (ام‌آی‌۶) در دیدار ویلیام برنز، رئیس سازمان اطلاعات مرکزی آمریکا (سیا) گفت همچنان احتمال اقدام ایران علیه اسرائیل در واکنش به ترور اسماعیل هنیه، رهبر حماس وجود دارد. آقای برنز نیز در این دیدار فاش کرد که در سال اول جنگ اوکراین، «خطر واقعی» وجود داشت که روسیه به استفاده از «تسلیحات هسته‌ای تاکتیکی» متوسل شود. این دو مقام امنیتی هشدار دادند که «نظم جهانی» از زمان جنگ سرد تا کنون تا این حد «در معرض تهدید» نبوده است.", "en"],
                ["Clément Delangue est, avec Julien Chaumond et Thomas Wolf, l’un des trois Français cofondateurs de Hugging Face, une start-up d’intelligence artificielle (IA) de premier plan. Valorisée à 4,2 milliards d’euros après avoir levé près de 450 millions d’euros depuis sa création en 2016, cette société de droit américain est connue comme la plate-forme de référence où développeurs et entreprises publient des outils et des modèles pour faire de l’IA en open source, c’est-à-dire accessible gratuitement et modifiable.", "en"],
                ["يُعد تفشي مرض جدري القردة قضية صحية عالمية خطيرة، ومن المهم محاولة منع انتشاره للحفاظ على سلامة الناس وتجنب العدوى. د. صموئيل بولاند، مدير الحوادث الخاصة بمرض الجدري في المكتب الإقليمي لمنظمة الصحة العالمية في أفريقيا، يتحدث من كينشاسا في جمهورية الكونغو الديمقراطية، ولديه بعض النصائح البسيطة التي يمكن للناس اتباعها لتقليل خطر انتشار المرض.", "en"],
                ["【ワシントン＝冨山優介】米ボーイングの新型宇宙船「スターライナー」は７日午前０時（日本時間７日午後１時）過ぎ、米ニューメキシコ州のホワイトサンズ宇宙港に着地し、地球に帰還した。スターライナーは米宇宙飛行士２人を乗せて６月に打ち上げられ、国際宇宙ステーション（ＩＳＳ）に接続したが、機体のトラブルが解決できず、無人でＩＳＳから離脱した。", "en"],
                ["張先生稱，奇瑞已經凖備在西班牙生產汽車，並決心採取「本地化」的方式進入歐洲市場。此外，他也否認該公司的出口受益於不公平補貼。奇瑞成立於1997年，是中國最大的汽車公司之一。它已經是中國最大的汽車出口商，並且制定了進一步擴張的野心勃勃的計劃。", "en"],
                ["ברוכה הבאה, קיטי: בית הקפה החדש בלוס אנג'לס החתולה האהובה והחברים שלה מקבלים בית קפה משלהם בשדרות יוניברסל סיטי, שם תוכלו למצוא מגוון של פינוקים מתוקים – החל ממשקאות ועד עוגות", "en"],
            ],
            inputs=[input_text, tgt_lang],
            outputs=[output_text,],
            fn=translate_text,
            cache_examples=False,
            label="Examples"
        )

    # Click actions
    translate_btn.click(
        fn=translate_text,
        inputs=[input_text, tgt_lang],
        outputs=[output_text,]
    )
    translate_google_btn.click(
        fn=translate_text_with_google,
        inputs=[input_text, tgt_lang],
        outputs=[output_text_google,]
    )
    clear_btn.click(
        fn=lambda : ('', '', ''), # input_text, output_text, output_text_google
        inputs=[],
        outputs=[input_text, output_text, output_text_google]
    )

    with gr.Accordion("Documentation", open=False):
        gr.Markdown(f"""
            - Model: serving {vlm.model_id} and Google Translate.
        """)