import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from PIL import Image, ImageFilter


def get_prompt_list(prompt: str):
    return [txt.strip() for txt in prompt.split(".") if txt.strip() != ""]


def plot_groundingdino_boxes(img: Image.Image, result, figsize=(12, 8), show_scores=True):
    """
    img: PIL Image
    result: results[0] dict containing 'boxes', 'scores', 'labels'
    """
    img_np = np.array(img)

    plt.figure(figsize=figsize)
    plt.imshow(img_np)
    ax = plt.gca()

    boxes = result["boxes"]
    scores = result.get("scores", None)
    labels = result.get("labels", None)

    for i, box in enumerate(boxes):
        x1, y1, x2, y2 = box

        rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor="red", facecolor="none")
        ax.add_patch(rect)

        text = ""
        if labels is not None:
            text += str(labels[i])
        if show_scores and scores is not None:
            text += f" ({scores[i]:.2f})"

        if text:
            ax.text(x1, y1 - 5, text, color="red", fontsize=12, bbox=dict(facecolor="black", alpha=0.5, pad=2))

    plt.axis("off")
    plt.show()


MASK_OVERLAY_RGBA = (225, 29, 72, int(0.58 * 255))


def save_visual_mask_from_binary(mask, out_path):
    mask_bool = np.asarray(mask)
    if mask_bool.ndim == 3:
        mask_bool = np.any(mask_bool > 0, axis=0)
    elif mask_bool.dtype != bool:
        mask_bool = mask_bool > 0

    H, W = mask_bool.shape
    out = np.zeros((H, W, 4), dtype=np.uint8)
    out[mask_bool] = MASK_OVERLAY_RGBA
    Image.fromarray(out, mode="RGBA").save(out_path)


def save_visual_mask(mask_data, out_path):
    masks = np.asarray(mask_data)
    if masks.ndim == 2:
        combined_mask = masks > 0
    else:
        masks = (masks > 0.5) if masks.dtype != bool else masks.astype(bool)
        combined_mask = np.any(masks, axis=0) if masks.size else np.zeros(masks.shape[-2:], dtype=bool)

    save_visual_mask_from_binary(combined_mask, out_path)


def mask_image_to_binary(mask_img: Image.Image, size=None) -> Image.Image:
    rgba = mask_img.convert("RGBA")
    if size is not None and rgba.size != size:
        rgba = rgba.resize(size, Image.Resampling.NEAREST)

    rgba_arr = np.asarray(rgba)
    alpha = rgba_arr[..., 3]

    if np.any(alpha < 255):
        mask_bool = alpha > 12
    else:
        gray = np.asarray(rgba.convert("L"))
        mask_bool = gray > 127

    return Image.fromarray(mask_bool.astype(np.uint8) * 255, mode="L")


def blur_img_with_mask(
    original_img: Image.Image,
    mask_img: Image.Image,
    blur_radius: int = 15,
) -> Image.Image:
    original_image = original_img.convert("RGB")
    mask_image = mask_img.convert("L")

    # Create blurred version of image
    blurred_image = original_img.filter(ImageFilter.GaussianBlur(radius=blur_radius))

    result = Image.composite(
        blurred_image,
        original_image,
        mask_img,
    )

    return result