"""Tests that need neither a dataflow nor a display."""

import io

import cv2
import numpy as np
import pytest
from PIL import Image

from opencv_plot import bbox, decode
from opencv_plot.plot import Bboxes, Plot

WIDTH, HEIGHT = 12, 8


class FakeStorage:
    """Stands in for the Arrow array dora hands to the node."""

    def __init__(self, data: np.ndarray) -> None:
        self._data = data

    def to_numpy(self) -> np.ndarray:
        return self._data


@pytest.fixture
def frame() -> np.ndarray:
    rng = np.random.default_rng(seed=0)
    return rng.integers(0, 255, (HEIGHT, WIDTH, 3), dtype=np.uint8)


def test_bgr8_round_trips_unchanged(frame):
    decoded = decode.to_bgr(FakeStorage(frame.ravel()), "bgr8", WIDTH, HEIGHT)

    np.testing.assert_array_equal(decoded, frame)


def test_rgb8_round_trips_unchanged(frame):
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    decoded = decode.to_bgr(FakeStorage(rgb.ravel()), "rgb8", WIDTH, HEIGHT)

    np.testing.assert_array_equal(decoded, frame)


@pytest.mark.parametrize(("suffix", "encoding"), [
    (".png", "png"),
    (".jpg", "jpeg"),
    (".bmp", "bmp"),
    (".webp", "webp"),
])
def test_self_describing_encodings_need_no_dimensions(frame, suffix, encoding):
    ok, buffer = cv2.imencode(suffix, frame)
    assert ok

    decoded = decode.to_bgr(FakeStorage(buffer.ravel()), encoding, None, None)

    assert decoded.shape == (HEIGHT, WIDTH, 3)


def test_yuv420_is_decoded(frame):
    planes = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV_I420)

    decoded = decode.to_bgr(FakeStorage(planes.ravel()), "yuv420", WIDTH, HEIGHT)

    assert decoded.shape == (HEIGHT, WIDTH, 3)


def test_avif_is_decoded(frame):
    buffer = io.BytesIO()
    Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).save(buffer, format="AVIF")
    payload = np.frombuffer(buffer.getvalue(), dtype=np.uint8)

    decoded = decode.to_bgr(FakeStorage(payload), "avif", None, None)

    assert decoded.shape == (HEIGHT, WIDTH, 3)


def test_decoded_frame_is_writable():
    storage = FakeStorage(np.arange(2 * 3 * 3, dtype=np.uint8))

    decoded = decode.to_bgr(storage, "bgr8", width=3, height=2)

    # The frame is annotated in place, so it must not be a read-only view.
    decoded[0, 0] = 0


def test_raw_encodings_need_dimensions():
    storage = FakeStorage(np.zeros(18, dtype=np.uint8))
    with pytest.raises(RuntimeError, match="width"):
        decode.to_bgr(storage, "bgr8", width=None, height=None)


def test_short_payload_is_rejected():
    storage = FakeStorage(np.zeros(4, dtype=np.uint8))
    with pytest.raises(RuntimeError, match="expected at least"):
        decode.to_bgr(storage, "bgr8", width=3, height=2)


def test_unsupported_encoding_is_rejected():
    storage = FakeStorage(np.zeros(4, dtype=np.uint8))
    with pytest.raises(RuntimeError, match="unsupported image encoding"):
        decode.to_bgr(storage, "tiff", width=None, height=None)


def test_xywh_is_converted_to_corners():
    import pyarrow as pa

    value = pa.array(
        [
            {
                "bbox": np.array([10.0, 20.0, 4.0, 6.0]),
                "conf": np.array([0.5]),
                "labels": np.array(["cat"]),
            },
        ],
    )

    boxes = bbox.parse(value, "xywh")

    np.testing.assert_allclose(boxes.corners, [[8.0, 17.0, 12.0, 23.0]])
    assert list(boxes.labels) == ["cat"]


def test_headless_show_draws_without_a_window():
    plot = Plot(headless=True, text="hello")
    plot.bboxes = Bboxes(
        corners=np.array([[1.0, 1.0, 5.0, 5.0]]),
        conf=np.array([0.9]),
        labels=np.array(["cat"]),
    )
    frame = np.zeros((32, 32, 3), dtype=np.uint8)

    assert plot.show(frame) is False
    assert frame.any(), "the overlays should have marked the frame"