from io import BytesIO

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


@pytest.mark.parametrize("loaded_png", [False, True])
@pytest.mark.parametrize("mode_before_pack", [False, True])
def test_premul_byte_reload_requires_straight_for_original_rgb(
    tmp_path, loaded_png, mode_before_pack
):
    """Characterize the reload mismatch blocking automatic PREMUL defaults."""
    rgba = np.array(
        [[51, 102, 204, alpha] for alpha in (0, 64, 128, 255)], dtype=np.uint8
    ).reshape(1, 4, 4)
    expected = rgba.reshape(4, 4) / 255
    path = tmp_path / "source.png"
    Image.fromarray(rgba).save(path)
    if loaded_png:
        image = bpy.data.images.load(str(path), check_existing=False)
    else:
        image = bpy.data.images.new("Premul probe", width=4, height=1, alpha=True)
        image.pixels.foreach_set(expected.astype(np.float32).ravel())
    image.use_fake_user = True
    image.alpha_mode = "STRAIGHT"
    if mode_before_pack:
        image.alpha_mode = "PREMUL"
    image.pack()
    image.alpha_mode = "PREMUL"
    packed = np.asarray(Image.open(BytesIO(image.packed_file.data)))
    np.testing.assert_array_equal(packed, rgba)
    name = image.name
    blend_path = str(tmp_path / "retention.blend")
    bpy.ops.wm.save_as_mainfile(filepath=blend_path)
    bpy.ops.wm.open_mainfile(filepath=blend_path)
    restored = bpy.data.images[name]
    assert restored.alpha_mode == "PREMUL"
    interpreted = np.array(restored.pixels[:]).reshape(4, 4)
    np.testing.assert_allclose(interpreted[0, :3], 1)
    assert not np.allclose(interpreted, expected, atol=1 / 255)
    restored.alpha_mode = "STRAIGHT"
    np.testing.assert_allclose(
        np.array(restored.pixels[:]).reshape(4, 4), expected, atol=1 / 255
    )
    restored.alpha_mode = "PREMUL"
    np.testing.assert_allclose(
        np.array(restored.pixels[:]).reshape(4, 4), interpreted, atol=1 / 255
    )