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

def process_and_save(image_dict):
    if image_dict is None:
        return None, None, None

    # ImageEditor returns a dict with 'background' and 'layers' keys
    if 'background' in image_dict:
        original = image_dict['background']
    else:
        return None, None, None

    # Create mask from layers (drawn parts)
    if 'layers' in image_dict and len(image_dict['layers']) > 0:
        # Combine all layers to create mask
        mask = np.zeros_like(original[:, :, 0])
        for layer in image_dict['layers']:
            if layer is not None:
                # Convert layer to grayscale if needed
                if layer.ndim == 3:
                    layer_gray = np.any(layer[:, :, :3] > 0, axis=2)
                else:
                    layer_gray = layer > 0
                mask = np.logical_or(mask, layer_gray)
    else:
        # If no layers, create empty mask
        mask = np.zeros_like(original[:, :, 0])

    # Convert to binary mask
    binary_mask = (mask > 0).astype(np.uint8) * 255

    mask_img = Image.fromarray(binary_mask, mode='L')
    original_img = Image.fromarray(original)

    mask_img.save("mask.png")

    return mask_img, original_img, "mask.png"

with gr.Blocks() as demo:
    gr.Markdown("# Image Masking Tool")

    with gr.Row():
        with gr.Column():
            image_input = gr.ImageEditor(
                label="Draw Mask (Upload image, then draw with brush)",
                type="numpy",
                height=512,
                brush=gr.Brush(colors=["#ffffff"], default_color="#ffffff", color_mode="fixed")
            )

            process_btn = gr.Button("Generate Mask", variant="primary")

        with gr.Column():
            mask_output = gr.Image(label="Binary Mask", type="pil")
            original_output = gr.Image(label="Original Image", type="pil")
            mask_download = gr.File(label="Download Mask")

    process_btn.click(
        fn=process_and_save,
        inputs=image_input,
        outputs=[mask_output, original_output, mask_download]
    )

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