import os
import torch
import spaces
import gradio as gr
from diffusers import FluxFillPipeline
import random
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps


CSS = """
h1 {
    margin-top: 10px
}
"""

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
MAX_SEED = np.iinfo(np.int32).max

repo_id = "black-forest-labs/FLUX.1-Fill-dev"

if torch.cuda.is_available():
    pipe = FluxFillPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda")

@spaces.GPU()
def gen(
        prompt,
        image,
        mask_image,
        width,
        height,
        num_inference_steps,
        seed,
        guidance_scale, 
):
    generator = torch.Generator("cpu").manual_seed(seed)
    result = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask_image,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale,
        max_sequence_length=512,
    ).images[0]
    
    return result


def inpaintGen(
        imgMask,
        inpaint_prompt: str,
        guidance: float,
        num_steps: int,
        seed: int,
        randomize_seed: bool,
        progress=gr.Progress(track_tqdm=True)):

    source_path = imgMask["background"]
    mask_path = imgMask["layers"][0]

    if not source_path:
        raise gr.Error("Please upload an image.")

    if not mask_path:
        raise gr.Error("Please draw a mask on the image.")

    source_img = Image.open(source_path).convert("RGB")
    mask_img = Image.open(mask_path)
    alpha_channel=mask_img.split()[3]
    binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
    
    width, height = source_img.size

    new_width = (width // 16) * 16
    new_height = (height // 16) * 16
        
    # If the image size is not already divisible by 16, resize it
    if width != new_width or height != new_height:
        source_img = source_img.resize((new_width, new_height), Image.LANCZOS)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator("cpu").manual_seed(seed)

    result = gen(
        inpaint_prompt,
        source_img,
        binary_mask,
        new_width,
        new_height,
        num_steps,
        seed,
        guidance,
    )
    
    return result, seed


def add_border_and_mask(image, zoom_all=1.0, zoom_left=0, zoom_right=0, zoom_up=0, zoom_down=0, overlap=0.01):
    """Adds a black border around the image with individual side control and mask overlap"""
    orig_width, orig_height = image.size

    # Calculate padding for each side (in pixels)
    left_pad = int(orig_width * zoom_left)
    right_pad = int(orig_width * zoom_right)
    top_pad = int(orig_height * zoom_up)
    bottom_pad = int(orig_height * zoom_down)

    # Calculate overlap in pixels
    overlap_left = int(orig_width * overlap)
    overlap_right = int(orig_width * overlap)
    overlap_top = int(orig_height * overlap)
    overlap_bottom = int(orig_height * overlap)

    # If using the all-sides zoom, add it to each side
    if zoom_all > 1.0:
        extra_each_side = (zoom_all - 1.0) / 2
        left_pad += int(orig_width * extra_each_side)
        right_pad += int(orig_width * extra_each_side)
        top_pad += int(orig_height * extra_each_side)
        bottom_pad += int(orig_height * extra_each_side)

    # Calculate new dimensions (ensure they're multiples of 32)
    new_width = 32 * round((orig_width + left_pad + right_pad) / 32)
    new_height = 32 * round((orig_height + top_pad + bottom_pad) / 32)

    # Create new image with black border
    bordered_image = Image.new("RGB", (new_width, new_height), (0, 0, 0))
    # Paste original image in position
    paste_x = left_pad
    paste_y = top_pad
    bordered_image.paste(image, (paste_x, paste_y))

    # Create mask (white where the border is, black where the original image was)
    mask = Image.new("L", (new_width, new_height), 255)  # White background
    # Paste black rectangle with overlap adjustment
    mask.paste(
        0,
        (
            paste_x + overlap_left,  # Left edge moves right
            paste_y + overlap_top,  # Top edge moves down
            paste_x + orig_width - overlap_right,  # Right edge moves left
            paste_y + orig_height - overlap_bottom,  # Bottom edge moves up
        ),
    )

    return bordered_image, mask


def outpaintGen(
    img,
    outpaint_prompt: str,
    overlap: float,
    zoom_all: float,
    zoom_left: float,
    zoom_right: float,
    zoom_up: float,
    zoom_down: float,
    guidance: float,
    num_steps: int,
    seed: int,
    randomize_seed: bool
):
    image = Image.open(img)

    new_image, mask_image = add_border_and_mask(
            image,
            zoom_all=zoom_all,
            zoom_left=zoom_left,
            zoom_right=zoom_right,
            zoom_up=zoom_up,
            zoom_down=zoom_down,
            overlap=overlap,
    )

    width, height = new_image.size
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    result = gen(
        outpaint_prompt,
        new_image,
        mask_image,
        width,
        height,
        num_steps,
        seed,
        guidance,
    )
    
    return result, seed


with gr.Blocks(theme="ocean", title="Flux.1 Fill dev", css=CSS) as demo:
    gr.HTML("<h1><center>Flux.1 Fill dev</center></h1>")
    gr.HTML("""
        <p>
            <center>
                FLUX.1 Fill [dev] is a 12 billion parameter rectified flow transformer capable of filling areas in existing images based on a text description.
            </center>
        </p>
    """)
    with gr.Tab("Inpainting"):
        with gr.Row():
            with gr.Column():
                imgMask = gr.ImageMask(type="filepath", label="Image", layers=False, height=800)
                inpaint_prompt = gr.Textbox(label='Prompts ✏️', placeholder="A hat...")
                with gr.Row():
                    Inpaint_sendBtn = gr.Button(value="Submit", variant='primary')
                    Inpaint_clearBtn = gr.ClearButton([imgMask, inpaint_prompt], value="Clear")
            image_out = gr.Image(type="pil", label="Output", height=960)
        with gr.Accordion("Advanced ⚙️", open=False):
            guidance = gr.Slider(label="Guidance scale", minimum=1, maximum=50, value=30.0, step=0.1)
            num_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
            seed = gr.Number(label="Seed", value=42, precision=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
    
        gr.on(
            triggers = [
                inpaint_prompt.submit,
                Inpaint_sendBtn.click,
            ],
            fn = inpaintGen,
            inputs = [
                imgMask,
                inpaint_prompt,
                guidance,
                num_steps,
                seed,
                randomize_seed
            ],
            outputs = [image_out, seed]
        )

    with gr.Tab("Outpainting"):
        with gr.Row():
            with gr.Column():
                img = gr.Image(type="filepath", label="Image", height=800)
                outpaint_prompt = gr.Textbox(label='Prompts ✏️', placeholder="In city...")
                with gr.Row():
                    outpaint_sendBtn = gr.Button(value="Submit", variant='primary')
                    outpaint_clearBtn = gr.ClearButton([img, outpaint_prompt], value="Clear")
            image_exp = gr.Image(type="pil", label="Output", height=960)
        with gr.Accordion("Advanced ⚙️", open=False):
            overlap = gr.Slider(label="Overlap", minimum=0.01, maximum=0.25, value=0.01, step=0.01)
            zoom_all = gr.Slider(label="Zoom Out Amount (All Sides)", minimum=1.0, maximum=3.0, value=1.0, step=0.1)
            with gr.Row():
                zoom_left = gr.Slider(label="Left", minimum=0.0, maximum=1.0, value=0.0, step=0.1)
                zoom_right = gr.Slider(label="Right", minimum=0.0, maximum=1.0, value=0.0, step=0.1)
            with gr.Row():
                zoom_up = gr.Slider(label="Up", minimum=0.0, maximum=1.0, value=0.0, step=0.1)
                zoom_down = gr.Slider(label="Down", minimum=0.0, maximum=1.0, value=0.0, step=0.1)
            op_guidance = gr.Slider(label="Guidance scale", minimum=1, maximum=50, value=30.0, step=0.1)
            op_num_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
            op_seed = gr.Number(label="Seed", value=42, precision=0)
            op_randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
    
        gr.on(
            triggers = [
                outpaint_prompt.submit,
                outpaint_sendBtn.click,
            ],
            fn = outpaintGen,
            inputs = [
                img,
                outpaint_prompt,
                overlap,
                zoom_all,
                zoom_left,
                zoom_right,
                zoom_up,
                zoom_down,
                op_guidance,
                op_num_steps,
                op_seed,
                op_randomize_seed
            ],
            outputs = [image_exp, op_seed]
        )


if __name__ == "__main__":
    demo.launch(show_api=False, share=False)