## 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
from omni.kit.test.async_unittest import AsyncTestCase
import omni.usd
from omni.kit import ui_test
from pxr import Sdf, UsdShade
from omni.kit.test_suite.helpers import (
    open_stage,
    get_test_data_path,
    select_prims,
    wait_stage_loading,
    build_sdf_asset_frame_dictonary,
    arrange_windows
    )
from omni.kit.window.content_browser.test_helper import ContentBrowserTestHelper


class DragDropTextureProperty(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows("Stage", 64)
        await open_stage(get_test_data_path(__name__, "bound_material.usda"))

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

    async def test_l1_drag_drop_texture_property(self):
        await ui_test.find("Stage").focus()
        await ui_test.find("Content").focus()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        to_select = ["/World/Looks/OmniHair/Shader"]

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

        # select prim
        await select_prims(to_select)
        await ui_test.human_delay()

        # open all CollapsableFrames
        for frame in ui_test.find_all("Property//Frame/**/CollapsableFrame[*]"):
            if frame.widget.title != "Raw USD Properties":
                frame.widget.scroll_here_y(0.5)
                frame.widget.collapsed = False
                await ui_test.human_delay()

        # prep content window for drag/drop
        texture_path = get_test_data_path(__name__, "textures/granite_a_mask.png")
        content_browser_helper = ContentBrowserTestHelper()

        # NOTE: cannot keep using widget as they become stale as window refreshes due to set_value

        # do tests
        widget_table = await build_sdf_asset_frame_dictonary()
        for frame_name in widget_table.keys():
            for field_name in widget_table[frame_name].keys():
                if field_name in ["sdf_asset_info:mdl:sourceAsset"]:
                    continue

                # NOTE: sdf_asset_ could be used more than once, if so skip
                if widget_table[frame_name][field_name] == 1:
                    def get_widget():
                        return ui_test.find(f"Property//Frame/**/CollapsableFrame[*].title=='{frame_name}'").find(f"**/StringField[*].identifier=='{field_name}'")

                    # scroll to widget
                    get_widget().widget.scroll_here_y(0.5)
                    await ui_test.human_delay()

                    # reset for test
                    get_widget().model.set_value("")

                    # drag/drop
                    await content_browser_helper.drag_and_drop_tree_view(texture_path, drag_target=get_widget().center)

                    # verify dragged item(s)
                    for prim_path in to_select:
                        prim = stage.GetPrimAtPath(Sdf.Path(prim_path))
                        self.assertTrue(prim.IsValid())
                        attr = prim.GetAttribute(field_name[10:])
                        self.assertTrue(attr.IsValid())
                        asset_path = attr.Get()
                        self.assertTrue(asset_path != None)
                        self.assertTrue(asset_path.resolvedPath.replace("\\", "/").lower() == texture_path.replace("\\", "/").lower())

                    # reset for next test
                    get_widget().model.set_value("")
