import os
import numpy as np
from prodiapy import Prodia
import gradio as gr
from style_template import styles
import base64

prodia = Prodia()
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "Photographic (Default)"
MAX_SEED = np.iinfo(np.int32).max


def generate_image(upload_images, prompt, negative_prompt, style_name, steps, strength, seed, progress=gr.Progress(track_tqdm=True)):
    error_if_no_img(prompt)
    p, n = apply_style(style_name, prompt, negative_prompt)

    job = prodia.create("/photomaker",
                        imageData=[file_to_base64(img) for img in upload_images],
                        prompt=p,
                        negative_prompt=n,
                        steps=steps,
                        strength=strength,
                        seed=seed if seed != 0 else None
                        )
    result = prodia.wait(job)

    return result.image_url


def error_if_no_img(prompt):
    if "img" not in prompt:
        raise gr.Error("Prompt must contain 'img'")



def swap_to_gallery(images):
    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)


def upload_example_to_gallery(images, prompt, style, negative_prompt):
    return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)


def remove_back_to_files():
    return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)


def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n + ' ' + negative


def get_image_path_list(folder_name):
    image_basename_list = os.listdir(folder_name)
    image_path_list = sorted([os.path.join(folder_name, basename) for basename in image_basename_list])
    return image_path_list


def file_to_base64(file_path):
    with open(file_path, "rb") as file:
        file_data = file.read()
        base64_string = base64.b64encode(file_data).decode('utf-8')
        return base64_string


def get_example():
    case = [
        [
            get_image_path_list('./examples/scarletthead_woman'),
            "instagram photo, portrait photo of a woman img, colorful, perfect face, natural skin, hard shadows, film grain",
            "(No style)",
            "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",
        ],
        [
            get_image_path_list('./examples/newton_man'),
            "sci-fi, closeup portrait photo of a man img wearing the sunglasses in Iron man suit, face, slim body, high quality, film grain",
            "(No style)",
            "(asymmetry, worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch), open mouth",
        ],
    ]
    return case


title = r"""
<h1 align="center">PhotoMaker: Customizing Realistic Human Photos via Stacked ID Embedding</h1>
"""

css = '''
.gradio-container {width: 85% !important}
'''
with gr.Blocks(css=css) as demo:
    gr.Markdown(title)

    with gr.Row():
        with gr.Column():
            files = gr.File(
                label="Drag (Select) 1 or more photos of your face",
                file_types=["image"],
                file_count="multiple"
            )
            uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200)
            with gr.Column(visible=False) as clear_button:
                remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
            prompt = gr.Textbox(label="Prompt",
                                info="Try something like 'a photo of a man/woman img', 'img' is the trigger word.",
                                placeholder="A photo of a [man/woman img]...")
            style = gr.Dropdown(label="Style template", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
            submit = gr.Button("Submit")

            with gr.Accordion(open=False, label="Advanced Options"):
                negative_prompt = gr.Textbox(
                    label="Negative Prompt", 
                    placeholder="low quality",
                    value="nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry",
                )
                steps = gr.Slider(
                    label="Number of sample steps",
                    minimum=20,
                    maximum=50,
                    step=1,
                    value=40,
                )
                strength_ratio = gr.Slider(
                    label="Strength (%)",
                    minimum=15,
                    maximum=50,
                    step=1,
                    value=20,
                )
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
        with gr.Column():
            result_image = gr.Image(label="Generated Image")

        files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
        remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])

        submit.click(
            fn=generate_image,
            inputs=[files, prompt, negative_prompt, style, steps, strength_ratio, seed],
            outputs=[result_image]
        )

    gr.Examples(
        examples=get_example(),
        inputs=[files, prompt, style, negative_prompt],
        run_on_click=True,
        fn=upload_example_to_gallery,
        outputs=[uploaded_files, clear_button, files],
    )
    
demo.queue(max_size=20).launch(show_api=False)