"""
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
"""
from __future__ import annotations
from typing import Iterable
import gradio as gr
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
import argparse

import langid
from transformers import pipeline


class myTheme(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.red,
        secondary_hue: colors.Color | str = colors.blue,
        neutral_hue: colors.Color | str = colors.orange,
        spacing_size: sizes.Size | str = sizes.spacing_md,
        radius_size: sizes.Size | str = sizes.radius_md,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("handjet"),
            "cursive",
            # "sans-serif",
        ),
        font_mono: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"),
            "ui-monospace",
            "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            spacing_size=spacing_size,
            radius_size=radius_size,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )
        super().set(
            body_background_fill="repeating-linear-gradient(135deg, *primary_200, *primary_200 10px, *primary_50 10px, *primary_50 20px)",
            body_background_fill_dark="repeating-linear-gradient(135deg, *primary_800, *primary_800 10px, *primary_900 10px, *primary_900 20px)",
            button_primary_background_fill="linear-gradient(45deg, *primary_300, *secondary_400)",
            button_primary_background_fill_hover="linear-gradient(45deg, *primary_200, *secondary_300)",
            button_primary_text_color="white",
            button_primary_background_fill_dark="linear-gradient(90deg, *primary_600, *secondary_800)",
            slider_color="*secondary_300",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
            button_shadow="*shadow_drop_lg",
            button_large_padding="24px",
        )


def detect_lang(article):
    """
    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)
    return result_lang[0]


def opus_trans(article, target_language):
    """
    Translation by Helsinki-NLP model

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

    Returns:
        string: translated piece of article based off target_language
    """
    print(target_language)
    result_lang = detect_lang(article, target_language)

    if target_language == "English":
        target_lang = "en"
    elif target_language == "Chinese":
        target_lang = "zh"
    # result_lang = detect_lang(article, target_language)
    print(target_lang)
    if result_lang != 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(article)[0]["translation_text"]
        # print(translated)
    else:
        translated = "Error: You chose the same language as the article detected language. Please reselect target_language and try again."
    return translated


myTheme = myTheme()

with gr.Blocks(theme=myTheme) as demo:
    article = gr.Textbox(label="Article")
    lang_select = gr.Radio(["English", "Chinese"], label="Select Desired Language")
    result = gr.Textbox(label="Translated Result")
    trans_btn = gr.Button("Translate")
    trans_btn.click(fn=opus_trans, inputs=[article, lang_select], outputs=result)

demo.launch()