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

import io

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

from dora_image_view import bbox, decode, depth
from dora_image_view.view import Bboxes, View

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():
    view = View(headless=True, text="hello")
    view.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)
    view.panels["image"] = frame

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


def test_show_with_no_panels_is_a_no_op():
    assert View(headless=True).show() is False


def test_panels_are_tiled_left_to_right_at_a_common_height():
    view = View(headless=True)
    view.panels["image"] = np.zeros((20, 10, 3), dtype=np.uint8)
    # Half the height and half the width: scaling to the common height should
    # double both, so it contributes 20 columns, not 10.
    view.panels["depth"] = np.zeros((10, 10, 3), dtype=np.uint8)

    composite = view.compose()

    assert composite.shape[0] == 20
    assert composite.shape[1] == 10 + 20


def test_depth_decodes_millimetres_and_reshapes():
    values = np.array([100.0, 200.0, 300.0, 400.0], dtype=np.float32)

    decoded = depth.to_millimetres(FakeStorage(values), {"shape": [2, 2]})

    np.testing.assert_allclose(decoded, [[100.0, 200.0], [300.0, 400.0]])


def test_depth_reinterprets_a_float32_payload_sent_as_bytes():
    values = np.array([1500.0, 2500.0], dtype=np.float32)

    decoded = depth.to_millimetres(
        FakeStorage(values.view(np.uint8)),
        {"shape": [1, 2], "dtype": "float32"},
    )

    np.testing.assert_allclose(decoded, [[1500.0, 2500.0]])


def test_depth_converts_metres_to_millimetres():
    decoded = depth.to_millimetres(
        FakeStorage(np.array([1.5, 2.0], dtype=np.float32)),
        {"shape": [1, 2], "units": "m"},
    )

    np.testing.assert_allclose(decoded, [[1500.0, 2000.0]])


def test_depth_normalises_non_finite_values_to_missing():
    values = np.array([np.inf, np.nan, 500.0, 0.0], dtype=np.float32)

    decoded = depth.to_millimetres(FakeStorage(values), {"shape": [2, 2]})

    np.testing.assert_allclose(decoded, [[0.0, 0.0], [500.0, 0.0]])


def test_depth_without_a_shape_is_an_error():
    with pytest.raises(RuntimeError, match="shape"):
        depth.to_millimetres(FakeStorage(np.zeros(4, dtype=np.float32)), {})


def test_colorize_leaves_missing_depth_black_and_colours_the_rest():
    depth_mm = np.array([[0.0, 500.0], [1000.0, 1500.0]], dtype=np.float32)

    frame = depth.colorize(depth_mm)

    assert frame.shape == (2, 2, 3)
    np.testing.assert_array_equal(frame[0, 0], [0, 0, 0])
    assert frame[1, 1].any(), "valid depth should be coloured"


def test_depth_range_is_pinned_by_the_environment(monkeypatch):
    monkeypatch.setenv("VIEW_DEPTH_MIN_MM", "200")
    monkeypatch.setenv("VIEW_DEPTH_MAX_MM", "600")

    assert depth.resolve_range(np.array([[1000.0]], dtype=np.float32)) == (200.0, 600.0)


def test_one_pinned_end_leaves_the_other_auto_ranged(monkeypatch):
    monkeypatch.setenv("VIEW_DEPTH_MIN_MM", "0")

    low, high = depth.resolve_range(np.array([[100.0, 900.0]], dtype=np.float32))

    assert low == 0.0
    assert high > 0.0


def test_depth_range_of_an_empty_frame_does_not_divide_by_zero():
    low, high = depth.resolve_range(np.zeros((2, 2), dtype=np.float32))

    assert high > low