import gradio as gio
from inference import Translator


LANGUAGES = {
    "German": "deu_Latn",
    "English": "eng_Latn",
    "Danish": "dan_Latn",
    "North Frisian (Mooring)": "moo_Latn"
}

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=LANGUAGES[src_lang],
            tgt_lang=LANGUAGES[tgt_lang])
        return result[0]
    
    interface = TranslationInterface(
        list(LANGUAGES.keys()),
        generate_translation
    )
    interface.launch()