import importlib
import os
from PIL import Image
import gradio as gr
from omegaconf import OmegaConf

config = OmegaConf.load("config/annotator.yaml")

package_annotator = "processor"


def process_image(cls: str, fg: Image.Image, bg: Image.Image, *kwargs):
    if fg.size != bg.size:
        fg = fg.resize(bg.size)
    module_imp = importlib.import_module(package_annotator)
    model = getattr(module_imp, cls)
    image_processor = model()
    result = image_processor(fg, bg, *kwargs)
    if type(result) == tuple:
        return result
    return [result]


def process(cls):
    def process_fc(img, res, *args):
        return process_image(cls, img, res, *args)

    return process_fc


block = gr.Blocks().queue()
examples = [[os.path.join(os.path.dirname(__file__), "example/product.png"),
             os.path.join(os.path.dirname(__file__), "example/bg.png")]]
with block:
    for key in config.keys():
        cls, input_element = config[key]["process"], config[key].get("input")
        input_append = []
        with gr.Tab(key):
            with gr.Row():
                gr.Markdown("## " + key)
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(label="foreground", source='upload', type="pil", image_mode="RGBA")
                    bg_image = gr.Image(label="background", source='upload', type="pil", image_mode="RGBA")
                    if input_element is not None:
                        for item in input_element:
                            input_append.append(getattr(gr, item["attr"])(**item["args"]))
                    run_button = gr.Button(label="Run")
                    gr.Examples(examples, [input_image, bg_image])
                with gr.Column():
                    gallery = gr.Gallery(label="Generated images", show_label=False).style(height="auto")

            run_button.click(fn=process(cls),
                             inputs=[input_image, bg_image] + input_append,
                             outputs=[gallery])

block.launch(server_port=7867)