import gradio as gr
from transformers import pipeline
import json
from langdetect import detect

# 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 detect_language(text):
    """
    Detects the language of the given text.

    Args:
        text (str): The text to detect language from.

    Returns:
        str: The detected language code.
    """
    try:
        lang_code = detect(text)
        return lang_code
    except Exception as e:
        return str(e)

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

    Args:
        text (str): The text to translate.
        source_language (str): The source language code.
        destination_language (str): The target language code.

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

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

    try:
        # Perform translation using T5 model pipeline
        translation = text_translator(text,
                                      src_lang=src_code,
                                      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..."),
            gr.Dropdown([code['Language'] for code in language_data], label="Select source language"),
            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 text between multiple languages using the T5 model. Select the source and destination languages from the dropdown menus.",
        theme=gr.themes.Soft(),
        live=True  # Enable live updates
    )

    # Launch the Gradio interface
    interface.launch()

if __name__ == "__main__":
    main()