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


class DragDropMaterialViewport(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows()
        await open_stage(get_test_data_path(__name__, "bound_shapes.usda"))

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

    def _verify_material(self, stage, mtl_name, to_select):
            # verify material created
            prim_paths = [p.GetPrimPath().pathString for p in stage.Traverse()]
            self.assertTrue(f"/World/Looks/{mtl_name}" in prim_paths)
            self.assertTrue(f"/World/Looks/{mtl_name}/Shader" in prim_paths)

            # verify bound material
            if to_select:
                for prim_path in to_select:
                    prim = stage.GetPrimAtPath(prim_path)
                    bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
                    self.assertTrue(bound_material.GetPrim().IsValid() == True)
                    self.assertTrue(bound_material.GetPrim().GetPrimPath().pathString == f"/World/Looks/{mtl_name}")

    # NOTE: This crashes in hydra so disabled
    async def __test_l1_drag_drop_material_viewport(self):
        await ui_test.find("Content").focus()
        await ui_test.find("Stage").focus()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        to_select = ["/World/Cube", "/World/Cone", "/World/Sphere", "/World/Cylinder"]

        await wait_stage_loading()

        # drag to center of viewport window
        # NOTE: Material binding is only done if dropped over prim. This just creates materials as center coordinates are used
        content_browser_helper = ContentBrowserTestHelper()
        material_test_helper = MaterialLibraryTestHelper()

        viewport_window = ui_test.find("Viewport")
        await viewport_window.focus()
        drag_target = viewport_window.center
        mdl_list = await omni.kit.material.library.get_mdl_list_async()
        for mtl_name, mdl_path, submenu in mdl_list:
            # delete anything in Looks
            await delete_prim_path_children("/World/Looks")

            # select prim
            await select_prims(to_select)
            # wait for UI to update
            await ui_test.human_delay()

            # drag/drop
            await content_browser_helper.drag_and_drop_tree_view(mdl_path, drag_target=drag_target)

            # handle create material dialog
            await material_test_helper.handle_create_material_dialog(mdl_path, mtl_name)

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

            # verify material
            self._verify_material(stage, mtl_name, [])

