"""
translation program for simple text
1. detect language from langdetect
2. translate to target language given by user

Example from
https://www.thepythoncode.com/article/machine-translation-using-huggingface-transformers-in-python 

user_input:
    string: string to be translated
    target_lang: language to be translated to

Returns:
    string: translated string of text
"""

import gradio as gr
import argparse

import langid
from transformers import pipeline


def lang_detect(article, target_lang):
    """
    Language Detection using library langid

    Args:
        article (string): article that user wish to translate
        target_lang (string): language user want to translate article into

    Returns:
        string: detected language short form
    """

    result_lang = langid.classify(article)
    print(result_lang[0])
    if result_lang == target_lang:
        return result_lang[0]
    else:
        return result_lang[0]


def opus_trans(message, result_lang, target_lang):
    """
    Translation by Helsinki-NLP model

    Args:
        article (string): article that user wishes to translate
        result_lang (string): detected language in short form
        target_lang (string): language that user wishes to translate article into

    Returns:
        string: translated piece of article based off target_lang
    """

    task_name = f"translation_{result_lang}_to_{target_lang}"
    model_name = f"Helsinki-NLP/opus-mt-{result_lang}-{target_lang}"
    translator = pipeline(task_name, model=model_name, tokenizer=model_name)
    translated = translator(message)[0]["translation_text"]
    print(translated)
    return translated


# textbox = gr.Textbox()

translate = gr.Interface(
    opus_trans,
    [
        gr.Textbox(),
        lang_detect([gr.Textbox(), gr.Dropdown(["en", "zh"])]),
        gr.Dropdown(["en", "zh"]),
    ],
    outputs=gr.Textbox(),
)

translate.launch()