## Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
##
## NVIDIA CORPORATION and its licensors retain all intellectual property
## and proprietary rights in and to this software, related documentation
## and any modifications thereto.  Any use, reproduction, disclosure or
## distribution of this software and related documentation without an express
## license agreement from NVIDIA CORPORATION is strictly prohibited.
##
import weakref
import omni.kit.app
import omni.kit.test
import omni.kit.commands
import omni.ui as ui
from omni.kit import ui_test
from omni.kit.test_suite.helpers import wait_stage_loading, select_prims
from omni.ui.tests.test_base import OmniUiTest
from omni.kit.property.usd.prim_selection_payload import PrimSelectionPayload
from pathlib import Path
from pxr import UsdGeom, Sdf
from ..scripts.transform_builder import TransformWidgets, TransformWatchModel, USDXformOpWidget
from ..scripts.transform_widget import TransformAttributeWidget
from ..scripts.transform_properties import g_ext

class TestTransformBuilder(OmniUiTest):

    xform_op_names = ["xformOp:translate", "xformOp:rotateXYZ", "xformOp:orient", "xformOp:scale", "xformOp:translate:pivot", "xformOp:transform"]

    # Before running each test
    async def setUp(self):
        await super().setUp()
        extension_root_folder = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))

        self._usd_path = extension_root_folder.joinpath("data/test_data/test_map")

    # After running each test
    async def tearDown(self):
        await super().tearDown()

    def get_index_op_code(self, xform : UsdGeom.Xformable, op_name):
        xform_ops = xform.GetOrderedXformOps()
        for i, op in enumerate(xform_ops):
            if op.GetOpName() == op_name:
                return i

        return -1

    async def test_builder_widgets(self):
        usd_context = omni.usd.get_context()

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()

        prim_path = '/World/Capsule'
        prim = stage.GetPrimAtPath(prim_path)
        xform = UsdGeom.Xformable(prim)
        transformWidget = TransformAttributeWidget("Test", True)

        await select_prims([prim_path])

        collapsable_frame = ui.Frame(height=0, style={"Frame": {"padding": 0}})

        # widget delete test
        widgets = []
        op_order_attr = prim.GetAttribute('XformOpOrder')
        for op_name in self.xform_op_names:
            #valid means disabled or not
            widget = transformWidget._transform_widget._create_xform_op_widget(
                        [Sdf.Path(prim_path)], collapsable_frame, stage, prim.GetAttribute(op_name), op_name,
                        True, op_order_attr, self.get_index_op_code(xform, op_name)
                    )
            widgets.append(widget)

        await ui_test.human_delay(3)

        transformWidget = None
        widgets = []

        transformWidget = TransformAttributeWidget("Test2", True)

        collapsable_frame = ui.Frame(height=0, style={"Frame": {"padding": 0}})

        widgets = []
        op_order_attr = prim.GetAttribute('XformOpOrder')
        for op_name in self.xform_op_names:
            #valid means disabled or not
            widget = transformWidget._transform_widget._create_xform_op_widget(
                        [Sdf.Path(prim_path)], collapsable_frame, stage, prim.GetAttribute(op_name), op_name,
                        False, op_order_attr, self.get_index_op_code(xform, op_name)
                    )
            widgets.append(widget)

        await ui_test.human_delay(3)
        transformWidget = None
        widgets = []

        #watch model
        watch_model = TransformWatchModel(0, stage)
        await ui_test.human_delay(3)
        watch_model.set_value(0.3)
        self.assertTrue(watch_model.get_value_as_float() == 0.3)
        self.assertTrue(watch_model.get_value_as_string() == "0.3")
        watch_model.clean()
        watch_model = None


        # USD Xform Op Widget
        for op_name in self.xform_op_names:
            #valid means disabled or not
            widget = USDXformOpWidget(
                        [Sdf.Path(prim_path)], collapsable_frame, stage, prim.GetAttribute(op_name).GetPath(), op_name,
                        True, op_order_attr, self.get_index_op_code(xform, op_name))

            #widget._inverse_op()
            widget._add_key()
            widget._show_right_click_menu(button=1)


    async def test_add_buttons(self):

        usd_context = omni.usd.get_context()

        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

        #get prim for sphere
        stage = usd_context.get_stage()

        prim_paths =["/World/Capsule", "/World/Sphere", "/World/Cube", "/World/Cylinder", "/World/Cone"]

        for prim_path in prim_paths:
            await select_prims([prim_path])

            prim = stage.GetObjectAtPath(prim_path)
            payload = PrimSelectionPayload( weakref.ref(stage), [prim.GetPath()])
            for button_menu in g_ext._add_button_menu:
                if button_menu.onclick_fn:
                    button_menu.onclick_fn(payload)
                    await ui_test.human_delay(2)
