## 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 omni.kit.app
import omni.kit.commands
import omni.kit.test
import omni.ui as ui
from omni.ui.tests.test_base import OmniUiTest
from pxr import Kind, Sdf, Gf


class TestBasicTypesRange(OmniUiTest):
    def __init__(self, tests=()):
        super().__init__(tests)
        self._widget_compare_table = {
            Sdf.ValueTypeNames.Half2.type: (Gf.Vec2h(0.12345, 0.12345), 3),
            Sdf.ValueTypeNames.Float2.type: (Gf.Vec2f(0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Double2.type: (Gf.Vec2d(0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Half3.type: (Gf.Vec3h(0.12345, 0.12345, 0.12345), 3),
            Sdf.ValueTypeNames.Float3.type: (Gf.Vec3f(0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Double3.type: (Gf.Vec3d(0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Half4.type: (Gf.Vec4h(0.12345, 0.12345, 0.12345, 0.12345), 3),
            Sdf.ValueTypeNames.Float4.type: (Gf.Vec4f(0.12345, 0.12345, 0.12345, 0.12345), 4),
            Sdf.ValueTypeNames.Double4.type: (Gf.Vec4d(0.12345, 0.12345, 0.12345, 0.12345), 4),
        }

    # Before running each test
    async def setUp(self):
        from omni.kit.test_suite.helpers import open_stage, get_test_data_path, arrange_windows

        await arrange_windows(topleft_window="Property", topleft_height=64, topleft_width=800.0)
        await open_stage(get_test_data_path(__name__, "usd/types.usda"))
        await self._show_raw(False)

    # After running each test
    async def tearDown(self):
        from omni.kit.test_suite.helpers import wait_stage_loading

        await wait_stage_loading()
        await self._show_raw(True)

    async def _show_raw(self, new_state):
        from omni.kit import ui_test
        from omni.kit.test_suite.helpers import select_prims

        await select_prims(["/defaultPrim/in_0"])
        for w in ui_test.find_all("Property//Frame/**/CollapsableFrame[*]"):
            if w.widget.title == "Raw USD Properties":
                w.widget.collapsed = new_state
        # change prim selection immediately after collapsed state change can result in one bad frame
        # i.e. selection change event happens before frame UI rebuild, and old widget models are already destroyed.
        await ui_test.human_delay()
        await select_prims([])

    async def test_l1_basic_type_coverage_ui(self):
        from omni.kit.test_suite.helpers import select_prims, wait_stage_loading
        from omni.kit import ui_test
        from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidgetBuilder

        UsdPropertiesWidgetBuilder.reset_builder_coverage_table()

        stage = omni.usd.get_context().get_stage()

        # wait for material to load & UI to refresh
        await wait_stage_loading()

        for prim_path in ["/defaultPrim/in_0", "/defaultPrim/in_1", "/defaultPrim/out_0"]:
            await select_prims([prim_path])

            # scroll so all prims are queryable
            for w in ui_test.find_all("Property//Frame/**/.identifier!=''"):
                w.widget.scroll_here_y(0.5)
                await ui_test.human_delay()

            # create attribute used table
            prim = stage.GetPrimAtPath(Sdf.Path(prim_path))
            attr_list = {}
            for attr in prim.GetAttributes():
                attr_list[attr.GetPath().name] = attr.GetTypeName().type.typeName

            # remove all widgets from attribute used table
            for w in ui_test.find_all("Property//Frame/**/.identifier!=''"):
                id = w.widget.identifier
                for type in [
                    "bool_",
                    "float_slider_",
                    "integer_slider_",
                    "drag_per_channel_",
                    "boolean_per_channel_",
                    "token_",
                    "string_",
                    "timecode_",
                    "sdf_asset_array_",  # must before "sdf_asset_"
                    "sdf_asset_",
                    "fallback_",
                ]:
                    if id.startswith(type):
                        attr_id = id[len(type) :]
                        # ignore placeholder widgets
                        if attr_id in attr_list:
                            del attr_list[attr_id]
                        break

            # check attribute used table is empty
            self.assertTrue(attr_list == {}, f"USD attribute {attr_list} have no widgets")

        widget_builder_coverage_table = UsdPropertiesWidgetBuilder.widget_builder_coverage_table
        for widget_type, builder_state in widget_builder_coverage_table.items():
            self.assertTrue(builder_state, f"Failed to test {widget_type}")

    async def test_l2_float_rounding(self):
        from omni.kit.test_suite.helpers import select_prims, wait_stage_loading
        from omni.kit import ui_test

        # wait for material to load & UI to refresh
        await wait_stage_loading()

        stage = omni.usd.get_context().get_stage()

        for prim_path in ["/defaultPrim/in_0", "/defaultPrim/in_1", "/defaultPrim/out_0"]:
            await select_prims([prim_path])

            # test float types
            for w in ui_test.find_all("Property//Frame/**/.identifier!=''"):
                id = w.widget.identifier
                if id.startswith("float_slider_"):
                    w.widget.scroll_here_y(0.5)
                    await ui_test.human_delay()
                    await w.input("0.12345")
                    await ui_test.human_delay()
                    attr = stage.GetPrimAtPath(prim_path).GetAttribute(id[13:])
                    self.assertAlmostEqual(attr.Get(), 0.12345, places=4)
                elif id.startswith("drag_per_channel_"):
                    if id.startswith("drag_per_channel_int"):
                        continue

                    w.widget.scroll_here_y(0.5)
                    await ui_test.human_delay()
                    sub_widgets = w.find_all("**/FloatSlider[*]")
                    if sub_widgets == []:
                        sub_widgets = w.find_all("**/FloatDrag[*]")
                    self.assertNotEqual(sub_widgets, [])

                    for child in sub_widgets:
                        child.model.set_value(0)
                        await child.input("0.12345")
                        await ui_test.human_delay()

                    attr = stage.GetPrimAtPath(prim_path).GetAttribute(id[17:])
                    metadata = attr.GetAllMetadata()
                    type_name = Sdf.ValueTypeNames.Find(metadata.get(Sdf.PrimSpec.TypeNameKey, "unknown type")).type
                    value, places = self._widget_compare_table[type_name]
                    self.assertAlmostEqual(attr.Get(), value, places=places)

    async def test_l2_float_rounding_tab(self):
        from omni.kit.test_suite.helpers import select_prims, wait_stage_loading
        from omni.kit import ui_test
        from carb.input import KeyboardInput

        # wait for material to load & UI to refresh
        await wait_stage_loading()

        stage = omni.usd.get_context().get_stage()

        for prim_path in ["/defaultPrim/in_0", "/defaultPrim/in_1", "/defaultPrim/out_0"]:
            await select_prims([prim_path])

            # test float types
            for w in ui_test.find_all("Property//Frame/**/.identifier!=''"):
                id = w.widget.identifier
                if id.startswith("float_slider_"):
                    w.widget.scroll_here_y(0.5)
                    await ui_test.human_delay()
                    await w.input("0.12345", end_key=KeyboardInput.TAB)
                    await ui_test.human_delay()
                    attr = stage.GetPrimAtPath(prim_path).GetAttribute(id[13:])
                    self.assertAlmostEqual(attr.Get(), 0.12345, places=4)
                elif id.startswith("drag_per_channel_"):
                    if id.startswith("drag_per_channel_int"):
                        continue

                    w.widget.scroll_here_y(0.5)
                    await ui_test.human_delay()
                    sub_widgets = w.find_all("**/FloatSlider[*]")
                    if sub_widgets == []:
                        sub_widgets = w.find_all("**/FloatDrag[*]")
                    self.assertNotEqual(sub_widgets, [])

                    for child in sub_widgets:
                        child.model.set_value(0)
                        await child.input("0.12345", end_key=KeyboardInput.TAB)
                        await ui_test.human_delay()

                    attr = stage.GetPrimAtPath(prim_path).GetAttribute(id[17:])
                    metadata = attr.GetAllMetadata()
                    type_name = Sdf.ValueTypeNames.Find(metadata.get(Sdf.PrimSpec.TypeNameKey, "unknown type")).type
                    value, places = self._widget_compare_table[type_name]
                    self.assertAlmostEqual(attr.Get(), value, places=places)
