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

# Use a pipeline as a high-level helper
from transformers import pipeline

pipe = pipeline("translation", model="facebook/nllb-200-distilled-600M",torch_dtype=torch.bfloat16)

# Load language data
with open('le.json','r',encoding='utf-8') as f:
    language_data = json.load(f)

# Prepare language map and list
lang_map = {item['Language'].lower(): item['FLORES-200 code'] for item in language_data}
all_languages = sorted(lang_map.keys())

def lang_to_code(lang):
    try:
        return lang_map[lang.lower()]
    except KeyError:
        raise ValueError(f"Unsupported language: {lang}")

# Find default languages in the list
def find_default(lang_name, choices):
    for lang in choices:
        if lang_name.lower() in lang.lower():
            return lang
    return choices[0]

default_src = find_default("English", all_languages)
default_tgt = find_default("Bengali", all_languages)

# Translator function
def translator(text, sourc_lang, desti_lang):
    if not sourc_lang:
        sourc_lang = default_src
    if not desti_lang:
        desti_lang = default_tgt
    sourc_lang = lang_to_code(sourc_lang)
    desti_lang = lang_to_code(desti_lang)
    translated_text = pipe(text, src_lang=sourc_lang, tgt_lang=desti_lang)
    return translated_text[0]['translation_text']

# --- Gradio UI: Must be outside the translator function ---
with gr.Blocks(title='Multi-Language Translator') as demo:
    gr.Markdown('## 🌍 Multi-Language Translator')

    with gr.Row():
        with gr.Column():
            src_text = gr.Textbox(label='Source Text', lines=6, placeholder='Type text here...')
            src_lang = gr.Dropdown(choices=all_languages, value=default_src, label='Source Language', filterable=True)

        with gr.Column():
            tgt_text = gr.Textbox(label='Translated Text', lines=6)
            tgt_lang = gr.Dropdown(choices=all_languages, value=default_tgt, label='Target Language', filterable=True)

    translate_btn = gr.Button("Translate 🚀")
    translate_btn.click(fn=translator, inputs=[src_text, src_lang, tgt_lang], outputs=[tgt_text])

# Launch the app
demo.launch()