import requests
from PIL import Image
import os
import io
import base64
from pathlib import Path
import gradio as gr

# parameters which can be customized in settings.json of webui

initial_width = 512
initial_height = 512

params = {
    "enable_SD_api": True,
    "address": "https://dd31-2001-d08-e3-56b9-d182-69e5-40c4-e46d.ngrok-free.app",
    "save_img": True,
    "SD_model": "Chilloutmix-Ni-prune-fp32-fix.safetensors",
    "prompt_prefix": "detailed portrait of a asian man in black jacket, Futuristic sci-fi fashion, fine details, realistic shaded, fine-face, pretty face cyberpunk, neotokyo, synthwave, aesthetics, futuristic, low-emission-neon, (blade runner movie scene)",
    "negative_prompt": "jpeg artifacts, low quality, lowres, 3d, render, doll, plastic, blur, haze, monochrome, b&w, text, (ugly:1.2), unclear eyes, no arms, bad anatomy, cropped, censoring, asymmetric eyes, bad anatomy, bad proportions, cropped, cross-eyed, deformed, extra arms, extra fingers, extra limbs, fused fingers, malformed, mangled hands, misshapen body, missing arms, missing fingers, missing hands, missing legs, poorly drawn, tentacle finger, too many arms, too many fingers, watermark, logo, text, letters, signature, username, words, blurry, cropped",
    "width": initial_width,
    "height": initial_height,
    "restore_faces": False,
}

pic_id = 0


# Get and save the Stable Diffusion-generated picture
def get_SD_pictures(description):
    global params, pic_id
    payload = {
        "prompt": params["prompt_prefix"] + description,
        "seed": -1,
        "sampler_name": "Euler a",
        "steps": 20,
        "cfg_scale": 7,
        "width": params["width"],
        "height": params["height"],
        "restore_faces": params["restore_faces"],
        "negative_prompt": params["negative_prompt"],
    }

    try:
        response = requests.post(
            url=f'{params["address"]}/sdapi/v1/txt2img', json=payload, timeout=100
        )
        response.raise_for_status()  # Raises stored HTTPError, if one occurred
    except requests.exceptions.HTTPError as http_err:
        return f"HTTP error occurred: {http_err}"
    except Exception as err:
        return f"An error occurred: {err}"

    r = response.json()

    visible_result = ""
    for img_str in r["images"]:
        image = Image.open(io.BytesIO(base64.b64decode(img_str.split(",", 1)[0])))

        if not os.path.exists("outputs"):
            os.makedirs("outputs")

        if params["save_img"]:
            output_file = Path(f"outputs/{pic_id:06d}.png")
            image.save(output_file.as_posix())
            pic_id += 1

        # lower the resolution of received images for the chat
        # otherwise the log size gets out of control quickly
        # with all the base64 values in visible history
        image.thumbnail((512, 512))
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        buffered.seek(0)
        image_bytes = buffered.getvalue()
        img_str = "data:image/jpeg;base64," + base64.b64encode(image_bytes).decode()
        visible_result = visible_result + f'<img src="{img_str}" alt="{description}">\n'

    return visible_result


def display_image(description: str, artist: str) -> gr.Image:
    if artist == "Kashif":
        description += " <lora:Kashif_v1:0.9>"
    elif artist == "Daniel Ho":
        description += " <lora:DanielHo_v1:0.9>"
    elif artist == "Adrian":
        description += " <lora:AdrianSeowV1:0.9>"
    elif artist == "Tony":
        description += " <lora:TonyYap4:0.9>"
    elif artist == "Kevin":
        description += " <lora:KevinLowV1:0.9>"
    elif artist == "Aaron":
        description += " <lora:AaronSiowV1:0.9>"

    if artist == "none":
        artist = False

    if artist:
        params["width"] = 768
        params["height"] = 768
        params[
            "prompt_prefix"
        ] = "detailed portrait of a asian man in black jacket, Futuristic sci-fi fashion, fine details, realistic shaded, fine-face, pretty face cyberpunk, neotokyo, synthwave, aesthetics, futuristic, low-emission-neon, (blade runner movie scene)"
    else:
        params["width"] = initial_width
        params["height"] = initial_height
        params["prompt_prefix"] = "4k"

    visible_result = get_SD_pictures(description)
    if "error occurred" in visible_result:
        return visible_result

    image_data = base64.b64decode(visible_result.split(",", 1)[1])
    image = Image.open(io.BytesIO(image_data))
    return image


inputs = [
    gr.inputs.Textbox(lines=2, label="Trial period has ended. Follow daveckw instagram for future updates"),
    gr.inputs.Radio(
        ["Kashif", "Daniel Ho", "Adrian", "Tony", "Kevin", "Aaron", "none"],
        label="Artist",
    ),
]

outputs = gr.outputs.Image(type="pil", label="Generated Image")

interface = gr.Interface(
    fn=display_image, inputs=inputs, outputs=outputs, title="DC Image Generator"
)

interface.launch()