"""Tests for the post render graph scaffolding APIs"""
import carb
import inspect
import pathlib
import omni.graph.core as og
import omni.graph.core.tests as ogts
import omni.hydratexture
from omni.kit.viewport.utility import get_active_viewport
import omni.kit.test
from omni.kit.test_helpers_gfx.compare_utils import finalize_capture_and_compare, ComparisonMetric
from omni.rtx.tests import RtxTest, postLoadTestSettings, testSettings, OUTPUTS_DIR
from omni.rtx.tests.test_common import wait_for_update
from pxr import Gf, Sdf, Usd, 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")
THRESHOLD = 1e-4


class TestPostRenderGraphs(RtxTest):

    @staticmethod
    def get_aov_list(stage: Usd.Stage, rp_prim_path) -> list[str]:
        aovs : list[str] = []
        render_vars = stage.GetPrimAtPath(rp_prim_path).GetRelationship("orderedVars").GetForwardedTargets()
        for render_var in render_vars:
            aov_name = stage.GetPrimAtPath(render_var).GetAttribute("sourceName").Get()
            aovs.append(aov_name)
        return aovs

    @property
    def __test_name(self) -> str:
        """
        The full name of the test.

        It has the name of the module, class and the current test function. We use the stack to get the name of the test
        function and since it's only called from create_test_window and finalize_test, we get the third member.
        """
        return f"{self.__module__}.{self.__class__.__name__}.{inspect.stack()[2][3]}"

    def render_product_path(self, hydra_texture) -> str:
        '''Return a string to the UsdRender.Product used by the texture'''
        render_product = hydra_texture.get_render_product_path()
        if render_product and (not render_product.startswith('/')):
            render_product = '/Render/RenderProduct_' + render_product
        return render_product

    def compare_output_image(self,
                             golden_img_name,
                             threshold,
                             output_img_dir: pathlib.Path,
                             golden_img_dir: pathlib.Path,
                             metric: ComparisonMetric = ComparisonMetric.MEAN_ERROR_SQUARED):
        """Compare an arbitrary image with a golden image. The image generated in the test could be produced by OgnGpuInteropCpuToDisk or some other method."""

        diff = finalize_capture_and_compare(golden_img_name, threshold, output_img_dir, golden_img_dir, metric=metric)

        if diff != 0:
            carb.log_warn(f"[{self.__test_name}] the generated image {golden_img_name} has max difference {diff}")

        if (diff is not None) and diff >= threshold:
            self._failedImages.append(golden_img_name)
        return diff

    async def setUpProps(self, stage):
        # add a couple of prims in the scene
        cube = ogts.create_cube(stage, "Cube", (255, 0, 0))
        sphere = ogts.create_sphere(stage, "Sphere", (0, 255, 0))

        xform_cube = UsdGeom.Xformable(cube)
        xform_cube.AddTranslateOp().Set((0, 30, 40))
        xform_cube.AddScaleOp().Set(Gf.Vec3f(30.0))
        xform_sphere = UsdGeom.Xformable(sphere)
        xform_sphere.AddTranslateOp().Set((0, 30, -40))
        xform_sphere.AddScaleOp().Set(Gf.Vec3f(25.0))

        await wait_for_update(self.ctx, 10)

    async def setUp(self):
        await super().setUp()
        # must apply the test settings to the renderer to ensure that the test setup matches the SDG setup
        self.set_settings(testSettings)
        # set up the stage
        omni.usd.get_context().new_stage()
        self.add_dir_light()
        self.add_floor()
        await omni.kit.app.get_app().next_update_async()
        self.set_settings(postLoadTestSettings)

        # renderer
        renderer = "rtx"
        if renderer not in self.ctx.get_attached_hydra_engine_names():
            omni.usd.add_hydra_engine(renderer, self.ctx)

        self._stage = self.ctx.get_stage()

        session_layer = self._stage.GetSessionLayer()
        with Usd.EditContext(self._stage, session_layer):
            # set the render product on the active viewport
            viewport_api = get_active_viewport()
            if viewport_api:
                self._render_product_path_0 = viewport_api.render_product_path

            # create the post render graph
            pipeline_stage = og.GraphPipelineStage.GRAPH_PIPELINE_STAGE_POSTRENDER
            execution_model = "push"

            orchestration_graphs = og.get_global_orchestration_graphs_in_pipeline_stage(pipeline_stage)
            top_graph = orchestration_graphs[0]

            self._og_graph_path = f"{self._render_product_path_0}/TestPostRender"
            compound_node1 = top_graph.create_graph_as_node(
                "TestPostRender",
                self._og_graph_path,
                evaluator=execution_model,
                is_global_graph=True,
                is_backed_by_usd=True,
                backing_type=og.GraphBackingType.GRAPH_BACKING_TYPE_FABRIC_SHARED,
                pipeline_stage=pipeline_stage)
            self._post_render_graph = compound_node1.get_wrapped_graph()

            # set the ogPostProcessPath attribute to associate the graph with this render product
            attr_name = "ogPostProcessPath"
            self._stage.GetPrimAtPath(self._render_product_path_0).CreateAttribute(attr_name, Sdf.ValueTypeNames.String).Set(self._og_graph_path)

            # debug sanity check
            # carb.log_info(f"!!! rp_path:{self._render_product_path_0}")
            # carb.log_info(f"!!! og_path:{self._og_graph_path}")
            # carb.log_info(f"!!! rp_aovs: {self.get_aov_list(self._stage, self._render_product_path_0)}")
            # carb.log_info(f"!!! rp.ogPostProcessPath={self._stage.GetPrimAtPath(self._render_product_path_0).GetAttribute(attr_name).Get()}")
            # carb.log_info(f"!!! og.nodes={self._post_render_graph.get_nodes()}")

    async def tearDown(self):
        await super().tearDown()

    # -----------------------------------------------------------------------------
    async def test_modify_aov(self):
        """Validate that the post render graph can modify an AOV in place."""

        await self.setUpProps(self._stage)

        await self.screenshot_and_diff(
            GOLDEN_DIR,
            output_subdir="omni.graph.rtxtest",
            golden_img_name="test_modify_aov.initial.png",
            threshold=THRESHOLD)

        keys = og.Controller.Keys
        (_, (_, test_node), _, _) = og.Controller.edit(
            self._post_render_graph,
            {
                keys.CREATE_NODES: [
                    ("RenderProductEntry", "omni.graph.nodes.GpuInteropRenderProductEntry"),
                    ("TestNode", "omni.graph.rtxtest.ScaleAOVTextureInPlace"),
                ],
                keys.CONNECT: [
                    ("RenderProductEntry.outputs:gpu", "TestNode.inputs:gpu"),
                    ("RenderProductEntry.outputs:rp", "TestNode.inputs:rp"),
                    ("RenderProductEntry.outputs:exec", "TestNode.inputs:exec"),
                ],
                keys.SET_VALUES: [
                    ("TestNode.inputs:inputAOV", "LdrColor"),
                    ("TestNode.inputs:multiplier", 1.5),
                    ("TestNode.inputs:width", self.WINDOW_SIZE[0]),
                    ("TestNode.inputs:height", self.WINDOW_SIZE[1]),
                ]
            },
        )

        await wait_for_update(self.ctx, 10)

        # Validate that the outputAOVPtr attribute is set.
        # This attribute is set after getting from the RenderProduct
        # the pointer of the AOV with the token set on the "inputAOV" attribute.
        actual = og.Controller.get(og.Controller.attribute("outputs:outputAOVPtr", test_node))
        self.assertNotEqual(0, actual)

        # validate that the scheduleCudaTask() method returned true
        schedule_result = og.Controller.get(og.Controller.attribute("outputs:scheduleResult", test_node))
        self.assertEqual(True, schedule_result)

        await self.screenshot_and_diff(
            GOLDEN_DIR,
            output_subdir="omni.graph.rtxtest",
            golden_img_name="test_modify_aov.scaled.png",
            threshold=THRESHOLD)

    # -----------------------------------------------------------------------------
    async def test_create_new_aov(self):
        """Validate that the post render graph can create new AOVs without throwing."""

        # This test allocates a new AOV, saves it to disk and compares it with a golden image
        # NOTE: the image produced by the test seems to be missing the expected details, indicating
        # a bug in the test nodes. Since this test does not rely on the actual data in the buffer
        # being correct, the test is considered valid.

        await self.setUpProps(self._stage)

        output_image_dir = OUTPUTS_DIR.joinpath("omni.graph.rtxtest")

        keys = og.Controller.Keys
        (_, (_, test_node, _, _), _, _) = og.Controller.edit(
            self._post_render_graph,
            {
                keys.CREATE_NODES: [
                    ("RenderProductEntry", "omni.graph.nodes.GpuInteropRenderProductEntry"),
                    ("TestNode", "omni.graph.rtxtest.ScaleAOVTexture"),
                    ("GpuToCpuCopy", "omni.graph.examples.cpp.GpuInteropGpuToCpuCopy"),
                    ("CpuToDisk", "omni.graph.examples.cpp.GpuInteropCpuToDisk"),
                ],
                keys.CONNECT: [
                    ("RenderProductEntry.outputs:gpu", "TestNode.inputs:gpu"),
                    ("RenderProductEntry.outputs:rp", "TestNode.inputs:rp"),
                    ("RenderProductEntry.outputs:exec", "TestNode.inputs:exec"),
                    ("TestNode.outputs:gpu", "GpuToCpuCopy.inputs:gpu"),
                    ("TestNode.outputs:rp", "GpuToCpuCopy.inputs:rp"),
                    ("GpuToCpuCopy.outputs:gpu", "CpuToDisk.inputs:gpu"),
                    ("GpuToCpuCopy.outputs:rp", "CpuToDisk.inputs:rp"),
                    ("GpuToCpuCopy.outputs:aovCpu", "CpuToDisk.inputs:aovCpu"),
                ],
                keys.SET_VALUES: [
                    ("TestNode.inputs:inputAOV", "LdrColor"),
                    ("TestNode.inputs:outputAOV", "scaledColor"),
                    ("TestNode.inputs:multiplier", 1.5),
                    ("TestNode.inputs:width", self.WINDOW_SIZE[0]),
                    ("TestNode.inputs:height", self.WINDOW_SIZE[1]),
                    ("GpuToCpuCopy.inputs:aovGpu", "scaledColor"),
                    ("CpuToDisk.inputs:aovGpu", "scaledColor"),
                    ("CpuToDisk.inputs:frameCount", 1),
                    ("CpuToDisk.inputs:saveLocation", str(output_image_dir)),
                    ("CpuToDisk.inputs:fileName", "test_create_new_aov"),  # final name: test_create_new_aov_scaledColor.png
                    ("CpuToDisk.inputs:autoFileNumber", 0),
                ]
            },
        )

        await wait_for_update(self.ctx, 10)

        # validate that the outputAOVPtr attribute is set - this attribute is set after the new AOV is allocated
        actual = og.Controller.get(og.Controller.attribute("outputs:outputAOVPtr", test_node))
        self.assertNotEqual(0, actual)

        # validate that the scheduleCudaTask() method returned true
        schedule_result = og.Controller.get(og.Controller.attribute("outputs:scheduleResult", test_node))
        self.assertEqual(True, schedule_result)

        self.compare_output_image(
            "test_create_new_aov_scaledColor.png",
            THRESHOLD,
            output_image_dir,
            GOLDEN_DIR)

    # -----------------------------------------------------------------------------
    async def test_invalid_input_aov_param__fails(self):
        """Validate that nodes accessing invalid AOVs do not schedule cuda tasks"""

        keys = og.Controller.Keys
        (_, (_, test_node), _, _) = og.Controller.edit(
            self._post_render_graph,
            {
                keys.CREATE_NODES: [
                    ("RenderProductEntry", "omni.graph.nodes.GpuInteropRenderProductEntry"),
                    ("TestNode", "omni.graph.rtxtest.ScaleAOVTextureInPlace"),
                ],
                keys.CONNECT: [
                    ("RenderProductEntry.outputs:gpu", "TestNode.inputs:gpu"),
                    ("RenderProductEntry.outputs:rp", "TestNode.inputs:rp"),
                    ("RenderProductEntry.outputs:exec", "TestNode.inputs:exec"),
                ],
                keys.SET_VALUES: [
                    ("TestNode.inputs:inputAOV", "InvalidAOVToken"),
                    ("TestNode.inputs:multiplier", 1.5),
                    ("TestNode.inputs:width", self.WINDOW_SIZE[0]),
                    ("TestNode.inputs:height", self.WINDOW_SIZE[1]),
                ]
            },
        )

        await wait_for_update(self.ctx, 10)

        # the builder will log a warning if an AOV is not found
        self.assertEqual(len(test_node.get_compute_messages(og.WARNING)), 1)

        # validate that the scheduleCudaTask() method returned false
        schedule_result = og.Controller.get(og.Controller.attribute("outputs:scheduleResult", test_node))
        self.assertEqual(False, schedule_result)
