import omni.kit.test

from pathlib import Path
from pxr import Sdf, UsdShade, UsdGeom
from omni.ui.tests.test_base import OmniUiTest
from omni.kit import ui_test
from omni.kit.test_suite.helpers import select_prims, arrange_windows
from omni.kit.window.content_browser.test_helper import ContentBrowserTestHelper


class TestDragDropMDLFile(OmniUiTest):
    async def test_drag_drop_single_mdl_file_target(self):
        await arrange_windows()

        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        mdl_path1 = Path(extension_path).joinpath("data/tests/TESTEXPORT.mdl").absolute()

        # setup
        await ui_test.find("Stage").focus()
        await ui_test.find("Content").focus()

        # new stage
        usd_context = omni.usd.get_context()
        await usd_context.new_stage_async()
        stage = usd_context.get_stage()

        # create prims
        omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
        await select_prims(["/Sphere"])
        await ui_test.human_delay()

        # drag/drop
        async with ContentBrowserTestHelper() as content_browser_helper:
            property_widget = ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'")
            await content_browser_helper.drag_and_drop_tree_view(str(mdl_path1), drag_target=property_widget.center)

        # verify - TESTEXPORT should be bound
        # NOTE: TESTEXPORT.mdl material is named "Material" so that is the prim created
        prim = stage.GetPrimAtPath("/Sphere")
        bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
        self.assertTrue(bound_material.GetPrim().IsValid() == True)
        self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/Looks/Material")

    async def test_drag_drop_multi_mdl_file_target(self):
        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        mdl_path = str(Path(extension_path).joinpath("data/tests").absolute())

        # setup
        await ui_test.find("Stage").focus()
        await ui_test.find("Content").focus()

        # new stage
        usd_context = omni.usd.get_context()
        await usd_context.new_stage_async()
        stage = usd_context.get_stage()

        # create prims
        omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
        await select_prims(["/Sphere"])
        await ui_test.human_delay()

        # drag/drop
        async with ContentBrowserTestHelper() as content_browser_helper:
            property_widget = ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'")
            await content_browser_helper.drag_and_drop_tree_view(
                mdl_path, names=["TESTEXPORT.mdl", "TESTEXPORT2.mdl"], drag_target=property_widget.center)

        # verify - nothing should be bound as 2 items where selected
        prim = stage.GetPrimAtPath("/Sphere")
        bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
        self.assertFalse(bound_material.GetPrim().IsValid())

    async def test_drag_drop_single_mdl_file_combo(self):
        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        mdl_path1 = Path(extension_path).joinpath("data/tests/TESTEXPORT.mdl").absolute()

        # setup
        await ui_test.find("Stage").focus()
        await ui_test.find("Content").focus()

        # new stage
        usd_context = omni.usd.get_context()
        await usd_context.new_stage_async()
        stage = usd_context.get_stage()

        # create prims
        omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
        await select_prims(["/Sphere"])
        await ui_test.human_delay()

        # drag/drop
        async with ContentBrowserTestHelper() as content_browser_helper:
            property_widget = ui_test.find("Property//Frame/**/StringField[*].identifier=='combo_drop_target'")
            await content_browser_helper.drag_and_drop_tree_view(str(mdl_path1), drag_target=property_widget.center)

        # verify - TESTEXPORT should be bound
        # NOTE: TESTEXPORT.mdl material is named "Material" so that is the prim created
        prim = stage.GetPrimAtPath("/Sphere")
        bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
        self.assertTrue(bound_material.GetPrim().IsValid() == True)
        self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/Looks/Material")

    async def test_drag_drop_multi_mdl_file_combo(self):
        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        mdl_path = str(Path(extension_path).joinpath("data/tests").absolute())

        # setup
        await ui_test.find("Stage").focus()
        await ui_test.find("Content").focus()

        # new stage
        usd_context = omni.usd.get_context()
        await usd_context.new_stage_async()
        stage = usd_context.get_stage()

        # create prims
        omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
        await select_prims(["/Sphere"])
        await ui_test.human_delay()

        # drag/drop
        async with ContentBrowserTestHelper() as content_browser_helper:
            property_widget = ui_test.find("Property//Frame/**/StringField[*].identifier=='combo_drop_target'")
            await content_browser_helper.drag_and_drop_tree_view(
                mdl_path, names=["TESTEXPORT.mdl", "TESTEXPORT2.mdl"], drag_target=property_widget.center)

        # verify - nothing should be bound as 2 items where selected
        prim = stage.GetPrimAtPath("/Sphere")
        bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
        self.assertFalse(bound_material.GetPrim().IsValid())
