import torch
import gradio as gr
import json
from transformers import pipeline

# Initialize translation pipeline (online or offline)
text_translator = pipeline("translation", model="facebook/nllb-200-distilled-600M",
                           torch_dtype=torch.bfloat16)

# Load the Flores-200 language codes from a JSON file
with open('language.json', 'r') as file:
    language_data = json.load(file)

# Helper function to get the Flores code from the full language name
def get_FLORES_code_from_language(language):
    for entry in language_data:
        if entry['Language'].lower() == language.lower():
            return entry['FLORES-200 code']
    return None

# Translation function used by the Gradio interface
def translate_text(text, destination_language):
    dest_code = get_FLORES_code_from_language(destination_language)
    if not dest_code:
        return f"❌ Language code for '{destination_language}' not found."
    translation = text_translator(text, src_lang="eng_Latn", tgt_lang=dest_code)
    return translation[0]["translation_text"]

# Build the dropdown list from the loaded language.json
language_list = [entry["Language"] for entry in language_data]

# Create a polished Gradio UI
with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown(
        """
        # 🌐 Multilingual Translator — Powered by Meta NLLB-200  
        Translate English text into 200+ languages using the FLORES-200 dataset.
        """
    )

    with gr.Row():
        input_text = gr.Textbox(label="📝 English Input", lines=6, placeholder="Type your English sentence here...")
        output_lang = gr.Dropdown(choices=language_list, label="🌍 Target Language")

    output_text = gr.Textbox(label="✅ Translated Output", lines=4)

    translate_btn = gr.Button("🔄 Translate")

    translate_btn.click(fn=translate_text, inputs=[input_text, output_lang], outputs=output_text)

    gr.Markdown(
        """
        ---
        🧠 Model: `facebook/nllb-200-distilled-600M`  
        🌐 Language codes from [FLORES-200](https://github.com/facebookresearch/flores)
        """
    )

demo.launch()