import gradio as gr
from nllb import translation
from flores200_codes import flores_codes

lang_codes = list(flores_codes.keys())

def text_translation(model, text):
    """
    Determines languages based on model, calls the translation function,
    and returns only the translated text.
    """
    if model == "Ɔbaa Panin MT: Akan-English":
        source_lang = "Akan"
        target_lang = "English"
    elif model == "Ɔbaa Panin MT: English-Akan":
        source_lang = "English"
        target_lang = "Akan"
    else:
        return "Error: Invalid model selected."
    
    result = translation(model, source_lang, target_lang, text)
    
    # Extract the translated text using the "Translation" key
    if "Translation" in result:
        return result["Translation"]
    elif "error" in result:
        return f"An error occurred: {result['error']}"
    else:
        return "Translation failed to produce a valid output."


with gr.Blocks(title="Ɔbaa Panin Maternal Health MT") as demo:
    gr.Markdown("# Ɔbaa Panin Maternal Health Machine Translation (MT)")
    gr.Markdown("Translate text between English and Akan, designed for maternal health care")
    with gr.Row():
        with gr.Column():
            model_dd = gr.Dropdown(
                ["Ɔbaa Panin MT: Akan-English", "Ɔbaa Panin MT: English-Akan"],
                label="Model",
                value="Ɔbaa Panin MT: Akan-English",
            )
            input_tb = gr.Textbox(lines=5, label="Input text")
            submit_btn = gr.Button("Translate", variant="huggingface")
        with gr.Column():
            output_tb = gr.Textbox(lines=5, label="Translation Output")

    submit_btn.click(
        fn=text_translation,
        inputs=[model_dd, input_tb],
        outputs=output_tb
    )

    gr.Markdown("Powered by [Human Computer Interaction Lab](https://hci-lab-dcsug.github.io/hci-ugdcs) (Department of Computer Science) at the University of Ghana")

demo.launch()