"""test viewport nodes"""

import asyncio
import inspect
import pathlib

import carb.input

import omni.graph.core as og
import omni.usd
from omni.kit.ui_test.input import emulate_keyboard_press
from omni.rtx.tests import RtxTest, postLoadTestSettings, testSettings
from omni.rtx.tests.test_common import wait_for_update
from omni.ui.tests.compare_utils import CompareMetric, capture_and_compare

from pxr import UsdGeom

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


# ======================================================================
class TestViewportNodes(RtxTest):

    def __init__(self, tests=()):
        super().__init__(tests)
        self.test_graph_path = "/World/TestGraph"

    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()

    @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]}"

    async def capture_and_compare(self, image_name: str, threshold=RtxTest.THRESHOLD, cmp_metric=CompareMetric.MEAN_ERROR_SQUARED):
        """
        RtxTest.capture_and_compare() captures directly from the viewport's texture, however,
        we need to capture from the swap chain in order to take into account the overlay image
        that hides the current frame during viewport render lock.
        """
        diff = await capture_and_compare(image_name, threshold, GOLDEN_IMG_DIR, True, cmp_metric)

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

        if diff >= threshold:
            self._failedImages.append(image_name)

        return diff

    def create_mesh(self, name):
        """Create a cube mesh"""
        cube = UsdGeom.Mesh.Define(self.ctx.get_stage(), name)
        cube.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),
            ]
        )
        cube.CreateFaceVertexCountsAttr([4, 4, 4, 4, 4, 4])
        cube.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])
        cube.CreateSubdivisionSchemeAttr("none")
        return cube

    async def test_lock_viewport_render(self):
        """Validate viewport render lock & unlock for OgnLockViewportRender"""
        # create a cube mesh
        cube = self.create_mesh("/World/cube")

        # create action graph
        controller = og.Controller()
        graph = controller.create_graph({"graph_path": self.test_graph_path, "evaluator_name": "execution"})
        controller.edit(
            graph,
            {
                og.Controller.Keys.CREATE_NODES: [
                    ("on_keyboard_input_l", "omni.graph.action.OnKeyboardInput"),
                    ("on_keyboard_input_u", "omni.graph.action.OnKeyboardInput"),
                    ("lock_viewport_render", "omni.graph.ui_nodes.LockViewportRender"),
                ],
                og.Controller.Keys.CONNECT: [
                    ("on_keyboard_input_l.outputs:released", "lock_viewport_render.inputs:lock"),
                    ("on_keyboard_input_u.outputs:released", "lock_viewport_render.inputs:unlock"),
                ],
                og.Controller.Keys.SET_VALUES: [
                    ("on_keyboard_input_l.inputs:keyIn", "L"),
                    ("on_keyboard_input_l.inputs:onlyPlayback", False),
                    ("on_keyboard_input_u.inputs:keyIn", "U"),
                    ("on_keyboard_input_u.inputs:onlyPlayback", False),
                ],
            },
        )

        # Initial
        await wait_for_update(self.ctx)
        await self.capture_and_compare("test_lock_viewport_render.initial.png")

        # Lock viewport render
        await emulate_keyboard_press(carb.input.KeyboardInput.L)
        await wait_for_update(self.ctx)
        await self.capture_and_compare("test_lock_viewport_render.locked.png")

        # Scale the cube
        scale_op = UsdGeom.Xformable(cube).AddScaleOp()
        scale = 2.0
        scale_op.Set((scale, scale, scale))
        await wait_for_update(self.ctx)
        await self.capture_and_compare("test_lock_viewport_render.scaled.png")

        # Unlock viewport render and wait for the fading duration (1 second by default) to complete
        await emulate_keyboard_press(carb.input.KeyboardInput.U)
        await asyncio.sleep(1.1)
        await self.capture_and_compare("test_lock_viewport_render.unlocked.png")

        # Lock viewport render and scale the cube again
        await emulate_keyboard_press(carb.input.KeyboardInput.L)
        scale = 3.0
        scale_op.Set((scale, scale, scale))
        await wait_for_update(self.ctx)
        await self.capture_and_compare("test_lock_viewport_render.relocked.png")

        # Remove the node. The viewport render will be automatically unlocked.
        controller.edit(graph, { og.Controller.Keys.DELETE_NODES: "lock_viewport_render" })
        await wait_for_update(self.ctx)
        await self.capture_and_compare("test_lock_viewport_render.removed.png")
