import os
import torch
import uuid
from pathlib import Path
import gradio as gr
from gradio_imageslider import ImageSlider  # Ensure this library is installed

if os.environ.get("SPACES_ZERO_GPU") is not None:
    import spaces
else:
    class spaces:
        @staticmethod
        def GPU(func):
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)

            return wrapper

from diffusers.utils import check_min_version
from flux.controlnet_flux import FluxControlNetModel
from flux.transformer_flux import FluxTransformer2DModel
from flux.pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline

# Import configuration
import config

# Define the output folder
output_folder = Path('output_images')
output_folder.mkdir(exist_ok=True)

# Login to Hugging Face Hub
# huggingface_hub.login(os.getenv('HF_TOKEN_FLUX'))
check_min_version("0.30.2")

# Load models
transformer = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dtype=torch.bfloat16
)
controlnet = FluxControlNetModel.from_pretrained(
    "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", torch_dtype=torch.bfloat16
)

# Build pipeline
pipe = FluxControlNetInpaintingPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    controlnet=controlnet,
    transformer=transformer,
    torch_dtype=torch.bfloat16
).to("cuda")
pipe.transformer.to(torch.bfloat16)
pipe.controlnet.to(torch.bfloat16)


@spaces.GPU()
def process(input_image_editor):
    # Use default values from config
    negative_prompt = config.DEFAULT_NEGATIVE_PROMPT
    controlnet_conditioning_scale = config.DEFAULT_CONTROLNET_CONDITIONING_SCALE
    guidance_scale = config.DEFAULT_GUIDANCE_SCALE
    seed = config.DEFAULT_SEED
    num_inference_steps = config.DEFAULT_NUM_INFERENCE_STEPS
    true_guidance_scale = config.DEFAULT_TRUE_GUIDANCE_SCALE

    # Process image and mask
    image = input_image_editor['background']
    mask = input_image_editor['layers'][0]
    size = (768, 768)
    image_or = image.copy()
    image = image.convert("RGB").resize(size)
    mask = mask.convert("RGB").resize(size)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    # Generate result
    result = pipe(
        prompt="nothing",#os.getenv('MAGIC_PROMPT'),
        height=size[1],
        width=size[0],
        control_image=image,
        control_mask=mask,
        num_inference_steps=num_inference_steps,
        generator=generator,
        controlnet_conditioning_scale=controlnet_conditioning_scale,
        guidance_scale=guidance_scale,
        negative_prompt=negative_prompt,
        true_guidance_scale=true_guidance_scale
    ).images[0]

    processed_image = result.resize(image_or.size[:2])

    # Save the processed image
    output_folder = Path("output")  # Make sure this folder exists or create it
    output_folder.mkdir(parents=True, exist_ok=True)
    image_path = output_folder / f"no_bg_image_{uuid.uuid4().hex}.png"
    processed_image.save(image_path)

    return (processed_image, image), str(image_path)

#################################################################

# Define inputs and outputs for the Gradio interface
image = gr.ImageEditor(
                label='Image',
                type='pil',
                sources=["upload", "webcam"],
                image_mode='RGB',
                layers=False,
                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed")
            )
output_slider = ImageSlider(label="Processed photo", type="pil")

demo = gr.Interface(
    fn=process,
    inputs=image,
    outputs=[output_slider, gr.File(label="output png file")],
    #title="🫧 Snap Clean 🧽",
    description=config.DESCRIPTION
)


demo.launch(debug=False, show_error=True, share=True)