"""
Tests for _plotly_utils.data_utils.image_array_to_data_uri.

Uses the pypng backend so the tests do not require Pillow. The generated
data URIs are decoded with the vendored pypng Reader so the tests assert on
the encoded pixels rather than only on the data URI prefix.
"""

import base64

import numpy as np
import pytest

from _plotly_utils.data_utils import image_array_to_data_uri
from _plotly_utils.png import Reader

PNG_PREFIX = "data:image/png;base64,"


def decode_data_uri(uri):
    """Decode a PNG data URI back into an array plus the PNG header info."""
    assert uri.startswith(PNG_PREFIX)
    png_bytes = base64.b64decode(uri[len(PNG_PREFIX) :])
    width, height, rows, info = Reader(bytes=png_bytes).read_flat()
    if info["greyscale"]:
        channels = 1
    elif info["alpha"]:
        channels = 4
    else:
        channels = 3
    decoded = np.array(list(rows), dtype=np.uint8).reshape(height, width, channels)
    if channels == 1:
        decoded = decoded[:, :, 0]
    return decoded, info


def test_greyscale_array_round_trips():
    img = np.array([[0, 255], [128, 64]], dtype=np.uint8)

    decoded, info = decode_data_uri(image_array_to_data_uri(img, backend="pypng"))

    assert info["greyscale"] is True
    assert info["alpha"] is False
    np.testing.assert_array_equal(decoded, img)


def test_rgb_array_round_trips():
    img = np.array(
        [[[255, 0, 0], [0, 255, 0]], [[0, 0, 255], [10, 20, 30]]], dtype=np.uint8
    )

    decoded, info = decode_data_uri(image_array_to_data_uri(img, backend="pypng"))

    assert info["greyscale"] is False
    assert info["alpha"] is False
    np.testing.assert_array_equal(decoded, img)


def test_rgba_array_round_trips():
    img = np.array(
        [[[255, 0, 0, 255], [0, 255, 0, 128]], [[0, 0, 255, 0], [10, 20, 30, 40]]],
        dtype=np.uint8,
    )

    decoded, info = decode_data_uri(image_array_to_data_uri(img, backend="pypng"))

    assert info["greyscale"] is False
    assert info["alpha"] is True
    np.testing.assert_array_equal(decoded, img)


@pytest.mark.parametrize("compression", [-1, 10])
def test_invalid_compression_raises(compression):
    img = np.zeros((2, 2), dtype=np.uint8)
    with pytest.raises(ValueError, match="compression level"):
        image_array_to_data_uri(img, backend="pypng", compression=compression)


def test_invalid_shape_raises():
    img = np.zeros(5, dtype=np.uint8)
    with pytest.raises(ValueError, match="Invalid image shape"):
        image_array_to_data_uri(img, backend="pypng")


def test_jpg_without_pil_backend_raises():
    img = np.zeros((2, 2), dtype=np.uint8)
    with pytest.raises(ValueError, match="jpg binary strings"):
        image_array_to_data_uri(img, backend="pypng", ext="jpg")