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

Two payload shapes are accepted, because dora has both in the field. The typed
one is `std/media/v1/Image`: an Arrow struct carrying width, height, encoding
and data together, so the frame is self-describing. The older one is a flat
UInt8 buffer whose dimensions and encoding travel beside it in the message
metadata, which is what a node declaring `std/core/v1/UInt8` sends -- dora-zed
among them. Everything after unpacking is identical.
"""

import io

import cv2
import numpy as np
import pillow_avif  # noqa: F401  # registers the AVIF plugin with Pillow
import pyarrow as pa
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"}


# std/media/v1/Image. A payload carrying these fields describes itself, so its
# message metadata is not consulted. Exported as an Arrow type so a producer --
# the example source, the tests -- builds it from one definition.
IMAGE_TYPE = pa.struct([
    ("width", pa.uint32()),
    ("height", pa.uint32()),
    ("encoding", pa.utf8()),
    ("data", pa.large_binary()),
])
IMAGE_FIELDS = tuple(field.name for field in IMAGE_TYPE)


def to_bgr(storage, metadata: dict) -> np.ndarray:
    """Decode one `image` payload into a writable BGR array.

    Accepts either an `std/media/v1/Image` struct or a flat buffer described by
    the message metadata.
    """
    buffer, encoding, width, height = unpack(storage, metadata)

    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 unpack(storage, metadata: dict) -> tuple[np.ndarray, str, int | None, int | None]:
    """Split a payload of either shape into buffer, encoding, width and height."""
    if is_image_struct(storage):
        # One image per message: the struct is a single-row array.
        row = storage[0].as_py()
        return (
            np.frombuffer(row["data"], dtype=np.uint8),
            str(row["encoding"]).lower(),
            row["width"],
            row["height"],
        )

    return (
        storage.to_numpy().astype(np.uint8),
        str(metadata.get("encoding", "bgr8")).lower(),
        metadata.get("width"),
        metadata.get("height"),
    )


def is_image_struct(storage) -> bool:
    """Whether the payload is an Arrow struct with the std/media/v1/Image fields."""
    fields = getattr(getattr(storage, "type", None), "fields", None)

    if fields is None:
        return False

    return all(name in {field.name for field in fields} for name in IMAGE_FIELDS)


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]