import sys
import tempfile
from pathlib import Path

import bpy
import numpy as np


PROJECT_ROOT = Path(__file__).parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from snapform.image_depth import (
    MOGE_DEPTH_PLANE_NODE_GROUP_NAME,
    MOGE_MATERIAL_NODE_GROUP_NAME,
    build_moge2_depth_plane,
)


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    image = bpy.data.images.new("MoGe Test", width=64, height=48, alpha=True)
    source = bpy.data.objects.new("MoGe Test", None)
    source.empty_display_type = "IMAGE"
    source.data = image
    bpy.context.collection.objects.link(source)
    depth = np.linspace(1.0, 2.0, 48 * 64, dtype=np.float32).reshape(1, 48, 64)
    normal = np.zeros((1, 48, 64, 3), dtype=np.float32)
    normal[..., 2] = -1.0
    mask = np.ones((1, 48, 64), dtype=bool)
    mask[:, :5] = False
    intrinsics = np.array(
        [[[1.0, 0.0, 0.5], [0.0, 1.0, 0.5], [0.0, 0.0, 1.0]]],
        dtype=np.float32,
    )
    with tempfile.TemporaryDirectory() as temporary:
        path = Path(temporary) / "results.npz"
        np.savez(
            path,
            depth=depth,
            normal=normal,
            mask=mask,
            intrinsics=intrinsics,
        )
        plane = build_moge2_depth_plane(
            bpy.context,
            source,
            path,
            delete_source=True,
        )
    assert plane.modifiers[0].node_group.name == MOGE_DEPTH_PLANE_NODE_GROUP_NAME
    assert MOGE_MATERIAL_NODE_GROUP_NAME in bpy.data.node_groups
    geometry_group = plane.modifiers[0].node_group
    inputs = {
        item.name
        for item in geometry_group.interface.items_tree
        if getattr(item, "in_out", None) == "INPUT"
    }
    assert "Displacement" in inputs
    assert "Background Depth" in inputs
    assert "Mask Image" not in inputs
    assert "Use Mask" not in inputs
    sockets = {
        item.name: item
        for item in geometry_group.interface.items_tree
        if getattr(item, "in_out", None) == "INPUT"
    }
    assert sockets["Subdivision"].default_value == 6
    assert sockets["Depth Scale"].structure_type == "SINGLE"
    assert sockets["Background Depth"].structure_type == "SINGLE"
    assert sockets["Background Depth"].subtype == "DISTANCE"
    assert sockets["Background Depth"].default_value == 1.0
    material_group = bpy.data.node_groups[MOGE_MATERIAL_NODE_GROUP_NAME]
    assert not any(node.bl_idname == "ShaderNodeNormalMap" for node in material_group.nodes)
    assert any(node.bl_idname == "ShaderNodeBump" for node in material_group.nodes)
    material = plane.data.materials[0]
    image_textures = [
        node
        for node in material.node_tree.nodes
        if node.bl_idname == "ShaderNodeTexImage"
    ]
    assert image_textures
    assert all(node.extension == "EXTEND" for node in image_textures)
    assert any(node.image.name == "MoGe Test.depth" for node in image_textures)
    geometry_textures = [
        node
        for node in geometry_group.nodes
        if node.bl_idname == "GeometryNodeImageTexture"
    ]
    assert geometry_textures
    assert all(node.extension == "EXTEND" for node in geometry_textures)
    print("MoGe-2 depth plane validation passed")


if __name__ == "__main__":
    main()