import spaces
import gradio as gr
import torch
import uuid
import os

from PIL import Image
from enhance_utils import enhance_image

DEFAULT_SRC_PROMPT = "a woman, photo"
DEFAULT_EDIT_PROMPT = "a beautiful woman, photo, hollywood style face, 8k, high quality"

device = "cuda" if torch.cuda.is_available() else "cpu"

def create_demo() -> gr.Blocks:
    from inversion_run_base import run as base_run

    @spaces.GPU(duration=10)
    def image_to_image(
        input_image_path: str,
        input_image_prompt: str,
        edit_prompt: str,
        seed: int,
        w1: float,
        num_steps: int,
        start_step: int,
        guidance_scale: float,
        enhance_face: bool = True,
    ):
        w2 = 1.0

        input_image = Image.open(input_image_path)

        run_model = base_run
        res_image = run_model(
            input_image,
            input_image_prompt,
            edit_prompt,
            seed,
            w1,
            w2,
            num_steps,
            start_step,
            guidance_scale,
        )
        enhanced_image = enhance_image(res_image, enhance_face)

        tmpPrefix = "/tmp/gradio/"

        extension = 'png'
        if enhanced_image.mode == 'RGBA':
            extension = 'png'
        else:
            extension = 'jpg'

        targetDir = f"{tmpPrefix}output/"
        if not os.path.exists(targetDir):
            os.makedirs(targetDir)

        enhanced_path = f"{targetDir}{uuid.uuid4()}.{extension}"
        enhanced_image.save(enhanced_path, quality=100)

        return enhanced_path

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                input_image_path = gr.File(label="Input Image")
            with gr.Column():
                generated_image_path = gr.File(label="Download the segment image", interactive=False)
        with gr.Row():
            with gr.Column():
                input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
                with gr.Accordion("Advanced Options", open=False):
                    guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale")
                    enhance_face = gr.Checkbox(label="Enhance Face", value=False)
                    seed = gr.Number(label="Seed", value=8)
            with gr.Column():
                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
                start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step")
                w1 = gr.Number(label="W1", value=2)
                g_btn = gr.Button("Edit Image")
                
        
        g_btn.click(
            fn=image_to_image,
            inputs=[input_image_path, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, enhance_face],
            outputs=[generated_image_path],
        )

    return demo