# app.py
from PIL import Image, ImageDraw
import gradio as gr

MAX_SIDE = 900

def fit_canvas(img):
    if img is None:
        return MAX_SIDE, MAX_SIDE
    w, h = img.size
    scale = MAX_SIDE / max(w, h)
    return int(w * scale), int(h * scale)

def snap_last_stroke(val):
    """Convert any roughly-dragged stroke into a filled rectangle of its color."""
    if not val or not val["layers"]:
        return val
    last = val["layers"][-1]
    box = last.getbbox()
    if box:
        mx, my = (box[0]+box[2])//2, (box[1]+box[3])//2
        try:
            color = last.getpixel((mx, my))
        except:
            color = (255,255,255,255)
        rect = Image.new("RGBA", last.size, (0,0,0,0))
        ImageDraw.Draw(rect).rectangle(box, fill=color)
        val["layers"][-1] = rect
        base = val["background"].copy()
        for layer in val["layers"]:
            base.alpha_composite(layer)
        val["composite"] = base
        # collapse layers so editor.value is always the merged image
        val["background"], val["layers"] = base, []
    return val

def load_image(img):
    return gr.update(value=img, canvas_size=fit_canvas(img))

def update_brush(color_hex):
    return gr.update(
        brush=gr.Brush(
            colors=[color_hex],
            default_color=color_hex,
            default_size=80,
            color_mode="fixed"
        )
    )

with gr.Blocks(css=".gradio-container { max-width:1200px; margin:auto; }") as demo:
    with gr.Row():
        color_picker = gr.ColorPicker("#FFFFFF", label="Brush Color")
        uploader     = gr.Image(type="pil", label="Upload Image")
        editor       = gr.ImageEditor(
            type="pil",  # emit PIL.Images so the download icon works
            brush=gr.Brush(colors=["#FFFFFF"], default_color="#FFFFFF", default_size=80, color_mode="fixed"),
            eraser=gr.Eraser(default_size=80),
            layers=False,
            canvas_size=(MAX_SIDE, MAX_SIDE),
        )

    # snapping callback
    editor.apply(snap_last_stroke, inputs=editor, outputs=editor)
    # wire color picker & uploader
    color_picker.change(update_brush, color_picker, editor)
    uploader.change(load_image, uploader, editor)

    demo.launch(share=True)