"""
Utilities for working with images.
"""

from io import BytesIO

from PIL import Image


def is_webp(image: bytes | BytesIO) -> bool:
    """Check if image is WebP format, analyzing the first bytes of the image.

    Args:
        image: The image to check.

    Returns:
        True if the image is in WebP format, False otherwise.
    """
    if isinstance(image, BytesIO):
        image = image.getvalue()
    return image[:4] == b"RIFF" and image[8:12] == b"WEBP"


def _convert_to(image: bytes | BytesIO, to_format: str) -> bytes:
    """Convert an image to a different format.

    Args:
        image: The image to convert.
        to_format: The format to convert the image to, e.g. "PNG".
    """
    image_data = image if isinstance(image, BytesIO) else BytesIO(image)
    with Image.open(image_data) as img:
        output = BytesIO()
        img.save(output, format=to_format)
        return output.getvalue()


def convert_to_png(image_bytes: bytes | BytesIO) -> bytes:
    """Convert an image to PNG format.

    Args:
        image_bytes: The image to convert.

    Returns:
        The image converted to PNG format.
    """
    return _convert_to(image_bytes, "PNG")