"""Image utilities for conversion, masking, and cropping."""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np
from numpy.typing import NDArray
from PIL import Image


def extract_masked_region(
    image: Image.Image,
    mask: NDArray[np.bool_],
) -> Image.Image:
    """Extract masked region from image.

    Args:
        image: Original PIL Image.
        mask: Binary mask as numpy array (True = keep).

    Returns:
        PIL Image with only the masked region visible.
    """
    image_np = np.array(image.convert("RGB"))
    masked_np = image_np * mask[:, :, np.newaxis]
    return Image.fromarray(masked_np.astype(np.uint8))


def crop_image_by_bbox(
    image: Image.Image,
    bbox: Sequence[float],
) -> Image.Image:
    """Crop an image by bounding box [x1, y1, x2, y2].

    Args:
        image: Source PIL image.
        bbox: OWLv2-style box coordinates [x1, y1, x2, y2].

    Returns:
        Cropped PIL image. Returns the original image when bbox is invalid.
    """
    if len(bbox) != 4:
        return image

    x1, y1, x2, y2 = tuple(float(v) for v in bbox)
    if not np.isfinite([x1, y1, x2, y2]).all():
        return image

    left = max(0, int(np.floor(x1)))
    top = max(0, int(np.floor(y1)))
    right = min(image.width, int(np.ceil(x2)))
    bottom = min(image.height, int(np.ceil(y2)))

    if right <= left or bottom <= top:
        return image

    return image.crop((left, top, right, bottom))


def numpy_to_pil(rgb: NDArray[np.uint8]) -> Image.Image:
    """Convert an RGB numpy array to a PIL Image.

    Handles arrays with 4 channels (RGBA) by dropping the alpha channel.

    Args:
        rgb: Numpy array of shape (H, W, C) with dtype uint8 or compatible.

    Returns:
        PIL Image in RGB mode.
    """
    rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
    return Image.fromarray(rgb3.astype(np.uint8))