from __future__ import annotations

import json
import sys
from pathlib import Path

import bpy


BLENDER_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BLENDER_ROOT))

import material_id_baker  # noqa: E402


def make_test_object() -> bpy.types.Object:
    mesh = bpy.data.meshes.new("MaterialIDSmokeMesh")
    mesh.from_pydata(
        [
            (-1.0, -1.0, 0.0),
            (0.0, -1.0, 0.0),
            (0.0, 1.0, 0.0),
            (-1.0, 1.0, 0.0),
            (0.0, -1.0, 0.0),
            (1.0, -1.0, 0.0),
            (1.0, 1.0, 0.0),
            (0.0, 1.0, 0.0),
        ],
        [],
        [(0, 1, 2, 3), (4, 5, 6, 7)],
    )
    mesh.update()

    left = bpy.data.materials.new("Left Material")
    right = bpy.data.materials.new("Right Material")
    mesh.materials.append(left)
    mesh.materials.append(right)
    mesh.polygons[0].material_index = 0
    mesh.polygons[1].material_index = 1

    uv_layer = mesh.uv_layers.new(name="ID UV")
    uv_coordinates = (
        (0.05, 0.05),
        (0.45, 0.05),
        (0.45, 0.95),
        (0.05, 0.95),
        (0.55, 0.05),
        (0.95, 0.05),
        (0.95, 0.95),
        (0.55, 0.95),
    )
    for loop, uv in zip(mesh.loops, uv_coordinates, strict=True):
        uv_layer.uv[loop.index].vector = uv

    obj = bpy.data.objects.new("MaterialIDSmokeObject", mesh)
    bpy.context.scene.collection.objects.link(obj)
    obj.modifiers.new(name="Triangulate for evaluated-mesh test", type="TRIANGULATE")
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    return obj


def pixel(image: bpy.types.Image, x: int, y: int) -> tuple[float, float, float, float]:
    offset = (y * image.size[0] + x) * 4
    return tuple(image.pixels[offset : offset + 4])


def close_rgb(actual: tuple[float, ...], expected: list[float], tolerance: float = 0.03) -> bool:
    return all(abs(actual[index] - expected[index]) <= tolerance for index in range(3))


def main() -> None:
    material_id_baker.register()
    source = make_test_object()
    original_engine = bpy.context.scene.render.engine

    settings = bpy.context.scene.material_id_baker
    settings.resolution = "256"
    settings.margin = 4
    settings.palette_mode = "DISTINCT"
    settings.apply_modifiers = True
    settings.image_name = "Material_ID_Smoke"
    settings.save_to_disk = True
    settings.filepath = "/tmp/material_id_baker_smoke.png"
    settings.write_legend = True

    result = bpy.ops.object.bake_material_id()
    assert result == {"FINISHED"}, result
    assert bpy.context.active_object == source
    assert source.select_get()
    assert bpy.context.scene.render.engine == original_engine
    assert not any(item.name.startswith("__MID_BAKER_") for item in bpy.data.objects)
    assert not any(item.name.startswith("__MID_BAKER_") for item in bpy.data.materials)

    image = bpy.data.images[settings.last_image_name]
    palette = json.loads(image["material_id_palette"])
    assert len(palette) == 2
    assert palette[0]["material"] == "Left Material"
    assert palette[1]["material"] == "Right Material"
    left_pixel = pixel(image, 64, 128)
    right_pixel = pixel(image, 192, 128)
    assert close_rgb(left_pixel, palette[0]["rgb"]), (left_pixel, palette[0])
    assert close_rgb(right_pixel, palette[1]["rgb"]), (right_pixel, palette[1])
    assert Path("/tmp/material_id_baker_smoke.png").is_file()
    assert Path("/tmp/material_id_baker_smoke.json").is_file()

    print("MATERIAL_ID_BAKER_SMOKE_TEST_OK")


if __name__ == "__main__":
    main()
