import gradio as gr
from transformers import pipeline

MODEL_SPECS = {
    "Spanish":   ("translation_en_to_es", "Helsinki-NLP/opus-mt-en-es"),
    "German":    ("translation_en_to_de", "Helsinki-NLP/opus-mt-en-de"),
    "Japanese":  ("translation_en_to_ja", "staka/fugumt-en-ja"),
    "Ukrainian": ("translation_en_to_uk", "Helsinki-NLP/opus-mt-en-uk"),
    "Russian":   ("translation_en_to_ru", "Helsinki-NLP/opus-mt-en-ru"),
}

CACHED_MODELS = {}

def get_model(language: str):
    if language not in CACHED_MODELS:
        task, model_name = MODEL_SPECS[language]
        CACHED_MODELS[language] = pipeline(task, model=model_name)
    return CACHED_MODELS[language]

def translate(text, language):
    text = (text or "").strip()
    if not text:
        return "Please enter an English sentence."

    model = get_model(language)
    result = model(text)
    return result[0]["translation_text"]

CSS = """
#output-box textarea {
    background-color: #800000 !important;
    color: white !important;
    font-size: 1.2em !important;
    font-weight: bold !important;
    text-align: center !important;
}
"""

with gr.Blocks(title="Short Translation Web App", css=CSS) as demo:
    gr.Markdown("# Short Translation")

    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                label="English Sentence",
                placeholder="Type your English sentence here...",
                lines=2
            )
            with gr.Row():
                translate_btn = gr.Button("Translate", variant="primary")
                clear_btn = gr.Button("Clear")

        with gr.Column(scale=1):
            lang_radio = gr.Radio(
                choices=list(MODEL_SPECS.keys()),
                label="Translation Language",
                value="Spanish"
            )

    output_text = gr.Textbox(
        label="Translation",
        interactive=False,
        container=True,
        elem_id="output-box"
    )

    translate_btn.click(fn=translate, inputs=[input_text, lang_radio], outputs=output_text)
    clear_btn.click(fn=lambda: ("", "Spanish", ""), inputs=None, outputs=[input_text, lang_radio, output_text])

if __name__ == "__main__":
    demo.launch()