import gradio as gio
from inference import Translator


class TranslationInterface:
    def __init__(self, languages, create_response) -> None:
        self.interface = gio.Interface(
            create_response,
            [
                gio.Dropdown(
                    languages,
                    label="Input Language",
                    info="The language to translate from",
                    value=languages[0]
                ),
                gio.Dropdown(
                    languages,
                    label="Output Language",
                    info="The language to translate to",
                    value=languages[-1]
                ),
                gio.Textbox(
                    label="Input text",
                    info="The text to be translated",
                    lines=2,
                    value="Momme wohnt in Niebüll."
                )
            ],
            "text"
        )
    
    def launch(self):
        self.interface.launch()

if __name__ == "__main__":
    translator = Translator.from_pretrained('CmdCody/nllb-deu-moo')
    
    def generate_translation(src_lang, tgt_lang, message):
        result = translator.translate(message, src_lang=src_lang, tgt_lang=tgt_lang)
        return result[0]
    
    interface = TranslationInterface(
        ["deu_Latn", "eng_Latn", "dan_Latn", "moo_Latn"],
        generate_translation
    )
    interface.launch()