import base64
from io import BytesIO
import spaces
import gradio as gr
import torch
from PIL import Image, ImageChops

from diffusers import Flux2KleinPipeline

from utils import process_source, process_reference, paste_back, binarize_mask

pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-base-4B", torch_dtype=torch.bfloat16
)
pipe.to("cuda")

pipe.load_lora_weights("LiXiY/Easy-Insert")

PROMPT = "Replace the white mask of image1 with the content in image2. Preserving the background, lighting, and surrounding elements, maintain a seamless and natural result."


# ===================== Example Data =====================

CANVAS_W, CANVAS_H = 1024, 1024

EXAMPLES = [
    ("examples/background_image/1.png", "examples/insert_mask/1.png", "examples/ref_image/1.png", "examples/ref_mask/1.png"),
    ("examples/background_image/2.png", "examples/insert_mask/2.png", "examples/ref_image/2.png", "examples/ref_mask/2.png"),
    ("examples/background_image/3.png", "examples/insert_mask/3.png", "examples/ref_image/3.png", "examples/ref_mask/3.png"),
    ("examples/background_image/4.png", "examples/insert_mask/4.png", "examples/ref_image/4.png", "examples/ref_mask/4.png"),
]


# ===================== Thumbnail HTML Generation =====================

def img_to_b64(path, size=(240, 240)):
    img = Image.open(path).convert("RGB")
    # Preserve the original aspect ratio; fit inside `size`. 2x the display
    # width (110 CSS px) so thumbnails stay sharp on HiDPI screens.
    img.thumbnail(size, Image.LANCZOS)
    buf = BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def build_examples_html():
    row_pairs = [EXAMPLES[0:2], EXAMPLES[2:4]]

    html = '<div class="ex-grid">'
    for row_idx, row_examples in enumerate(row_pairs):
        html += '<div class="ex-grid-row">'
        for col_idx, (bg, mask, ref, ref_mask) in enumerate(row_examples):
            i = row_idx * 2 + col_idx
            bg_b64 = img_to_b64(bg)
            mask_b64 = img_to_b64(mask)
            ref_b64 = img_to_b64(ref)
            ref_mask_b64 = img_to_b64(ref_mask)
            html += f'''
            <div class="ex-row" onclick="(function(){{var el=document.getElementById('ex_btn_{i}');if(!el)return;var btn=el.querySelector('button')||el;btn.dispatchEvent(new MouseEvent('click',{{bubbles:true,cancelable:true}}));}})()">
                <div class="ex-label">Example {i + 1}</div>
                <div class="ex-thumbs">
                    <div class="ex-thumb-wrap">
                        <img src="data:image/png;base64,{bg_b64}" class="ex-thumb" draggable="false"/>
                        <span class="ex-thumb-sublabel">Background</span>
                    </div>
                    <div class="ex-thumb-wrap">
                        <img src="data:image/png;base64,{mask_b64}" class="ex-thumb" draggable="false"/>
                        <span class="ex-thumb-sublabel">Mask</span>
                    </div>
                    <div class="ex-thumb-wrap">
                        <img src="data:image/png;base64,{ref_b64}" class="ex-thumb" draggable="false"/>
                        <span class="ex-thumb-sublabel">Reference</span>
                    </div>
                    <div class="ex-thumb-wrap">
                        <img src="data:image/png;base64,{ref_mask_b64}" class="ex-thumb" draggable="false"/>
                        <span class="ex-thumb-sublabel">Ref Mask</span>
                    </div>
                </div>
            </div>
            '''
        html += '</div>'
    html += '</div>'
    return html


# ===================== Utility Functions =====================

def extract_mask_from_layers(layers, target_size):
    mask = Image.new("L", target_size, 0)
    for layer in layers:
        if layer is not None:
            layer_rgba = layer.convert("RGBA").resize(target_size)
            alpha = layer_rgba.split()[3]
            alpha_binary = alpha.point(lambda x: 255 if x > 0 else 0)
            mask = Image.composite(Image.new("L", target_size, 255), mask, alpha_binary)
    return mask


EDITOR_SIZE = 800
# Matches Gradio light-mode image editor background so padded examples blend in.
EDITOR_BG_COLOR = (243, 244, 246)


def fit_transform(img, editor_size=EDITOR_SIZE):
    """Display transform the fixed-canvas editor applies to `img`: scale to fit
    inside the square canvas, then center-pad. Returns (disp_size, pad)."""
    scale = min(editor_size / img.width, editor_size / img.height, 1.0)
    disp_size = (
        max(1, round(img.width * scale)),
        max(1, round(img.height * scale)),
    )
    pad = (
        (editor_size - disp_size[0]) // 2,
        (editor_size - disp_size[1]) // 2,
    )
    return disp_size, pad


def build_display_bg(bg, editor_size=EDITOR_SIZE):
    """Rebuild the padded square background the editor shows for `bg`."""
    disp_size, pad = fit_transform(bg, editor_size)
    canvas = Image.new("RGB", (editor_size, editor_size), EDITOR_BG_COLOR)
    canvas.paste(bg.convert("RGB").resize(disp_size, Image.LANCZOS), pad)
    return canvas


def images_close(a, b, tol_frac=0.001):
    """True if two RGB images are (nearly) pixel-identical."""
    if a.size != b.size:
        return False
    diff = ImageChops.difference(a.convert("RGB"), b.convert("RGB"))
    changed = sum(diff.histogram()[1:])
    return changed <= tol_frac * a.size[0] * a.size[1] * 3


def make_editor_value(bg_pil, mask_pil=None, editor_size=EDITOR_SIZE):
    """Build an ImageEditor value on a fixed square canvas.

    The background is downscaled to fit inside `editor_size` and padded to a
    square; the mask is padded the same way and sent as a real layer. This
    keeps the mask and background perfectly aligned in the frontend (both are
    800x800 and the padding is identical), while the full-resolution image and
    exact mask are kept in gr.State for inference.

    Returns (editor_value, sent_mask, transform).
    """
    bg = bg_pil.convert("RGB")
    disp_size, pad = fit_transform(bg, editor_size)
    bg_sq = build_display_bg(bg, editor_size)

    transform = {
        "editor_size": editor_size,
        "disp_size": disp_size,
        "pad": pad,
    }

    if mask_pil is None:
        # composite=None avoids saving a second, redundant image file. The editor
        # can render the background directly; if it needs a composite it will build
        # it from background + (empty) layers.
        return {"background": bg_sq, "layers": [], "composite": None}, None, transform

    mask = mask_pil.convert("L").resize(disp_size, Image.NEAREST)
    mask_sq = Image.new("L", (editor_size, editor_size), 0)
    mask_sq.paste(mask, pad)

    transparent = Image.new("RGBA", (editor_size, editor_size), (0, 0, 0, 0))
    # Semi-transparent white (alpha 153 ≈ 60%) so the background shows through
    # the loaded mask. Mask extraction binarizes alpha > 0, so this still
    # resolves to a solid mask at inference time.
    white_solid = Image.new("RGBA", (editor_size, editor_size), (255, 255, 255, 153))
    mask_layer = Image.composite(white_solid, transparent, mask_sq)

    return {"background": bg_sq, "layers": [mask_layer], "composite": None}, mask_sq, transform


def masks_close(a, b, tol_frac=0.0):
    """True if two binarized masks are pixel-identical (within tol_frac).
    The fixed-canvas PNG round-trip is exact, so we treat any non-zero
    difference as a real user edit."""
    if a.size != b.size:
        b = b.resize(a.size, Image.NEAREST)
    diff = ImageChops.difference(binarize_mask(a), binarize_mask(b))
    changed = sum(diff.histogram()[1:])
    return changed <= tol_frac * a.size[0] * a.size[1]


def load_example(idx):
    bg_path, mask_path, ref_path, ref_mask_path = EXAMPLES[idx]
    bg_img = Image.open(bg_path).convert("RGB")
    ref_img = Image.open(ref_path).convert("RGB")
    base_val, base_sent, base_transform = make_editor_value(
        bg_img, Image.open(mask_path)
    )
    ref_val, ref_sent, ref_transform = make_editor_value(
        ref_img, Image.open(ref_mask_path)
    )
    # Keep the pristine full-res image + exact mask in State. The editor value
    # is only for display; its small size avoids slow frontend re-uploads.
    base_state = (bg_img, Image.open(mask_path).convert("L"), base_sent, base_transform)
    ref_state = (ref_img, Image.open(ref_mask_path).convert("L"), ref_sent, ref_transform)
    return base_val, ref_val, base_state, ref_state


def load_ex1():
    return load_example(0)


def load_ex2():
    return load_example(1)


def load_ex3():
    return load_example(2)


def load_ex4():
    return load_example(3)


# ===================== Generation Function =====================

def resolve_source(editor_value, state):
    """Return (full_res_image, full_res_mask) for one editor.

    `state` is (image, mask, sent_mask, transform) captured when an example was
    loaded. The editor value's background is display-only, so we don't trust it
    blindly: rebuild the padded display image from the state's pristine image
    and compare it with what the editor currently shows. If they match, the
    example is still loaded and the state's full-res image (+ exact mask) is
    used. If they differ, the user uploaded a different image - the editor
    background is then the full-res original, and the brush strokes are mapped
    back through the fixed-canvas transform.
    """
    bg_ed = editor_value.get("background")
    if bg_ed is None:
        return None, None
    img_ed = bg_ed.convert("RGB")

    ed_mask = extract_mask_from_layers(
        editor_value.get("layers", []), (EDITOR_SIZE, EDITOR_SIZE)
    )

    if state is not None:
        img, full_mask, sent_mask, transform = state
        img = img.convert("RGB")
        if images_close(img_ed, build_display_bg(img)):
            if full_mask is not None:
                # Example flow: exact full-res mask is in state. If the editor's
                # layer mask still matches what we sent (user hasn't brushed or
                # erased), use the state mask directly - bit-exact with inference.py.
                if sent_mask is not None and masks_close(ed_mask, sent_mask):
                    return img, binarize_mask(full_mask)
                # User edited the example's mask (erased and/or brushed): crop the
                # padding and map the edited mask back to the original resolution.
                if transform:
                    ed_mask_bin = binarize_mask(ed_mask)
                    px, py = transform["pad"]
                    dw, dh = transform["disp_size"]
                    cropped = ed_mask_bin.crop((px, py, px + dw, py + dh))
                    return img, binarize_mask(cropped.resize(img.size, Image.NEAREST))

    # Upload flow: the editor background is the user's full-res original.
    disp_size, pad = fit_transform(img_ed)
    ed_mask_bin = binarize_mask(ed_mask)
    cropped = ed_mask_bin.crop((pad[0], pad[1], pad[0] + disp_size[0], pad[1] + disp_size[1]))
    return img_ed, binarize_mask(cropped.resize(img_ed.size, Image.NEAREST))

@spaces.GPU
def run_local(base, ref, base_state, ref_state, seed, num_inference_steps, cfg_scale):
    if base is None or ref is None or not isinstance(base, dict) or not isinstance(ref, dict):
        return None, gr.update(visible=False)

    pil_bg, pil_mask = resolve_source(base, base_state)
    pil_ref, pil_ref_mask = resolve_source(ref, ref_state)
    if pil_bg is None or pil_ref is None:
        return None, gr.update(visible=False)

    if pil_mask.getextrema() == (0, 0) or pil_ref_mask.getextrema() == (0, 0):
        error_html = """
            <div class="error-overlay" style="
                position: fixed; top: 0; left: 0; width: 100%; height: 100%;
                background: rgba(0,0,0,0.5); display: flex; justify-content: center;
                align-items: center; z-index: 9999;
            ">
                <div style="
                    background: white; padding: 30px; border-radius: 10px;
                    text-align: center; font-size: 18px; box-shadow: 0 0 15px rgba(0,0,0,0.3);
                ">
                    <p style="color: red; margin-bottom: 20px;">
                        ⚠️ Please draw the mask on BOTH the background image and the reference image first, or click an example!
                    </p>
                    <button onclick="this.closest('.error-overlay').remove()"
                            style="padding: 8px 20px; cursor: pointer; border: none;
                                background: #eee; border-radius: 5px;">
                        OK
                    </button>
                </div>
            </div>
            """
        return None, gr.update(value=error_html, visible=True)

    background_image, _, crop_box, source_mask_cropped = process_source(pil_bg, pil_mask, CANVAS_W)
    ref_image = process_reference(pil_ref, pil_ref_mask, CANVAS_W)

    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    generated_image = pipe(
        image=[background_image, ref_image],
        prompt=PROMPT,
        height=CANVAS_H,
        width=CANVAS_W,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(cfg_scale),
        generator=generator,
    ).images[0]
    result_img = paste_back(generated_image, pil_bg, crop_box, source_mask_cropped, feather=0)

    return result_img, gr.update(visible=False)


# ===================== Gradio UI =====================

with gr.Blocks(css="""
    .input-row {
        overflow: visible !important;
    }

    .input-row .gr-image-editor {
        overflow: hidden !important;
    }
    .input-row .gr-image-editor .image-container,
    .input-row .gr-image-editor .canvas-container,
    .input-row .gr-image-editor canvas {
        max-width: 100% !important;
        max-height: 100% !important;
        object-fit: contain !important;
    }

    .ex-section-header {
        display: flex;
        align-items: center;
        gap: 10px;
        margin: 28px 0 14px 0;
        justify-content: center;
    }
    .ex-section-header::before {
        content: '';
        flex: 1;
        height: 1px;
        max-width: 180px;
        background: #e5e7eb;
    }
    .ex-section-header::after {
        content: '';
        flex: 1;
        height: 1px;
        max-width: 180px;
        background: #e5e7eb;
    }

    .ex-container {
        display: flex;
        flex-direction: column;
        align-items: center;
        padding-bottom: 20px;
    }

    .ex-grid {
        display: flex;
        flex-direction: column;
        align-items: center;
        gap: 12px;
        padding-bottom: 20px;
    }
    .ex-grid-row {
        display: flex;
        gap: 20px;
        justify-content: center;
        flex-wrap: wrap;
    }

    .ex-row {
        display: flex;
        align-items: center;
        gap: 20px;
        padding: 14px 28px;
        border: 2px solid #e5e7eb;
        border-radius: 12px;
        cursor: pointer;
        transition: all 0.25s ease;
        background: #ffffff;
        user-select: none;
        width: fit-content;
    }
    .ex-row:hover {
        border-color: #3b82f6;
        background: #f0f7ff;
        box-shadow: 0 4px 18px rgba(59, 130, 246, 0.15);
        transform: translateY(-2px);
    }
    .ex-row:active {
        transform: translateY(0);
        box-shadow: 0 2px 8px rgba(59, 130, 246, 0.2);
    }

    .ex-label {
        font-weight: 700;
        font-size: 15px;
        min-width: 62px;
        color: #1e40af;
        letter-spacing: 0.02em;
    }

    .ex-thumbs {
        display: flex;
        gap: 14px;
    }

    .ex-thumb-wrap {
        display: flex;
        flex-direction: column;
        align-items: center;
        gap: 6px;
    }

    .ex-thumb {
        width: 110px;
        height: auto;
        object-fit: contain;
        border-radius: 8px;
        border: 2px solid #e5e7eb;
        transition: all 0.25s ease;
        pointer-events: none;
    }
    .ex-row:hover .ex-thumb {
        border-color: #93c5fd;
    }

    .ex-thumb-sublabel {
        font-size: 12px;
        color: #6b7280;
        font-weight: 500;
    }

    /* Gradio >=5 does not mount components with visible=False into the DOM,
       so the example-thumbnail JS cannot find them. Hide via CSS instead:
       the buttons stay mounted (click handlers still fire) but invisible. */
    .ex-hidden-btn {
        display: none !important;
    }
""") as demo:

    gr.Markdown(
        "<h1 style='text-align: center;'>Reference-Based Object Insertion and Clothing Replacement</h1>"
        "<h3 style='text-align: center;'>Insert objects from reference images into masked regions of background images, and replace clothing in masked regions using reference images.</h3>"
        "<p style='text-align: center;'>"
        "<a href='https://github.com/huan-yin/Easy-Insert' target='_blank' "
        "style='color: #3b82f6; text-decoration: none; font-weight: 500;'>"
        "GitHub Repo: huan-yin/Easy-Insert</a></p>"
    )
    gr.Markdown(
        """
        **Instructions:**
        1. Upload a background image, then use the brush/eraser tools below the image to mark or refine the insertion region (mask); upload a reference image and brush/erase over the object to be inserted (ref mask).
        2. Or click any row of thumbnails in the "Examples" section below to automatically load background + mask + reference + ref mask.
        3. Click the "Generate" button, and the result will be displayed below.
        """
    )

    with gr.Row(elem_classes="input-row"):
        base = gr.ImageEditor(
            label="Background Image (brush the insertion region)",
            type="pil",
            format="png",
            width=420,
            height=450,
            sources=["upload"],
            canvas_size=(EDITOR_SIZE, EDITOR_SIZE),
            fixed_canvas=True,
            brush=gr.Brush(
                default_size=30,
                default_color="rgba(255, 255, 255, 0.6)",
                color_mode="fixed",
                colors=["rgba(255, 255, 255, 0.6)"],
            ),
            eraser=gr.Eraser(default_size=30),
        )
        ref = gr.ImageEditor(
            label="Reference Image (brush the object to insert)",
            type="pil",
            format="png",
            width=420,
            height=450,
            sources=["upload"],
            canvas_size=(EDITOR_SIZE, EDITOR_SIZE),
            fixed_canvas=True,
            brush=gr.Brush(
                default_size=30,
                default_color="rgba(255, 255, 255, 0.6)",
                color_mode="fixed",
                colors=["rgba(255, 255, 255, 0.6)"],
            ),
            eraser=gr.Eraser(default_size=30),
        )

    with gr.Row():
        seed = gr.Number(label="Seed", value=1, precision=0)
        num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=15)
        cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=10, step=0.1, value=4)

    with gr.Row():
        gen_btn = gr.Button("Generate", variant="primary")

    # ==================== Generation Result ====================
    with gr.Row():
        output_image = gr.Image(
            label="Generated Result",
            interactive=False,
            width=420,
            height=450,
        )

    with gr.Row():
        error_dialog = gr.HTML(visible=False)

    gr.HTML('<div class="ex-section-header"><span style="font-weight:700;font-size:16px;color:#374151;">Examples (click to load background + mask + reference + ref mask)</span></div>')

    base_state = gr.State(None)
    ref_state = gr.State(None)

    ex_btn0 = gr.Button("Example 1", elem_id="ex_btn_0", elem_classes=["ex-hidden-btn"])
    ex_btn1 = gr.Button("Example 2", elem_id="ex_btn_1", elem_classes=["ex-hidden-btn"])
    ex_btn2 = gr.Button("Example 3", elem_id="ex_btn_2", elem_classes=["ex-hidden-btn"])
    ex_btn3 = gr.Button("Example 4", elem_id="ex_btn_3", elem_classes=["ex-hidden-btn"])

    gr.HTML('<div class="ex-container">' + build_examples_html() + '</div>')

    # ==================== Event Bindings ====================

    # Single-step load: a clear-then-load chain sends the frontend two rapid
    # updates, and the upload round-trip in between references temp PNGs that
    # get overwritten mid-read (truncated/broken PNG errors on preprocess).
    ex_btn0.click(fn=load_ex1, outputs=[base, ref, base_state, ref_state])
    ex_btn1.click(fn=load_ex2, outputs=[base, ref, base_state, ref_state])
    ex_btn2.click(fn=load_ex3, outputs=[base, ref, base_state, ref_state])
    ex_btn3.click(fn=load_ex4, outputs=[base, ref, base_state, ref_state])

    gen_btn.click(
        fn=run_local,
        inputs=[base, ref, base_state, ref_state, seed, num_inference_steps, cfg_scale],
        outputs=[output_image, error_dialog],
    )

demo.launch(server_name="0.0.0.0", server_port=7860)