## Copyright (c) 2022, 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 os
import omni.kit.app
import omni.kit.test
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.test_suite.helpers import (
    open_stage,
    get_test_data_path,
    select_prims,
    wait_stage_loading,
    arrange_windows
)
from omni.kit.window.content_browser.test_helper import ContentBrowserTestHelper


class ZZViewportDragDropMaterialMulti(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.assertEqual(bound_material.GetPrim().IsValid(), True)
                    self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, f"/World/Looks/{mtl_name}")

    def choose_material(self, mdl_list, excluded_list):
        # need to skip any multi-subid materials
        exclude_multi = {}
        for _, mdl_path, _ in mdl_list:
            mdl_name = os.path.splitext(os.path.basename(mdl_path))[0]
            if not mdl_name in exclude_multi:
                exclude_multi[mdl_name] = 1
            else:
                exclude_multi[mdl_name] += 1

        for _, mdl_path, _ in mdl_list:
            mdl_name = os.path.splitext(os.path.basename(mdl_path))[0]
            if mdl_name not in excluded_list and exclude_multi[mdl_name] < 2:
                return mdl_path, mdl_name

        return None, None

    async def test_l1_drag_drop_single_material_viewport(self):
        from carb.input import KeyboardInput

        await ui_test.find("Content").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()

        # verify bound materials before DnD
        for prim_path, material_path in [("/World/Cube", "/World/Looks/OmniGlass"), ("/World/Cone", "/World/Looks/OmniPBR"), ("/World/Sphere", "/World/Looks/OmniGlass"), ("/World/Cylinder", "/World/Looks/OmniPBR")]:
            prim = stage.GetPrimAtPath(prim_path)
            bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
            self.assertEqual(bound_material.GetPrim().IsValid(), True)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, material_path)

        # 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        viewport_window = ui_test.find("Viewport")
        viewport_window = ui_test.find("Viewport")
        await viewport_window.focus()
        drag_target = viewport_window.center

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

        # get paths
        mdl_list = await omni.kit.material.library.get_mdl_list_async()
        mdl_path1, mdl_name = self.choose_material(mdl_list, excluded_list=["OmniGlass", "OmniPBR"])

        # drag and drop items from the tree view of content browser
        async with ContentBrowserTestHelper() as content_browser_helper:
            await content_browser_helper.drag_and_drop_tree_view(mdl_path1, drag_target=drag_target)

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

        # verify bound materials after DnD - Should be unchanged
        for prim_path, material_path in [("/World/Cube", "/World/Looks/OmniGlass"), ("/World/Cone", f"/World/Looks/{mdl_name}"), ("/World/Sphere", "/World/Looks/OmniGlass"), ("/World/Cylinder", "/World/Looks/OmniPBR")]:
            prim = stage.GetPrimAtPath(prim_path)
            bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
            self.assertEqual(bound_material.GetPrim().IsValid(), True)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, material_path)

    async def test_l1_drag_drop_multi_material_viewport(self):
        from carb.input import KeyboardInput

        await ui_test.find("Content").focus()

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

        await wait_stage_loading()

        # verify bound materials before DnD
        for prim_path, material_path in expected_list:
            prim = stage.GetPrimAtPath(prim_path)
            bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
            self.assertEqual(bound_material.GetPrim().IsValid(), True)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, material_path)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, material_path)

        # 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
        viewport_window = ui_test.find("Viewport")
        await viewport_window.focus()
        drag_target = viewport_window.center

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

        # get paths
        mdl_list = await omni.kit.material.library.get_mdl_list_async()
        mdl_path1, mdl_name1 = self.choose_material(mdl_list, excluded_list=["OmniGlass", "OmniPBR"])
        mdl_path2, mdl_name2 = self.choose_material(mdl_list, excluded_list=["OmniGlass", "OmniPBR", mdl_name1])

        # drag and drop items from the tree view of content browser
        async with ContentBrowserTestHelper() as content_browser_helper:
            dir_url = os.path.dirname(os.path.commonprefix([mdl_path1, mdl_path2]))
            filenames = [os.path.basename(url) for url in [mdl_path1, mdl_path2]]
            await content_browser_helper.drag_and_drop_tree_view(dir_url, names=filenames, drag_target=drag_target)

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

        # verify bound materials after DnD - Should be unchanged
        for prim_path, material_path in expected_list:
            prim = stage.GetPrimAtPath(prim_path)
            bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
            self.assertEqual(bound_material.GetPrim().IsValid(), True)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, material_path)
