import gradio as gr
from transformers import pipeline
import json

# Initialize the translation pipeline with the specific model
text_translator = pipeline("translation", model="facebook/nllb-200-distilled-600M")

# Load the JSON table containing language mappings
with open('language.json') as f:
    language_data = json.load(f)

def get_flores_200_code(language):
    """
    Retrieves the FLORES-200 code for a given language from the loaded JSON data.

    Args:
        language (str): The name of the language.

    Returns:
        str: The FLORES-200 code for the language, or None if not found.
    """
    for code in language_data:
        if code['Language'] == language:
            return code['FLORES-200 code']
    return None

def translate_text(text, destination_language):
    """
    Translates text from English to the selected destination language using the T5 model.

    Args:
        text (str): The text to translate.
        destination_language (str): The target language for translation.

    Returns:
        str: The translated text.
    """
    dest_code = get_flores_200_code(destination_language)

    if not dest_code:
        return "Unsupported language selected. Please choose a valid language."

    try:
        # Perform translation using T5 model pipeline
        translation = text_translator(text,
                                      src_lang="eng_Latn",
                                      tgt_lang=dest_code)
        translated_text = translation[0]["translation_text"]
        return translated_text
    except Exception as e:
        return f"Translation error: {str(e)}"

def main():
    # Create Gradio interface
    interface = gr.Interface(
        fn=translate_text,
        inputs=[
            gr.Textbox(label="Input text to translate", lines=6, placeholder="Enter text in English..."),
            gr.Dropdown([code['Language'] for code in language_data], label="Select destination language")
        ],
        outputs=[gr.Textbox(label="Translated text", lines=4)],
        title="Multi-Language Translator",
        description="Translate English text to multiple languages using the T5 model. Select the target language from the dropdown menu and enter the text to translate.",
        theme=gr.themes.Soft(),
        live=True  # Enable live updates
    )

    # Launch the Gradio interface
    interface.launch()

if __name__ == "__main__":
    main()