import gradio as gr
from PIL import Image
import numpy as np
import uuid

def editor_to_mask(value, invert=False):
    """
    Convert an ImageEditor EditorValue to a white-on-black mask.
    value: dict with keys: 'background', 'layers', 'composite' (Gradio 4/5 ImageEditor)
    Strategy:
      - Create an empty (black) mask the size of the background (or composite if background is None)
      - For each layer, use its alpha channel as painted region; OR it to the mask
    """
    if not value:
        return None, None

    bg = value.get("background")
    layers = value.get("layers") or []
    comp = value.get("composite")

    # Determine canvas size
    ref_img = bg or comp
    if ref_img is None:
        return None, None
    if not isinstance(ref_img, Image.Image):
        ref_img = Image.fromarray(ref_img)

    w, h = ref_img.size
    mask_acc = np.zeros((h, w), dtype=np.uint8)

    for layer in layers:
        if layer is None:
            continue
        if not isinstance(layer, Image.Image):
            layer = Image.fromarray(layer)
        # Ensure RGBA to get alpha; if no alpha, treat non-black as painted
        if layer.mode != "RGBA":
            layer = layer.convert("RGBA")
        la = np.array(layer)[:, :, 3]  # alpha
        # Any non-zero alpha counts as painted
        mask_acc = np.maximum(mask_acc, (la > 0).astype(np.uint8) * 255)

    if invert:
        mask_acc = 255 - mask_acc

    out = Image.fromarray(mask_acc, mode="L")
    fname = f"mask_{uuid.uuid4().hex}.png"
    out.save(fname)
    return out, fname

with gr.Blocks(css="footer {visibility:hidden}") as demo:
    gr.Markdown("# 🖌️ Quick Mask Maker (Gradio 5)\nUpload an image, draw on the canvas (white brush), then export a white-on-black PNG mask.")

    with gr.Row():
        with gr.Column(scale=3):
            brush_size = gr.Slider(5, 120, value=40, step=1, label="Brush radius")

            editor = gr.ImageEditor(
                label="Draw mask on its own layer (preselected).",
                type="pil",
                image_mode="RGBA",
                height=520,
                # Force a single "Mask" layer, disable adding extra layers
                layers=gr.LayerOptions(allow_additional_layers=False, layers=["Mask"]),
                # Force a white-only brush so painted area clearly becomes the mask
                brush=gr.Brush(default_size=40, colors=["#FFFFFF"], default_color="#FFFFFF", color_mode="fixed"),
                eraser=gr.Eraser(default_size=60),
            )

            invert = gr.Checkbox(False, label="Invert mask (white background)")
            gen = gr.Button("Generate Mask", variant="primary")

        with gr.Column(scale=2):
            out_mask = gr.Image(label="Mask Preview (PNG)", type="pil", image_mode="L")
            file_dl = gr.File(label="Download Mask")

    # Wire brush size slider to editor
    def _set_brush(r):
        return gr.update(brush=gr.Brush(default_size=int(r), colors=["#FFFFFF"], default_color="#FFFFFF", color_mode="fixed"))
    brush_size.change(_set_brush, inputs=brush_size, outputs=editor)

    gen.click(fn=editor_to_mask, inputs=[editor, invert], outputs=[out_mask, file_dl])

if __name__ == "__main__":
    demo.launch()