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

def brightness_to_opacity_overlay(foreground_img, background_img, invert_opacity=False):
    if foreground_img is None or background_img is None:
        return None

    # Resize foreground to match background
    bg = background_img.convert("RGBA")
    fg = foreground_img.convert("RGB").resize(bg.size)

    # Convert to luminance (grayscale)
    fg_array = np.array(fg)
    luminance = np.dot(fg_array[...,:3], [0.299, 0.587, 0.114]).astype(np.uint8)

    # Invert if requested
    alpha = 255 - luminance if invert_opacity else luminance

    # Create white RGBA image with computed alpha
    white_rgb = np.ones_like(fg_array) * 255
    rgba_array = np.dstack((white_rgb, alpha)).astype(np.uint8)
    overlay = Image.fromarray(rgba_array, mode="RGBA")

    # Manual alpha blending
    bg_arr = np.array(bg).astype(np.float32) / 255.0
    ov_arr = np.array(overlay).astype(np.float32) / 255.0

    # Extract alpha channels
    alpha_bg = bg_arr[..., 3:4]
    alpha_ov = ov_arr[..., 3:4]

    # Composite alpha and color
    out_alpha = alpha_ov + alpha_bg * (1 - alpha_ov)
    out_rgb = (
        ov_arr[..., :3] * alpha_ov +
        bg_arr[..., :3] * alpha_bg * (1 - alpha_ov)
    ) / np.clip(out_alpha, 1e-6, 1)

    # Combine and convert to final image
    out_image = np.dstack((out_rgb, out_alpha)).clip(0, 1) * 255
    out_image = Image.fromarray(out_image.astype(np.uint8), mode="RGBA")

    return out_image

# Gradio UI
iface = gr.Interface(
    fn=brightness_to_opacity_overlay,
    inputs=[
        gr.Image(type="pil", label="Mask Source (Brightness to Alpha)"),
        gr.Image(type="pil", label="Background Image (with or without transparency)"),
        gr.Checkbox(label="Invert Opacity", value=False)
    ],
    outputs=gr.Image(type="pil", label="Final Composite"),
    title="Brightness-to-Alpha Overlay Tool (with Transparency Fix)",
    description=(
        "Uses the brightness of the first image as alpha, resizes to match the second image, "
        "and overlays correctly even when the background has transparency."
    )
)

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