import zipfile
from pathlib import Path

import gradio as gr
from gradio_client import Client
from huggingface_hub import hf_hub_download

clientGFPGAN = Client("Xintao/GFPGAN")
clientSuperface = Client("leonelhs/superface")
clientZeroScratches = Client("leonelhs/ZeroScratches")
clientDeoldify = Client("leonelhs/deoldify")
clientEnhanceLight = Client("leonelhs/Zero-DCE")
clientZeroBackground = Client("leonelhs/remove-background")
clientFaceParser = Client("leonelhs/faceparser")

REPO_ID = "leonelhs/faceshine"


def download(file):
    file_zip = hf_hub_download(repo_id=REPO_ID, filename=file + ".zip")
    with zipfile.ZipFile(file_zip, "r") as zip_ref:
        zip_ref.extractall(file)
    return list(Path(file).iterdir())


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


def enhance_face(image, upsampler, face_enhancer, scale):
    return clientSuperface.predict(image, upsampler, face_enhancer, scale, api_name="/predict")


def zero_scratches(image):
    return clientZeroScratches.predict(image, api_name="/predict")


def colorize_photo(image):
    return clientDeoldify.predict(image, api_name="/predict")


def enhance_light(image):
    return clientEnhanceLight.predict(image, api_name="/predict")


def zero_background(image, new_bgr=None):
    # return clientZeroBackground.predict(image, new_bgr, api_name="/predict")[0]
    # Fixme: cant find predict function by name
    return clientZeroBackground.predict(image, new_bgr, fn_index=0)[1]


def parse_face(image):
    return clientFaceParser.predict(image, api_name="/predict")


def mirror(x):
    return x


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


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:
    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():
                with gr.Column(scale=1):
                    btn_hires = gr.Button(value="Enhance faces")
                    btn_eraser = gr.Button(value="Erase scratches")
                    btn_color = gr.Button(value="Colorize photo")
                    btn_light = gr.Button(value="Enhance light")
                    with gr.Accordion("New background", open=False):
                        img_newbgr = gr.Image(label="Left empty for a transparent background", type="filepath")
                        btn_newbgr = gr.Button(value="Clear background")

                with gr.Column(scale=4):
                    with gr.Row():
                        img_input = gr.Image(label="Input", type="filepath")
                        img_output = gr.Image(label="Result", type="filepath", interactive=False)
                    with gr.Row():
                        gr.ClearButton([img_input, img_output, img_newbgr], variant="stop")
                        btn_swap = gr.Button(value="Swap images", variant="primary")

        with gr.TabItem("Examples", id=1):
            gr.Examples(examples=download("lowres"), inputs=[img_input], label="Low resolution")
            gr.Examples(examples=download("gray"), inputs=[img_input], label="Gray scale")
            gr.Examples(examples=download("scratch"), inputs=[img_input], label="Scratched")
            gr.Examples(examples=download("backs"), inputs=[img_newbgr], label="Backgrounds")
            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'],
                    type="value", value='RestoreFormer', label='GFPGAN face restoration algorithm',
                    info="version")
                restorer = gr.Dropdown([
                    'RealESRGAN_x2plus',
                    'RealESRGAN_x4plus',
                    'RealESRNet_x4plus',
                    'AI-Forever_x2plus',
                    'AI-Forever_x4plus',
                    'RealESRGAN_x4plus_anime_6B',
                    'realesr-animevideov3',
                    'realesr-general-x4v3'],
                    type="value", value='RealESRGAN_x4plus',
                    label='Real-ESRGAN image restoration algorithm', info="version", visible=False)
                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(enhance_face, inputs=[img_input, restorer, enhancer, rescale], outputs=[img_output, text_logger])
    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_newbgr.click(zero_background, inputs=[img_input, img_newbgr], outputs=[img_output])
    btn_swap.click(mirror, inputs=[img_output], outputs=[img_input])

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

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