import gradio as gr
import numpy as np

import requests
from PIL import Image
from io import BytesIO
import time

# TODO: PNGinfo追加, 使用データの展開(コピーボタン)
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1344

model_url = {
    "ImageUltra": "https://api.stability.ai/v2beta/stable-image/generate/ultra",
    "ImageCore": "https://api.stability.ai/v2beta/stable-image/generate/core",
    "StableDiffusion3": "https://api.stability.ai/v2beta/stable-image/generate/sd3",
}

service_url = {
    "Conservative_Upscale": "https://api.stability.ai/v2beta/stable-image/upscale/conservative",
    "Creative_Upscale": "https://api.stability.ai/v2beta/stable-image/upscale/creative",
    "Erase": "https://api.stability.ai/v2beta/stable-image/edit/erase",
    "Inpaint": "https://api.stability.ai/v2beta/stable-image/edit/inpaint",
    "Outpaint": "https://api.stability.ai/v2beta/stable-image/edit/outpaint",
    "SR": "https://api.stability.ai/v2beta/stable-image/edit/search-and-replace",
    "RMBG": "https://api.stability.ai/v2beta/stable-image/edit/remove-background",
    "Sketch": "https://api.stability.ai/v2beta/stable-image/control/sketch",
    "Structure": "https://api.stability.ai/v2beta/stable-image/control/structure",
}


def bytes_to_image(image):
    image = BytesIO(image)
    image = Image.open(image).convert("RGB")
    return image


def image_to_bytes(image):
    byte_io = BytesIO()
    image.save(byte_io, format="PNG")
    byte_data = byte_io.getvalue()
    return byte_data


def send_request(url, api_key, file, data):
    response = requests.post(
        url,
        headers={"Authorization": f"Bearer {api_key}", "Accept": "image/*"},
        files=file,
        data=data,
    )
    return response


def generate(
    prompt,
    negative_prompt,
    seed,
    mode,
    submode,
    input_image,
    mask,
    CNstrength,
    search_prompt,
    op_left,
    op_right,
    op_up,
    op_down,
    randomize_seed,
    aspect,
    model,
    preset,
    api_key,
):
    if randomize_seed:
        seed = 0

    file = {}
    data = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "output_format": "png",
        "seed": seed,
        "aspect_ratio": aspect,
    }
    if input_image is not None:
        file["image"] = image_to_bytes(input_image)
        if mask is not None:
            file["mask"] = image_to_bytes(mask)

    if mode == "Generate":
        file["none"] = ""
        if model == "Stable Image Ultra (8B + workflow)":
            url = model_url["ImageUltra"]
        elif model == "Stable Image Core (2B + workflow)":
            url = model_url["ImageCore"]
            data["style_preset"] = preset
        elif model == "Stable Diffusion 3 Medium (2B)":
            url = model_url["StableDiffusion3"]
            data["model"] = "sd3-medium"
        elif model == "Stable Diffusion 3 Large (8B)":
            url = model_url["StableDiffusion3"]
            data["model"] = "sd3-large"
        elif model == "Stable Diffusion 3 Large Turbo (8B Turbo)":
            url = model_url["StableDiffusion3"]
            data["model"] = "sd3-large-turbo"

        else:
            raise ValueError("Invalid model type")

    elif mode == "Upscale":
        if submode == "Conservative":
            url = service_url["Conservative_Upscale"]
        elif submode == "Creative":
            url = service_url["Creative_Upscale"]

    elif mode == "Edit":
        if submode == "Erase":
            url = service_url["Erase"]
        elif submode == "Inpaint":
            url = service_url["Inpaint"]
        elif submode == "Outpaint":
            url = service_url["Outpaint"]
            data["left"] = op_left
            data["right"] = op_right
            data["up"] = op_up
            data["down"] = op_down
        elif submode == "Search and Replace":
            url = service_url["SR"]
            data["search_prompt"] = search_prompt
        elif submode == "Remove Background":
            url = service_url["RMBG"]

    elif mode == "Control":
        data["control_strength"] = CNstrength
        if submode == "Sketch":
            url = service_url["Sketch"]
        elif submode == "Structure":
            url = service_url["Structure"]

    response = send_request(url, api_key, file, data)

    if response.status_code == 200:
        if mode == "Upscale" and submode == "Creative":
            generation_id = response.json().get("id")
            if not generation_id:
                raise Exception("No generation ID returned for creative upscale")

            # Polling for the result
            result_url = f"https://api.stability.ai/v2beta/stable-image/upscale/creative/result/{generation_id}"
            while True:
                result_response = requests.get(
                    result_url,
                    headers={"accept": "image/*", "authorization": f"Bearer {api_key}"},
                )
                if result_response.status_code == 202:
                    print("Generation in-progress, try again in 10 seconds.")
                    time.sleep(10)
                elif result_response.status_code == 200:
                    print("Generation complete!")
                    image = result_response.content
                    image = bytes_to_image(image)
                    copy_filed_value = f"prompt:{prompt}, negative:{negative_prompt}, mode:{mode}, submode:{submode}"
                    return image, seed, copy_filed_value
                else:
                    raise Exception(str(result_response.json()))
        else:
            image = response.content
            image = bytes_to_image(image)
            copy_filed_value = f"prompt:{prompt}, negative:{negative_prompt}, mode:{mode}, submode:{submode}"
            return image, seed, copy_filed_value
    else:
        raise Exception(str(response.json()))


examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 50vw;
}
"""


def update_style_visibility(model):
    if model == "Stable Image Core (2B + workflow)":
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)


def update_mode(mode):
    submode_update = gr.update(choices=["None"], visible=False)
    image_label_update = gr.update(visible=False)
    img_input_update = gr.update(visible=False)
    mask_update = gr.update(visible=False)

    if mode == "Generate":
        submode_update = gr.update(visible=False)
    elif mode == "Upscale":
        submode_update = gr.update(
            choices=["Conservative", "Creative"], value="Conservative", visible=True
        )
        img_input_update = gr.update(visible=True)
        image_label_update = gr.update(visible=True)
    elif mode == "Edit":
        submode_update = gr.update(
            choices=[
                "Erase",
                "Inpaint",
                "Outpaint",
                "Search and Replace",
                "Remove Background",
            ],
            value="Erase",
            visible=True,
        )
        img_input_update = gr.update(visible=True)
        image_label_update = gr.update(visible=True)
    elif mode == "Control":
        submode_update = gr.update(
            choices=["Sketch", "Structure"], value="Sketch", visible=True
        )
        img_input_update = gr.update(visible=True)
        image_label_update = gr.update(visible=True)

    return submode_update, img_input_update, mask_update, image_label_update


def update_submode(submode):
    mask = gr.update(visible=False)
    outpaint = gr.update(visible=False)
    cn = gr.update(visible=False)
    search_prompt = gr.update(visible=False)

    if submode in ["Erase", "Inpaint"]:
        mask = gr.update(visible=True)

    else:
        if submode == "Outpaint":
            outpaint = gr.update(visible=True)

        elif submode == "Control":
            cn = gr.update(visible=True)

        elif submode == "Search and Replace":
            search_prompt = gr.update(visible=True)

    return mask, outpaint, cn, search_prompt


with gr.Blocks(css=css, theme="NoCrypt/miku") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
        f"""
        # Stability AI - Developer Platform WebUI
        **Overview**  
        Stability AI’s Stable Image services offer a growing set of APIs for developers to build the best in class image applications.

        - **Disrupting Content Creation:** Stability’s Image APIs are the foundation for applications disrupting publishing, media, gaming, marketing, advertising, design, and more.
        - **For Developers:** Application developers can build advanced features for designers, photographers, content creators, and a variety of B2C customers.
        - **Simple APIs:** Stability AI is focused on delivering simple APIs for easy integration into applications with a high bar for quality, alignment, speed, and safety.

        Get Started Now: https://platform.stability.ai/docs/getting-started/stable-image    

        Contact: D̷ELL@Stability AI - Advocate (https://x.com/xqdior) / Author: umise (https://x.com/UiE029)
        """
        )

        with gr.Row():
            api_key = gr.Text(
                label="API Key",
                type="password",
                placeholder="Enter your API key",
                max_lines=1,
                container=False,
            )

        with gr.Row():
            model = gr.Dropdown(
                label="Model",
                choices=["Stable Image Ultra (8B + workflow)", "Stable Image Core (2B + workflow)", "Stable Diffusion 3 Large Turbo (8B Turbo)", "Stable Diffusion 3 Large (8B)", "Stable Diffusion 3 Medium (2B)"],
                value="Ultra",
            )
            mode = gr.Dropdown(
                label="Mode",
                choices=["Generate", "Upscale", "Edit", "Control"],
                value="Generate",
            )

        submode = gr.Dropdown(
            label="Submode", choices=["None"], visible=False, value="None"
        )

        with gr.Row():
            with gr.Column():
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
                search_prompt = gr.Text(
                    label="search prompt",
                    visible=False,
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter a search prompt",
                )

            run_button = gr.Button("Run", scale=0)

        with gr.Row():
            with gr.Column():
                image_label = gr.Markdown(value="input image", visible=False)
                image = gr.Image(
                    type="pil",
                    label="img input",
                    width="20vw",
                    height="20vw",
                    show_label=True,
                    visible=False,
                    interactive=True,
                    container=False,
                )
            with gr.Column(visible=False) as mask:
                mask_label = gr.Markdown(value="input mask")
                mask_input = gr.Image(
                    type="pil",
                    label="mask",
                    width="20vw",
                    height="20vw",
                    show_label=True,
                    interactive=True,
                    container=False,
                )

        with gr.Row():
            result = gr.Image(label="Result", width="20vw", height="20%")

        with gr.Accordion("Advanced Settings", open=False):
            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
            )
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            CN_strength = gr.Slider(
                label="Control Strength",
                minimum=0,
                maximum=1,
                step=0.01,
                value=0.5,
                visible=False,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            aspect = gr.Radio(
                choices=[
                    "1:1",
                    "16:9",
                    "21:9",
                    "2:3",
                    "3:2",
                    "4:5",
                    "5:4",
                    "9:16",
                    "9:21",
                ],
                label="Aspect raito",
                value="1:1",
            )
            with gr.Row(visible=False) as style:
                style_preset = gr.Radio(
                    choices=[
                        "3d-model",
                        "analog-film",
                        "anime",
                        "cinematic",
                        "comic-book",
                        "digital-art",
                        "enhance",
                        "fantasy-art",
                        "isometric",
                        "line-art",
                        "low-poly",
                        "modeling-compound",
                        "neon-punk",
                        "origami",
                        "photographic",
                        "pixel-art",
                        "tile-texture",
                    ],
                    label="Style_preset",
                    value="anime",
                    info="This parameter is only available for ImageCore model.",
                )
            with gr.Row(visible=False) as outpaint_scale:
                paint = gr.Markdown(value="Outpain Scale")
                op_left = gr.Slider(
                    label="left", minimum=0, maximum=2000, step=4, value=200
                )
                op_right = gr.Slider(
                    label="right", minimum=0, maximum=2000, step=4, value=200
                )
                op_up = gr.Slider(
                    label="up", minimum=0, maximum=2000, step=4, value=200
                )
                op_down = gr.Slider(
                    label="down", minimum=0, maximum=2000, step=4, value=200
                )

        gr.Examples(examples=examples, inputs=[prompt])

        copy_filed = gr.TextArea(
            value="",
            label="Copy Field",
            max_lines=1,
            placeholder="Copy the field",
            show_copy_button=True,
            container=False,
        )
        gr.Markdown(
        f"""
        ## License
        This work is licensed under a
        [Creative Commons Attribution-NonCommercial 4.0 International License][cc-by-nc].

        [![CC BY-NC 4.0][cc-by-nc-image]][cc-by-nc]

        [cc-by-nc]: https://creativecommons.org/licenses/by-nc/4.0/
        [cc-by-nc-image]: https://licensebuttons.net/l/by-nc/4.0/88x31.png
        [cc-by-nc-shield]: https://img.shields.io/badge/License-CC%20BY--NC%204.0-lightgrey.svg

        **MIT Licensed Source Code**
        Portions of this work are licensed under the MIT License. For more details, please refer to the original source at: [stabilityai/stable-diffusion-3-medium](https://huggingface.co/spaces/stabilityai/stable-diffusion-3-medium) 
        """
        )
    gr.on(
        triggers=[run_button.click, prompt.submit, negative_prompt.submit],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            mode,
            submode,
            image,
            mask_input,
            CN_strength,
            search_prompt,
            op_left,
            op_right,
            op_up,
            op_down,
            randomize_seed,
            aspect,
            model,
            style_preset,
            api_key,
        ],
        outputs=[result, seed, copy_filed],
    )

    mode.change(
        fn=update_mode, inputs=mode, outputs=[submode, image, mask, image_label]
    )
    submode.change(
        fn=update_submode,
        inputs=submode,
        outputs=[mask, outpaint_scale, CN_strength, search_prompt],
    )
    model.change(fn=update_style_visibility, inputs=model, outputs=style)
demo.launch()