import os

import gradio as gr
import rembg
import spaces
import torch
from diffusers import DiffusionPipeline

from instantMesh.src.utils.infer_util import (
    remove_background, resize_foreground)


pipe = DiffusionPipeline.from_pretrained(
    "playgroundai/playground-v2.5-1024px-aesthetic",
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda")


def generate_prompt(subject, style, color_scheme, angle, lighting_type, additional_details):
    return f"A 3D cartoon render of {subject}, featuring the entire body and shape, on a transparent background. The style should be {style}, with {color_scheme} colors, emphasizing the essential features and lines. The pose should clearly showcase the full form of the {subject} from a {angle} perspective. Lighting is {lighting_type}, highlighting the volume and depth of the subject. {additional_details}. Output as a high-resolution PNG with no background."


@spaces.GPU
def generate_image(subject, style, color_scheme, angle, lighting_type, additional_details):
    prompt = generate_prompt(subject, style, color_scheme,
                             angle, lighting_type, additional_details)
    results = pipe(prompt, num_inference_steps=25, guidance_scale=7.5)
    return results.images[0]


def check_input_image(input_image):
    if input_image is None:
        raise gr.Error("No image selected!")


def preprocess(input_image):
    rembg_session = rembg.new_session()

    input_image = remove_background(input_image, rembg_session)
    input_image = resize_foreground(input_image, 0.85)

    return input_image


def image_generation_ui():
    with gr.Row():
        subject = gr.Textbox(label='Subject', scale=2)
        style = gr.Dropdown(
            label='Style',
            choices=['Pixar-like', 'Disney-esque', 'Anime-inspired'],
            value='Pixar-like',
            multiselect=False,
            scale=2
        )
        color_scheme = gr.Dropdown(
            label='Color Scheme',
            choices=['Vibrant', 'Pastel', 'Monochromatic', 'Black and White'],
            value='Vibrant',
            multiselect=False,
            scale=2
        )
        angle = gr.Dropdown(
            label='Angle',
            choices=['Front', 'Side', 'Three-quarter'],
            value='Front',
            multiselect=False,
            scale=2
        )
        lighting_type = gr.Dropdown(
            label='Lighting Type',
            choices=['Bright and Even', 'Dramatic Shadows', 'Soft and Warm'],
            value='Bright and Even',
            multiselect=False,
            scale=2
        )
        additional_details = gr.Textbox(label='Additional Details', scale=2)
        submit_prompt = gr.Button('Generate Image', scale=1, variant='primary')

    with gr.Row(variant="panel"):
        with gr.Column():
            with gr.Row():
                input_image = gr.Image(
                    label="Input Image",
                    image_mode="RGBA",
                    sources="upload",
                    type="pil",
                    elem_id="content_image",
                )
                processed_image = gr.Image(
                    label="Processed Image",
                    image_mode="RGBA",
                    type="pil",
                    interactive=False
                )
            with gr.Row():
                submit_process = gr.Button(
                    "Remove Background", elem_id="process", variant="primary")
            with gr.Row(variant="panel"):
                gr.Examples(
                    examples=[os.path.join("examples/images", img_name)
                              for img_name in sorted(os.listdir("examples/images"))],
                    inputs=[input_image],
                    label="Examples",
                    cache_examples=False,
                    examples_per_page=16
                )

    submit_prompt.click(fn=generate_image, inputs=[subject, style, color_scheme, angle, lighting_type, additional_details], outputs=input_image).success(
        fn=preprocess, inputs=[input_image], outputs=[processed_image]
    )
    submit_process.click(fn=check_input_image, inputs=[input_image]).success(
        fn=preprocess, inputs=[input_image], outputs=[processed_image],
    )

    return input_image, processed_image