import gradio as gr
from .translation_model import TranslationModel
import time

def translator_component():
    # Initialize the translation model
    model = TranslationModel()
    
    # Language mapping
    LANGUAGES = {
        "Afrikaans": "af",
        "English": "en",
    }
    
    def translate_text(text, source_lang, target_lang):
        if not text.strip():
            return "Please enter text to translate."
        
        start_time = time.time()
        src_code = LANGUAGES[source_lang]
        tgt_code = LANGUAGES[target_lang]
        
        result = model.translate(text, src_code, tgt_code)
        end_time = time.time()
        
        return f"{result}\n\nTranslation time: {round(end_time - start_time, 2)} seconds"
    
    with gr.Column() as translator:
        gr.Markdown("### Neural Machine Translation")
        gr.Markdown("Using M2M100 1.2B model for high-quality translations")
        
        input_text = gr.Textbox(
            label="Text to Translate",
            placeholder="Enter text here...",
            lines=3
        )
        
        with gr.Row():
            source_lang = gr.Dropdown(
                choices=list(LANGUAGES.keys()),
                value="English",
                label="From"
            )
            target_lang = gr.Dropdown(
                choices=list(LANGUAGES.keys()),
                value="Afrikaans",
                label="To"
            )
        
        translate_btn = gr.Button("Translate")
        output_text = gr.Textbox(label="Translation", lines=3)
        
        translate_btn.click(
            fn=translate_text,
            inputs=[input_text, source_lang, target_lang],
            outputs=output_text
        )
        
    return translator