import gradio as gr
from transformers import pipeline

import ast
import os

import utils

# Get pipeline from Hugging Face Hub
pipe = pipeline("translation", model="eepj/wstcg-mt-ja-en")
# Model to enter evaluation mode
pipe.model.eval()


def func(text_ja: str, sub_emoji: bool) -> str:
    # Format the input string to replace emoji
    text_ja = utils.format_input(text_ja)
    # Split text by line
    splits_ja = [s for s in text_ja.splitlines() if s]

    # Iterate through each sentence pair
    segs_en = []
    for s in splits_ja:
        # Join a sentence pair
        seg_ja = "".join(s)
        # Replace named entites with placeholder tokens (<NAME> and <TRAIT>)
        seg_ja, repls = utils.text_to_placeholder_tokens(seg_ja)
        # Translate the sentence pair
        seg_en = pipe.predict(seg_ja)[0]["translation_text"]
        # Replace placehoder tokens with the original text
        seg_en = utils.placeholder_tokens_to_text(seg_en, repls)
        # Format segment
        seg_en = utils.format_output(seg_en, sub_emoji)
        # Save the translated segment
        segs_en.append(seg_en)

    # Join the translated segments as one output strings
    text_en = "\n".join(segs_en)

    return text_en


with gr.Blocks() as app:
    gr.Markdown("# WS TCG Card Text Translator")

    with gr.Row():
        with gr.Column():
            input_box = gr.TextArea(label="Original Card Text",
                                    info="Put each ability on a new line")
            output_box = gr.TextArea(label="Translated Card Text")

            sub_emoji = gr.Checkbox(True, label="Show Trigger Icon Emojis", info="Optional")

            submit_btn = gr.Button(variant="primary")
            clear_btn = gr.ClearButton([input_box, output_box])

        try:
            example_text = [[e] for e in ast.literal_eval(os.environ["EXAMPLE_TEXT"])]
            if len(example_text) > 0:
                with gr.Column():
                    gr.Examples(
                        example_text,
                        inputs=[input_box, sub_emoji],
                        fn=func, outputs=[output_box],
                        label="Example Text", api_name=False, cache_examples=False)

        except (KeyError, SyntaxError, ValueError) as err:
            print(err)

    submit_btn.click(func, [input_box, sub_emoji], [output_box])

if __name__ == "__main__":
    app.launch(show_api=False)