import sys
from pathlib import Path

import bpy

PROJECT_ROOT = Path(__file__).parents[2]
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(
    0,
    str(next((PROJECT_ROOT / "src" / "anyimage" / "wheels").glob("blendjob-*.whl"))),
)

from anyimage.operators.common.mesh_result import (
    modifier_input_identifier,
    set_modifier_input,
)
from anyimage.operators.cutout_tool import object as cutout_objects
from anyimage.operators.cutout_tool import shape as shapes


def evaluated_mesh(obj):
    evaluated = obj.evaluated_get(bpy.context.evaluated_depsgraph_get())
    result = evaluated.to_mesh()
    return evaluated, result


def input_names(group):
    return [
        item.name
        for item in group.interface.items_tree
        if item.item_type == "SOCKET" and item.in_out == "INPUT"
    ]


def grid_object(name):
    vertices = [(0.0, float(y), float(z)) for z in range(3) for y in range(3)]
    faces = []
    for z in range(2):
        for y in range(2):
            a = z * 3 + y
            faces.extend(((a, a + 4, a + 1), (a, a + 3, a + 4)))
    data = bpy.data.meshes.new(name)
    data.from_pydata(vertices, (), faces)
    uv_layer = data.uv_layers.new(name="UVMap")
    for loop in data.loops:
        coordinate = data.vertices[loop.vertex_index].co
        uv_layer.data[loop.index].uv = (coordinate.y * 0.5, coordinate.z * 0.5)
    obj = bpy.data.objects.new(name, data)
    bpy.context.scene.collection.objects.link(obj)
    return obj


def add_modifier(obj, shape, thickness):
    group = cutout_objects.load_cutout_node_group()
    modifier = obj.modifiers.new(group.name, "NODES")
    modifier.node_group = group
    for name, value in (
        ("Shape", shapes.CUTOUT_SHAPE_INDICES[shape]),
        ("Cleanup", 0.0),
        ("Smooth", 0.0),
        ("Smooth Iterations", 0),
        ("Thickness", thickness),
    ):
        set_modifier_input(
            modifier,
            modifier_input_identifier(group, name),
            value,
        )
    return modifier


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    wrapper = cutout_objects.load_cutout_node_group()
    assert all(
        cutout_objects.load_cutout_node_group() == wrapper
        for shape in shapes.CUTOUT_SHAPES
    )
    assert wrapper.name == shapes.CUTOUT_NODE_GROUP_NAME
    assert wrapper.is_modifier
    groups = {
        shape: bpy.data.node_groups[name]
        for shape, name in shapes.CUTOUT_NODE_GROUP_NAMES.items()
    }
    assert all(not group.is_modifier for group in groups.values())

    menu_switch = next(
        node for node in wrapper.nodes if node.bl_idname == "GeometryNodeMenuSwitch"
    )
    assert [item.name for item in menu_switch.enum_items] == list(
        shapes.CUTOUT_SHAPE_LABELS.values()
    )
    assert any(
        link.from_socket.name == "Shape"
        and link.to_node == menu_switch
        and link.to_socket.name == "Menu"
        for link in wrapper.links
    )
    child_nodes = {
        node.node_tree: node
        for node in wrapper.nodes
        if node.bl_idname == "GeometryNodeGroup"
    }
    assert set(child_nodes) == set(groups.values())
    for shape, group in groups.items():
        assert any(
            link.from_node == child_nodes[group]
            and link.to_node == menu_switch
            and link.to_socket.name == shapes.CUTOUT_SHAPE_LABELS[shape]
            for link in wrapper.links
        )

    assert input_names(wrapper) == [
        "Geometry",
        "Shape",
        "Thickness",
        "Depth Scale",
        "Double Sided",
        "Smooth Iterations",
        "Cleanup",
        "Smooth",
        "Base Plane Depth",
        "Depth Direction",
        "Uniform Scale",
        "Depth Image",
    ]
    expected_inputs = {
        "SURFACE": ["Geometry", "Thickness"],
        "BALLOON": ["Geometry", "Thickness"],
        "DEPTH_BALLOON": [
            "Geometry",
            "Thickness",
            "Depth Scale",
            "Double Sided",
            "Base Plane Depth",
            "Depth Direction",
            "Uniform Scale",
            "Depth Image",
        ],
        "DEPTH_SURFACE": [
            "Geometry",
            "Thickness",
            "Depth Scale",
            "Base Plane Depth",
            "Uniform Scale",
            "Depth Image",
        ],
    }
    for shape, group in groups.items():
        assert input_names(group) == expected_inputs[shape]

    dependencies = {}
    for link in wrapper.links:
        if link.from_node.bl_idname != "NodeGroupInput":
            continue
        if link.to_node.bl_idname != "GeometryNodeGroup":
            continue
        shape = next(
            shape for shape, group in groups.items() if group == link.to_node.node_tree
        )
        dependencies.setdefault(link.from_socket.name, set()).add(shape)
    assert dependencies == {
        "Thickness": set(shapes.CUTOUT_SHAPES),
        "Depth Scale": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
        "Double Sided": {"DEPTH_BALLOON"},
        "Base Plane Depth": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
        "Depth Direction": {"DEPTH_BALLOON"},
        "Uniform Scale": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
        "Depth Image": {"DEPTH_BALLOON", "DEPTH_SURFACE"},
    }

    repeat_counts = {
        shape: sum(node.bl_idname == "GeometryNodeRepeatInput" for node in group.nodes)
        for shape, group in groups.items()
    }
    assert repeat_counts == {
        "SURFACE": 1,
        "BALLOON": 0,
        "DEPTH_BALLOON": 0,
        "DEPTH_SURFACE": 1,
    }
    assert (
        sum(node.bl_idname == "GeometryNodeRepeatInput" for node in wrapper.nodes) == 1
    )
    assert (
        sum(node.bl_idname == "GeometryNodeInputMeshIsland" for node in wrapper.nodes)
        == 1
    )
    assert not any(
        node.bl_idname == "GeometryNodeAttributeStatistic" for node in wrapper.nodes
    )
    assert all(
        not any(node.bl_idname == "GeometryNodeInputMeshIsland" for node in group.nodes)
        for group in groups.values()
    )
    assert any(node.bl_idname == "GeometryNodeBlurAttribute" for node in wrapper.nodes)
    smooth = next(
        item
        for item in wrapper.interface.items_tree
        if item.item_type == "SOCKET" and item.name == "Smooth"
    )
    iterations = next(
        item
        for item in wrapper.interface.items_tree
        if item.item_type == "SOCKET" and item.name == "Smooth Iterations"
    )
    cleanup = next(
        item
        for item in wrapper.interface.items_tree
        if item.item_type == "SOCKET" and item.name == "Cleanup"
    )
    assert cleanup.default_value == 0.0
    assert cleanup.subtype == "NONE"
    cleanup_scale = next(
        link.to_node
        for link in wrapper.links
        if link.from_socket.name == "Cleanup"
        and link.from_node.bl_idname == "NodeGroupInput"
    )
    assert cleanup_scale.bl_idname == "ShaderNodeMath"
    assert cleanup_scale.operation == "MULTIPLY"
    assert abs(cleanup_scale.inputs[1].default_value - 0.01) < 1e-7
    assert smooth.default_value == 1.0
    assert iterations.default_value == 0
    assert iterations.min_value == 0
    assert getattr(iterations.parent, "name", "") == ""
    assert all(
        not any(node.bl_idname == "GeometryNodeBlurAttribute" for node in group.nodes)
        for group in groups.values()
    )

    for group in (wrapper, *groups.values()):
        assert all(not node.label for node in group.nodes)
        assert all(
            link.from_node.location.x <= link.to_node.location.x for link in group.links
        )
        group_nodes = list(group.nodes)
        assert len({tuple(node.location) for node in group_nodes}) == len(group_nodes)
        for index, node in enumerate(group_nodes):
            for other in group_nodes[index + 1 :]:
                assert not (
                    abs(node.location.x - other.location.x) < 150.0
                    and abs(node.location.y - other.location.y) < 100.0
                )
        for group_input in (
            node for node in group_nodes if node.bl_idname == "NodeGroupInput"
        ):
            assert all(
                output.is_linked or output.hide
                for output in group_input.outputs
                if output.name
            )
        for item in group.interface.items_tree:
            if item.item_type == "SOCKET" and item.in_out == "INPUT":
                assert item.structure_type == "SINGLE"
    surface = grid_object("Surface Validation")
    add_modifier(surface, "SURFACE", 2.4)
    evaluated, result = evaluated_mesh(surface)
    try:
        assert len(result.vertices) == 34
        assert len(result.polygons) == 40
        side_faces = [
            polygon for polygon in result.polygons if abs(polygon.normal.x) < 0.1
        ]
        assert len(side_faces) == 24
    finally:
        evaluated.to_mesh_clear()

    balloon = grid_object("Balloon Validation")
    profile = balloon.data.attributes.new(
        shapes.BALLOON_ATTRIBUTE_NAME, "FLOAT", "POINT"
    )
    profile.data.foreach_set(
        "value",
        (0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0),
    )
    modifier = add_modifier(balloon, "BALLOON", 1.0)
    evaluated, result = evaluated_mesh(balloon)
    try:
        initial_faces = len(result.polygons)
        initial_extent = max(abs(vertex.co.x) for vertex in result.vertices)
    finally:
        evaluated.to_mesh_clear()
    set_modifier_input(
        modifier,
        modifier_input_identifier(modifier.node_group, "Thickness"),
        2.0,
    )
    balloon.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    evaluated, result = evaluated_mesh(balloon)
    try:
        assert len(result.polygons) == initial_faces
        extent = max(abs(vertex.co.x) for vertex in result.vertices)
        assert abs(extent - initial_extent * 2.0) < 1e-5
    finally:
        evaluated.to_mesh_clear()

    depth_balloon = grid_object("Depth Balloon Metric Validation")
    profile = depth_balloon.data.attributes.new(
        shapes.BALLOON_ATTRIBUTE_NAME, "FLOAT", "POINT"
    )
    profile.data.foreach_set("value", (1.0,) * 9)
    modifier = add_modifier(depth_balloon, "DEPTH_BALLOON", 0.0)
    point_image = bpy.data.images.new(
        "Depth Balloon Points",
        width=3,
        height=3,
        alpha=True,
        float_buffer=True,
    )
    point_image.colorspace_settings.name = "Non-Color"
    point_image.alpha_mode = "CHANNEL_PACKED"
    pixels = []
    for z in range(3):
        for y in range(3):
            if (y, z) == (1, 1):
                depth = 1.5
            elif (y, z) == (0, 0):
                depth = 2.5
            else:
                depth = 2.0
            pixels.extend((0.0, 0.0, depth, 0.0))
    point_image.pixels.foreach_set(pixels)
    for name, value in (
        ("Depth Image", point_image),
        ("Depth Scale", 1.0),
        ("Uniform Scale", 1.0),
        ("Base Plane Depth", 2.0),
        ("Depth Direction", (0.0, 0.0, 1.0)),
    ):
        set_modifier_input(
            modifier,
            modifier_input_identifier(modifier.node_group, name),
            value,
        )
    depth_balloon.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    evaluated, result = evaluated_mesh(depth_balloon)
    try:
        x_values = [vertex.co.x for vertex in result.vertices]
        assert abs(min(x_values) + 0.5) < 1e-5
        assert abs(max(x_values)) < 1e-5
    finally:
        evaluated.to_mesh_clear()

    pixels[(1 * 3 + 1) * 4 + 3] = -1.0
    point_image.pixels.foreach_set(pixels)
    point_image.update()
    depth_balloon.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    evaluated, result = evaluated_mesh(depth_balloon)
    try:
        x_values = [vertex.co.x for vertex in result.vertices]
        assert abs(min(x_values) + 0.5) < 1e-5
        assert abs(max(x_values)) < 1e-5
    finally:
        evaluated.to_mesh_clear()

    depth_surface = grid_object("Depth Surface Projection Validation")
    modifier = add_modifier(depth_surface, "DEPTH_SURFACE", 0.0)
    point_image = bpy.data.images.new(
        "Depth Surface Points",
        width=3,
        height=3,
        alpha=True,
        float_buffer=True,
    )
    point_image.colorspace_settings.name = "Non-Color"
    point_image.alpha_mode = "CHANNEL_PACKED"
    pixels = []
    for z in range(3):
        for y in range(3):
            pixels.extend((-y * 0.5, -z * 0.5, 2.0 + y * 0.25, 0.0))
    point_image.pixels.foreach_set(pixels)
    for name, value in (
        ("Depth Image", point_image),
        ("Depth Scale", 1.0),
        ("Uniform Scale", 1.0),
        ("Base Plane Depth", 2.0),
        ("Smooth Iterations", 0),
    ):
        set_modifier_input(
            modifier,
            modifier_input_identifier(modifier.node_group, name),
            value,
        )
    depth_surface.update_tag(refresh={"DATA"})
    bpy.context.view_layer.update()
    evaluated, result = evaluated_mesh(depth_surface)
    try:
        coordinates = [tuple(vertex.co) for vertex in result.vertices]
        assert all(
            all(
                abs(actual - expected) < 1e-5 for actual, expected in zip(point, target)
            )
            for point, target in zip(
                coordinates,
                ((y * 0.25, y * 0.5, z * 0.5) for z in range(3) for y in range(3)),
            )
        )
    finally:
        evaluated.to_mesh_clear()

    print(
        f"Validated Blender {bpy.app.version_string}: "
        f"{wrapper.name} wrapping "
        f"{', '.join(group.name for group in groups.values())}"
    )


if __name__ == "__main__":
    main()