"""Decoding of the `image` input into a BGR frame that can be drawn on."""

import io

import cv2
import numpy as np
import pillow_avif  # noqa: F401  # registers the AVIF plugin with Pillow
from PIL import Image

# Encodings that carry their own dimensions, so `width`/`height` metadata is not
# consulted for them.
SELF_DESCRIBING = frozenset({"jpeg", "jpg", "jpe", "bmp", "webp", "png"})

SUPPORTED = SELF_DESCRIBING | {"bgr8", "rgb8", "yuv420", "avif"}


def to_bgr(storage, encoding: str, width: int | None, height: int | None) -> np.ndarray:
    """Decode one `image` payload into a writable BGR array.

    `width` and `height` come from the message metadata and are only required by
    the raw encodings, which carry no dimensions of their own.
    """
    buffer = storage.to_numpy().astype(np.uint8)

    if encoding == "bgr8":
        # Copy so the frame can be annotated: to_numpy() may be a view on the
        # shared-memory payload.
        return _reshape(buffer, width, height, encoding).copy()

    if encoding == "rgb8":
        return cv2.cvtColor(_reshape(buffer, width, height, encoding), cv2.COLOR_RGB2BGR)

    if encoding == "yuv420":
        # A planar YUV 4:2:0 frame is 1.5 bytes per pixel, laid out as a
        # single-channel image of `height * 3 // 2` rows.
        width, height = _dimensions(width, height, encoding)
        plane = _take(buffer, width * height * 3 // 2, encoding).reshape(height * 3 // 2, width)
        # COLOR_YUV420p2RGB, not ...2BGR: the upstream node converts to RGB and
        # then displays it as if it were BGR. The channel swap is part of the
        # behaviour callers already see.
        return cv2.cvtColor(plane, cv2.COLOR_YUV420p2RGB)

    if encoding in SELF_DESCRIBING:
        frame = cv2.imdecode(buffer, cv2.IMREAD_COLOR)
        if frame is None:
            raise RuntimeError(f"could not decode {encoding} image")
        return frame

    if encoding == "avif":
        image = Image.open(io.BytesIO(buffer.tobytes())).convert("RGB")
        return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    raise RuntimeError(f"unsupported image encoding: {encoding}")


def _reshape(buffer: np.ndarray, width: int | None, height: int | None, encoding: str):
    width, height = _dimensions(width, height, encoding)
    return _take(buffer, width * height * 3, encoding).reshape(height, width, 3)


def _dimensions(width: int | None, height: int | None, encoding: str) -> tuple[int, int]:
    if not width or not height:
        raise RuntimeError(
            f"the {encoding} encoding needs positive `width` and `height` metadata",
        )
    return width, height


def _take(buffer: np.ndarray, expected: int, encoding: str) -> np.ndarray:
    if buffer.size < expected:
        raise RuntimeError(
            f"{encoding} payload is {buffer.size} bytes, "
            f"expected at least {expected} for the declared size",
        )
    return buffer[:expected]