"""Turning a captured BGR frame into the bytes of a chosen encoding."""

import cv2
import numpy as np

# Compressed encodings, and the OpenCV file suffix each is written through.
COMPRESSED = {
    "jpeg": ".jpeg",
    "jpg": ".jpg",
    "jpe": ".jpe",
    "png": ".png",
    "bmp": ".bmp",
    "webp": ".webp",
}

RAW = ("bgr8", "rgb8", "yuv420")

SUPPORTED = (*RAW, *COMPRESSED)


def encode(frame: np.ndarray, encoding: str, jpeg_quality: int) -> np.ndarray:
    """Convert a BGR frame into the bytes of ``encoding``.

    Parameters
    ----------
    frame : np.ndarray
        Captured frame, BGR, as OpenCV reads it.
    encoding : str
        One of ``SUPPORTED``.
    jpeg_quality : int
        Quality 0-100, consulted only for the JPEG encodings.

    Returns
    -------
    np.ndarray
        The bytes to publish. Its shape is not the image shape for a
        compressed encoding, which is why the dimensions are recorded
        separately.

    Raises
    ------
    RuntimeError
        If ``encoding`` is unsupported, or the encoder fails.

    """
    if encoding == "bgr8":
        return frame

    if encoding == "rgb8":
        return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    if encoding == "yuv420":
        return cv2.cvtColor(frame, cv2.COLOR_BGR2YUV_I420)

    if encoding in COMPRESSED:
        params = (
            [cv2.IMWRITE_JPEG_QUALITY, jpeg_quality]
            if encoding in ("jpeg", "jpg", "jpe")
            else []
        )
        ok, encoded = cv2.imencode(COMPRESSED[encoding], frame, params)

        if not ok:
            raise RuntimeError(f"could not encode the frame as {encoding}")

        return encoded

    raise RuntimeError(
        f"unsupported encoding: {encoding}; expected one of {', '.join(SUPPORTED)}",
    )