"""test gpu interop nodes"""

import pathlib

import omni.graph.core as og
import omni.kit.commands
import omni.kit.stage_templates
import omni.kit.test
import omni.usd
from omni.rtx.tests import RtxTest, postLoadTestSettings, testSettings
from omni.rtx.tests.test_common import wait_for_update
from pxr import UsdGeom

EXTENSION_FOLDER_PATH = pathlib.Path(
    omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
)
GOLDEN_DIR = EXTENSION_FOLDER_PATH.joinpath("omni/graph/rtxtest/tests/data/golden")


# ======================================================================
class TestPreRender(RtxTest):
    async def setUp(self):
        await super().setUp()
        self.set_settings(testSettings)
        omni.usd.get_context().new_stage()
        self.add_dir_light()
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

    async def tearDown(self):
        await super().tearDown()
        await omni.kit.stage_templates.new_stage_async()

    def create_mesh(self, name):
        box = UsdGeom.Mesh.Define(self.ctx.get_stage(), name)
        box.CreatePointsAttr(
            [
                (-50, -50, -50),
                (50, -50, -50),
                (-50, -50, 50),
                (50, -50, 50),
                (-50, 50, -50),
                (50, 50, -50),
                (50, 50, 50),
                (-50, 50, 50),
            ]
        )
        box.CreateFaceVertexCountsAttr([4, 4, 4, 4, 4, 4])
        box.CreateFaceVertexIndicesAttr([0, 1, 3, 2, 0, 4, 5, 1, 1, 5, 6, 3, 2, 3, 6, 7, 0, 2, 7, 4, 4, 7, 6, 5])
        box.CreateSubdivisionSchemeAttr("none")
        return box

    # ----------------------------------------------------------------------
    async def test_prerender_gpuinterop(self):
        controller = og.Controller()
        keys = og.Controller.Keys()
        (_, _, _, _) = controller.edit(
            {"graph_path": "/PushGraph", "pipeline_stage": og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_PRERENDER},
            {
                keys.CREATE_NODES: [
                    ("ConstantString", "omni.graph.nodes.ConstantString"),
                    ("ToToken", "omni.graph.nodes.ToToken"),
                    ("ReadPoints", "omni.graph.nodes.ReadPrimAttribute"),
                    ("Entry", "omni.graph.nodes.RenderPreProcessEntry"),
                    ("Allocator", "omni.graph.nodes.RpResourceExampleAllocator"),
                    ("Deformer", "omni.graph.nodes.RpResourceExampleDeformer"),
                    ("ToHydra", "omni.graph.nodes.RpResourceExampleHydra"),
                ],
                keys.CONNECT: [
                    ("ConstantString.inputs:value", "ToToken.inputs:value"),
                    ("ToToken.outputs:converted", "ReadPoints.inputs:primPath"),
                    ("ToToken.outputs:converted", "Allocator.inputs:primPath"),
                    ("ReadPoints.outputs:value", "Allocator.inputs:points"),
                    ("Entry.outputs:stream", "Allocator.inputs:stream"),
                    ("Allocator.outputs:pointCountCollection", "Deformer.inputs:pointCountCollection"),
                    ("Allocator.outputs:primPathCollection", "Deformer.inputs:primPathCollection"),
                    ("Allocator.outputs:resourcePointerCollection", "Deformer.inputs:resourcePointerCollection"),
                    ("Allocator.outputs:stream", "Deformer.inputs:stream"),
                    ("Deformer.outputs:pointCountCollection", "ToHydra.inputs:pointCountCollection"),
                    ("Deformer.outputs:primPathCollection", "ToHydra.inputs:primPathCollection"),
                    ("Deformer.outputs:resourcePointerCollection", "ToHydra.inputs:resourcePointerCollection"),
                ],
                keys.SET_VALUES: [
                    ("ConstantString.inputs:value", "/World/box"),
                    ("ReadPoints.inputs:usePath", True),
                    ("ReadPoints.inputs:name", "points"),
                    ("Deformer.inputs:deformScale", 50.0),
                    ("Deformer.inputs:positionScale", 100.0),
                    ("ToHydra.inputs:sendToHydra", True),
                ],
            },
        )

        self.create_mesh("/World/box")

        await wait_for_update(self.ctx, 30)

        await self.screenshot_and_diff(
            GOLDEN_DIR,
            output_subdir="omni.graph.rtxtest",
            golden_img_name="gpuinterop_prerender_deformer.png",
            threshold=1e-4,
        )
