import weakref
import platform
import unittest
import pathlib
import omni.kit.test
import omni.ui as ui
import carb
from pxr import Usd, Sdf, UsdGeom
from omni.ui.tests.test_base import OmniUiTest
from omni.kit import ui_test
from omni.kit.test_suite.helpers import arrange_windows, wait_stage_loading


class TestMaterialWidget(OmniUiTest):
    async def setUp(self):
        await super().setUp()
        await arrange_windows()

        from omni.kit.property.material.scripts.material_properties import TEST_DATA_PATH
        self._golden_img_dir = TEST_DATA_PATH.absolute().joinpath("golden_img").absolute()
        self._usd_path = TEST_DATA_PATH.absolute()
        self._icon_path = TEST_DATA_PATH.resolve().parent.joinpath("icons")

        from omni.kit.property.usd.usd_attribute_widget import UsdPropertiesWidget
        import omni.kit.window.property as p

        self._w = p.get_window()


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


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

        await usd_context.new_stage_async()
        material_prims, bound_prims = omni.kit.property.material.tests.create_test_stage()

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=285,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/LightPivot_01/AreaLight/Backside"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_material_ui.png")
        usd_context.get_selection().set_selected_prim_paths([], True)


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

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=520,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/LightPivot_01/AreaLight/Backside", "/LightPivot_01/AreaLight/EmissiveSurface"], True)

        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay(10)

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_material_multiselect_ui.png")
        usd_context.get_selection().set_selected_prim_paths([], True)


    async def test_material_popup_full_ui(self):
        from omni.kit.window.property.templates import SimplePropertyWidget, LABEL_WIDTH, LABEL_HEIGHT
        from omni.kit.material.library.listbox_widget import MaterialListBoxWidget

        usd_context = omni.usd.get_context()
        test_file_path = self._usd_path.joinpath("material_binding.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        # full window
        window = await self.create_test_window(width=460, height=380)
        with window.frame:
            with ui.VStack(width=0, height=0):
                with ui.HStack(width=0, height=0):
                    name_field = ui.StringField(width=400, height=20, enabled=False)
                    name_field.model.set_value("TEST"*4)

                await ui_test.human_delay(10)

                listbox_widget = MaterialListBoxWidget(icon_path=None, index=0, on_click_fn=None, theme={})
                listbox_widget.set_parent(name_field)
                listbox_widget.build_ui()
                # material loading is async so alow for loadng
                await ui_test.human_delay(50)

        await ui_test.human_delay(10)
        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_material_popup1_ui.png")


    async def test_material_popup_search_ui(self):
        from omni.kit.window.property.templates import SimplePropertyWidget, LABEL_WIDTH, LABEL_HEIGHT
        from omni.kit.material.library.listbox_widget import MaterialListBoxWidget

        usd_context = omni.usd.get_context()
        test_file_path = self._usd_path.joinpath("material_binding.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await omni.kit.app.get_app().next_update_async()

        # search
        window = await self.create_test_window(width=460, height=380)
        with window.frame:
            with ui.VStack(width=0, height=0):
                with ui.HStack(width=0, height=0):
                    name_field = ui.StringField(width=400, height=20, enabled=False)
                    name_field.model.set_value("SEARCH"*4)

                await ui_test.human_delay(10)

                listbox_widget = MaterialListBoxWidget(icon_path=None, index=0, on_click_fn=None, theme={})
                listbox_widget.set_parent(name_field)
                listbox_widget.build_ui()
                await ui_test.human_delay(10)
                listbox_widget._search_updated("PBR")
                # material loading is async so alow for loadng
                await wait_stage_loading()

        await ui_test.human_delay(10)
        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_material_popup2_ui.png")


    async def test_material_description_ui(self):
        from omni.kit import ui_test

        usd_context = omni.usd.get_context()

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

        await self.docked_test_window(
            window=self._w._window,
            width=450,
            height=285,
            restore_window = ui.Workspace.get_window("Layer") or ui.Workspace.get_window("Stage"),
            restore_position = ui.DockPosition.BOTTOM)

        # Select the prim.
        usd_context.get_selection().set_selected_prim_paths(["/World/Looks/OmniPBR"], True)
        # Need to wait for an additional frames for omni.ui rebuild to take effect
        await ui_test.human_delay()

        for w in ui_test.find_all("Property//Frame/**/CollapsableFrame[*]"):
            if w.widget.title != "Raw USD Properties":
                w.widget.collapsed = False
        await ui_test.human_delay()

        await self.finalize_test(golden_img_dir=self._golden_img_dir, golden_img_name="test_material_description_ui.png")
        usd_context.get_selection().set_selected_prim_paths([], True)


