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

# Use a pipeline as a high-level helper
text_translator = pipeline("translation", model="facebook/nllb-200-distilled-600M", torch_dtype=torch.bfloat16)

# Load the JSON data from the file
with open('language.json', 'r') as file:
    language_data = json.load(file)


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


def translate_text(text, destination_language):
    dest_code = get_FLORES_code_from_language(destination_language)
    print(f"Destination Language: {destination_language}, Code: {dest_code}")
    translation = text_translator(text, src_lang="eng_Latn", tgt_lang=dest_code)
    translated_text = translation[0]["translation_text"]
    print(f"Translated Text: {translated_text}")

    # For Arabic, add HTML to force RTL display
    if destination_language.lower() in ["egyptian arabic", "arabic"]:
        translated_text = f'<div style="text-align: right; direction: rtl;">{translated_text}</div>'

    return translated_text


# Create the Gradio interface
def translate(text, destination_language):
    return translate_text(text, destination_language)


language_options = [entry['Language'] for entry in language_data]

iface = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter text here..."),
        gr.Dropdown(choices=language_options, value="English", label="Destination Language")
    ],
    outputs=gr.HTML(label="Translated Text"),
    title="Text Translator",
    description="Enter text and choose the destination language to get the translation."
)

iface.launch()