from glob import glob

import gradio as gr
from gradio_client import Client

from utils import make_flatten_background

REPO_ID = "leonelhs/faceshine"

clients = {
    "GFPGAN": "leonelhs/GFPGAN",
    "ZeroScratches": "leonelhs/ZeroScratches",
    "Deoldify": "leonelhs/deoldify",
    "EnhanceLight": "leonelhs/Zero-DCE",
    "ZeroBackground": "leonelhs/rembg",
}


def load_client(space):
    try:
        return Client(space)
    except ValueError as err:
        print(err)
        logger.value.append(f"Space: {space}, log: {err}")
        pass


def gfpgan_face(image, version, scale):
    return clients["GFPGAN"].predict(image, version, scale, fn_index=0)[0]


def zero_scratches(image):
    return clients["ZeroScratches"].predict(image, api_name="/predict")


def colorize_photo(image):
    return clients["Deoldify"].predict(image, api_name="/predict")


def enhance_light(image):
    return clients["EnhanceLight"].predict(image, api_name="/predict")


def zero_background(image, new_bgr=None):
    # Fixme: cant find predict function by name
    # return clients["ZeroBackground"].predict(image, new_bgr, fn_index=0)[1]
    # return clients["ZeroBackground"].predict(image, fn_index=0)
    img, mask = clients["ZeroBackground"].predict(image, "U2NET Human Seg", False, fn_index=9)
    return make_flatten_background(img, mask)


def parse_face(image):
    return clients["FaceParser"].predict(image, api_name="/predict")


def mirror(x):
    return x


def active_first():
    return gr.Tabs.update(selected=0)


def clear():
    return None, None


footer = r"""
<center>
<p>This App is running on a CPU, help us to upgrade a GPU or just give us a <a href='https://github.com/leonelhs/face-shine' target='_blank'>Github ⭐</a></p>
</br>
<a href="https://www.buymeacoffee.com/leonelhs">
    <img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=leonelhs&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" />
</a>
</center>
</br>
<center><span>leonelhs@gmail.com</span></center>
"""

with gr.Blocks(title="Face Shine") as app:
    logger = gr.State(value=[])

    for client, endpoint in clients.items():
        clients[client] = load_client(endpoint)

    with gr.Row():
        gr.HTML("<center><h1>Face Shine</h1></center>")

    with gr.Tabs() as tabs:
        with gr.TabItem("Photo restorer", id=0):
            with gr.Row(equal_height=False):
                with gr.Column(scale=1):
                    btn_eraser = gr.Button(value="Erase scratches")
                    btn_color = gr.Button(value="Colorize photo")
                    btn_hires = gr.Button(value="Enhance face")
                    btn_light = gr.Button(value="Enhance light")
                    btn_clear = gr.Button(value="Flatten background")

                with gr.Column(scale=2):
                    with gr.Row():
                        img_input = gr.Image(label="Input", type="filepath")
                    with gr.Row():
                        btn_reset = gr.Button(value="Reset", variant="stop")
                        btn_swap = gr.Button(value="Ok", variant="primary")

                with gr.Column(scale=2):
                    with gr.Row():
                        img_output = gr.Image(label="Result", type="filepath", interactive=False)

        with gr.TabItem("Examples", id=1):
            gr.Examples(examples=glob("lowres/*"), inputs=[img_input], label="Low resolution")
            gr.Examples(examples=glob("gray/*"), inputs=[img_input], label="Gray scale")
            gr.Examples(examples=glob("scratch/*"), inputs=[img_input], label="Scratched")
            gr.Button(value="Ok", variant="primary").click(active_first, None, tabs)

        with gr.TabItem("Settings", id=2):
            with gr.Accordion("Image restoration settings", open=False):
                enhancer = gr.Dropdown(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer'],
                                       label='GFPGAN face restoration algorithm',
                                       type="value", value='RestoreFormer',
                                       info="version")
                rescale = gr.Dropdown(["1", "2", "3", "4"],
                                      type="value", value="2", label="Rescaling factor")
            with gr.Accordion("Logs info", open=False):
                text_logger = gr.Textbox(label="login", lines=5, show_label=False)
            gr.Button("Save settings")

    btn_hires.click(gfpgan_face, inputs=[img_input, enhancer, rescale], outputs=[img_output])
    btn_eraser.click(zero_scratches, inputs=[img_input], outputs=[img_output])
    btn_color.click(colorize_photo, inputs=[img_input], outputs=[img_output])
    btn_light.click(enhance_light, inputs=[img_input], outputs=[img_output])
    btn_clear.click(zero_background, inputs=[img_input], outputs=[img_output])
    btn_swap.click(mirror, inputs=[img_output], outputs=[img_input])
    btn_reset.click(clear, outputs=[img_input, img_output])

    with gr.Row():
        gr.HTML(footer)

app.launch(share=False, debug=True, show_error=True)
app.queue()