## 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.
##
from typing import Optional
from ..extension import ExampleExtension
from omni.ui.tests.test_base import OmniUiTest
from pathlib import Path
import omni.kit
import omni.ui as ui
import carb

CURRENT_PATH = Path(__file__).parent


# TreeView test does a golden-image comparison against model generated from omni.kit.commands.get_commands
# This will break the test whenever a command is added or removed, so hi-jack the function to
# return a known set of commands for the image
class TreeViewTestFix:
    def __init__(self):
        import omni.kit.commands
        # Build a static known command list and warn if the golden image should be updated
        known_cmds = ['Redo', 'Repeat', 'Undo', 'Command', 'ChangeSettingCommand', 'GroupPrimsCommand', 'CreatePrimWithDefaultXformCommand', 'CreatePrimCommand', 'CopyPrimCommand', 'CopyPrimsCommand', 'CreateInstanceCommand', 'CreateInstancesCommand', 'DeletePrimsCommand', 'CreatePrimsCommand', 'CreateDefaultXformOnPrimCommand', 'BindMaterialCommand', 'SetMaterialStrengthCommand', 'TransformPrimCommand', 'TransformPrimSRTCommand', 'TransformPrimsCommand', 'TransformPrimsSRTCommand', 'FramePrimsCommand', 'SelectPrimsCommand', 'ToggleVisibilitySelectedPrimsCommand', 'UnhideAllPrimsCommand', 'MovePrimCommand', 'MovePrimsCommand', 'ReplaceReferencesCommand', 'CreateUsdAttributeOnPathCommand', 'CreateUsdAttributeCommand', 'ChangePropertyCommand', 'RemovePropertyCommand', 'ChangeMetadataInPrimsCommand', 'ChangeMetadataCommand', 'ChangeAttributesColorSpaceCommand', 'CreateMdlMaterialPrimCommand', 'CreatePreviewSurfaceMaterialPrimCommand', 'CreatePreviewSurfaceTextureMaterialPrimCommand', 'ClearRefinementOverridesCommand', 'RelationshipTargetBase', 'AddRelationshipTargetCommand', 'RemoveRelationshipTargetCommand', 'SetRelationshipTargetsCommand', 'ReferenceCommandBase', 'AddReferenceCommand', 'RemoveReferenceCommand', 'ReplaceReferenceCommand', 'PayloadCommandBase', 'AddPayloadCommand', 'RemovePayloadCommand', 'ReplacePayloadCommand', 'CreatePrimCommandBase', 'CreateReferenceCommand', 'CreatePayloadCommand', 'CreateAudioPrimFromAssetPathCommand', 'SetDisplayColorCommand']
        cmd_dict = {}
        for cmd_list in omni.kit.commands.get_commands().values():
            for cmd in cmd_list.values():
                cmd_name = cmd.__name__
                if cmd_name in known_cmds:
                    cmd_dict[cmd_name] = type(cmd_name, (object,), {})
                else:
                    carb.log_warn(f'TreeView golden-image test needs an update because of command: {cmd_name}')
        if len(cmd_dict) != len(known_cmds):
            carb.log_warn('TreeView golden-image test needs an update because a command was aded or removed')

        # Hijack the get_commands function
        self.__orig_get_commands = omni.kit.commands.get_commands
        omni.kit.commands.get_commands = lambda: {'indirection': cmd_dict}

    def __del__(self):
        import omni.kit.commands
        omni.kit.commands.get_commands = self.__orig_get_commands


class TestExampleUi(OmniUiTest):
    # Before running each test
    async def setUp(self):
        await super().setUp()
        self._golden_img_dir = CURRENT_PATH.parent.parent.parent.parent.joinpath("data", "tests")
        ui.Workspace.show_window(ExampleExtension.WINDOW_NAME)
        ExampleExtension.ui_doc().omniverse_version = "test"
        ExampleExtension.ui_doc().vertical_scrollbar_on = False
        ExampleExtension.ui_doc().rebuild()

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

    async def capture_doc(self, topic: Optional[str] = None, navigate_to: Optional[str] = None):
        # One frame to rebuild the page
        await omni.kit.app.get_app().next_update_async()

        # Navigate to the specified topic
        if topic:
            entry = next(e for e in ExampleExtension.ui_doc().index_list if e[0] == topic)
            ExampleExtension.ui_doc().show_a_doc(entry, navigate_to)
            await omni.kit.app.get_app().next_update_async()

        window = ui.Workspace.get_window(ExampleExtension.WINDOW_NAME)

        await self.docked_test_window(
            window=window,
            width=1000,
            height=500,
            restore_window=ui.Workspace.get_window("Content"),
            restore_position=ui.DockPosition.SAME,
        )

    async def test_all(self):
        await self.capture_doc()
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Arrangement_of_elements(self):
        await self.capture_doc("Arrangement of elements")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Lengths(self):
        await self.capture_doc("Lengths")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Styling(self):
        await self.capture_doc("Styling")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Base_Widgets(self):
        await self.capture_doc("Base Widgets")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Menu(self):
        await self.capture_doc("Menu")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_ScrollingFrame(self):
        await self.capture_doc("ScrollingFrame")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_CanvasFrame(self):
        await self.capture_doc("CanvasFrame")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Image(self):
        await self.capture_doc("Image")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_CheckBox(self):
        await self.capture_doc("CheckBox")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_AbstractValueModel(self):
        await self.capture_doc("AbstractValueModel")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_ComboBox(self):
        await self.capture_doc("ComboBox")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_ColorWidget(self):
        await self.capture_doc("ColorWidget")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_MultiField(self):
        await self.capture_doc("MultiField")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_TreeView(self):
        tree_view_test_fix = TreeViewTestFix()
        await self.capture_doc("TreeView")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)
        del tree_view_test_fix

    async def test_Shapes(self):
        await self.capture_doc("Shapes")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Fields(self):
        await self.capture_doc("Fields")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Sliders(self):
        await self.capture_doc("Sliders")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_ProgressBar(self):
        await self.capture_doc("ProgressBar")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Plot(self):
        await self.capture_doc("Plot")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Drag_and_Drop(self):
        await self.capture_doc("Drag and Drop")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Windows(self):
        await self.capture_doc("Windows")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_MenuBar(self):
        await self.capture_doc("MenuBar")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)

    async def test_Workspace(self):
        await self.capture_doc("Workspace")
        await self.finalize_test(golden_img_dir=self._golden_img_dir)
