import gradio as gr
from transformers import pipeline

# Load models locally on the Hugging Face Space server
# Note: In a real space, you can also use Inference Endpoints for even more speed
models = {
    "Spanish": pipeline("translation_en_to_es", model="Helsinki-NLP/opus-mt-en-es"),
    "German": pipeline("translation_en_to_de", model="Helsinki-NLP/opus-mt-en-de"),
    "Japanese": pipeline("translation_en_to_ja", model="staka/fugumt-en-ja"),
    "Ukranian": pipeline("translation_en_to_uk", model="Helsinki-NLP/opus-mt-en-uk"),
    "Russian": pipeline("translation_en_to_ru", model="Helsinki-NLP/opus-mt-en-ru"),
}

def translate(text, language):
    if not text:
        return "Please enter an English sentence."
    
    # Perform translation using the selected model
    result = models[language](text)
    return result[0]['translation_text']

# Create the Gradio interface
with gr.Blocks(title="Short Translation Web App") 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=["Spanish", "German", "Japanese", "Ukranian", 
                         "Russian"], 
                label="Translation Language",
                value="Spanish"
            )
    
    # Output area styled to resemble the maroon box (using Gradio's color system)
    output_text = gr.Textbox(
        label="Translation", 
        interactive=False,
        container=True,
        elem_id="output-box"
    )

    # Custom CSS to mimic the maroon look from the Tkinter app
    demo.css = """
    #output-box textarea {
        background-color: #800000 !important;
        color: white !important;
        font-size: 1.2em !important;
        font-weight: bold !important;
        text-align: center !important;
    }
    """

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

# Launch the app
if __name__ == "__main__":
    demo.launch()