import base64
import io

from PIL import Image


def encode_image(img: Image.Image, image_format: str | None = None) -> str:
    """Encode a PIL image to a data URL in the specified format.

    image_format examples: "PNG", "JPEG", "WEBP"
    If not provided, fall back to WEBP.
    """

    params = {}

    if image_format is None:
        image_format = "WEBP"

    if image_format == "WEBP":
        params["quality"] = 90

    buffer = io.BytesIO()
    try:
        img.save(buffer, format=image_format, **params)
    except Exception as e:
        raise ValueError(f"Failed to encode image as {image_format}: {e}")

    raw = buffer.getvalue()
    b64 = base64.b64encode(raw).decode("ascii")
    mime = f"image/{image_format.lower()}"
    return f"data:{mime};base64,{b64}"


def decode_image(data: str) -> Image.Image:
    """Decode a data URL or base64 string into a PIL.Image.Image.

    Supports any format that Pillow supports (png, jpg, webp, etc).
    """

    if not isinstance(data, str):
        raise TypeError("decode_image expects a string")

    # strip data URL prefix if present
    if data.startswith("data:"):
        try:
            _, _, b64 = data.partition(",")
            data = b64
        except Exception:
            raise ValueError("Invalid data URL")

    # base64 decode
    try:
        raw = base64.b64decode(data.encode("ascii"))
    except Exception as e:
        raise ValueError(f"Base64 decode error: {e}")

    # open with Pillow (auto detects format)
    buffer = io.BytesIO(raw)
    try:
        img = Image.open(buffer)
        img.load()  # ensure fully loaded
        return img
    except Exception as e:
        raise ValueError(f"Failed to decode image: {e}")