import numpy as np
from PIL import Image
import gradio as gr

from utils.langs import languages
from main import predict

language_choices = [(name.title(), code) for name, code in languages.items()]


def process_image(image, target_lang):
    if image is not None:
        if not isinstance(image, np.ndarray):
            image = np.array(Image.open(image))

        translated_image = predict(image, target_lang=target_lang)
        return translated_image
    return None


with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div style="display: flex; align-items: center; flex-direction: row; justify-content: center; margin-bottom: 20px; text-align: center;">
            <a href="https://github.com/Detopall/manga-translator" target="_blank" rel="noopener noreferrer" style="text-decoration: none;">
                <h1 style="display: inline; margin-left: 10px; text-decoration: underline;">Manga Translator</h1>
            </a>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image()
            language_dropdown = gr.Dropdown(
                choices=language_choices,
                label="Target Language",
                value="en-GB",
            )
            submit_button = gr.Button("Translate")
        with gr.Column(scale=1):
            image_output = gr.Image()

    submit_button.click(
        process_image, inputs=[image_input, language_dropdown], outputs=image_output
    )

    examples = gr.Examples(
        examples=[
            ["./examples/ex1.jpg"],
            ["./examples/ex2.jpg"],
            ["./examples/ex3.jpg"],
            ["./examples/ex4.jpg"],
        ],
        inputs=image_input,
    )

if __name__ == "__main__":
    demo.launch()