import omni.kit.test
import omni.kit.ui_test as ui_test

from pathlib import Path
from omni.ui.tests.test_base import OmniUiTest
from omni.kit.test_suite.helpers import open_stage, wait_stage_loading, get_test_data_path, select_prims, arrange_windows
from omni.kit.window.content_browser.test_helper import ContentBrowserTestHelper


class TestDragDropFileToMaterial(OmniUiTest):
    # Before running each test
    async def setUp(self):
        await arrange_windows("Stage", 200)
        await open_stage(get_test_data_path(__name__, "usd/locate_file_material.usda"))
        await wait_stage_loading()

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

    async def test_drag_drop_single_mdl_asset_path(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        await select_prims(["/World/Looks/OmniPBR/Shader"])
        await wait_stage_loading()

        # verify material is correct
        shader = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"))
        asset = shader.GetSourceAsset("mdl")
        self.assertTrue(shader.GetPrim().IsValid())
        self.assertEqual(asset.path, "OmniPBR.mdl")

        # drag file from content window
        property_widget = ui_test.find("Property//Frame/**/StringField[*].identifier=='sdf_asset_info:mdl:sourceAsset'")
        drag_target = property_widget.center

        async with ContentBrowserTestHelper() as content_browser_helper:
            usd_path = get_test_data_path(__name__, "usd/TESTEXPORT.mdl")
            await content_browser_helper.drag_and_drop_tree_view(usd_path, drag_target=drag_target)

        # verify asset
        shader = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"))
        asset = shader.GetSourceAsset("mdl")
        self.assertTrue(shader.GetPrim().IsValid())
        self.assertTrue(asset.path.endswith("/TESTEXPORT.mdl"))

    async def test_drag_drop_single_png_asset_path(self):
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        await select_prims(["/World/Looks/OmniPBR/Shader"])
        await wait_stage_loading()

        # verify material is correct
        shader = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"))
        asset = shader.GetSourceAsset("mdl")
        self.assertTrue(shader.GetPrim().IsValid())
        self.assertEqual(asset.path, "OmniPBR.mdl")

        # drag file from content window
        property_widget = ui_test.find("Property//Frame/**/StringField[*].identifier=='sdf_asset_inputs:ao_texture'")
        drag_target = property_widget.center

        async with ContentBrowserTestHelper() as content_browser_helper:
            usd_path = get_test_data_path(__name__, "usd/textures/granite_a_mask.png")
            await content_browser_helper.drag_and_drop_tree_view(usd_path, drag_target=drag_target)

        # verify asset
        prim = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"), True)
        asset_path = prim.GetAttribute("inputs:ao_texture").Get()
        self.assertTrue(prim.IsValid())
        self.assertTrue(asset_path.path.endswith("textures/granite_a_mask.png"))

    async def test_drag_drop_multi_mdl_asset_path(self):
        # Test that dropping multiple material files results in no changes
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        await select_prims(["/World/Looks/OmniPBR/Shader"])
        await wait_stage_loading()

        # verify material is correct
        shader = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"))
        asset = shader.GetSourceAsset("mdl")
        self.assertTrue(shader.GetPrim().IsValid())
        self.assertEqual(asset.path, "OmniPBR.mdl")

        # drag file from content window
        property_widget = ui_test.find("Property//Frame/**/StringField[*].identifier=='sdf_asset_info:mdl:sourceAsset'")
        drag_target = property_widget.center

        async with ContentBrowserTestHelper() as content_browser_helper:
            usd_path = get_test_data_path(__name__, "usd")
            await content_browser_helper.drag_and_drop_tree_view(
                usd_path, names=["TESTEXPORT.mdl", "TESTEXPORT2.mdl"], drag_target=drag_target)

        # verify material not changed
        shader = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"))
        asset = shader.GetSourceAsset("mdl")
        self.assertTrue(shader.GetPrim().IsValid())
        self.assertEqual(asset.path, "OmniPBR.mdl")

    async def test_drag_drop_multi_png_asset_path(self):
        # Test that dropping multiple texture files results in no changes.
        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        await select_prims(["/World/Looks/OmniPBR/Shader"])
        await wait_stage_loading()

        # verify material is correct
        shader = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"))
        asset = shader.GetSourceAsset("mdl")
        self.assertTrue(shader.GetPrim().IsValid())
        self.assertEqual(asset.path, "OmniPBR.mdl")

        # drag file from content window
        property_widget = ui_test.find("Property//Frame/**/StringField[*].identifier=='sdf_asset_inputs:ao_texture'")
        drag_target = property_widget.center

        async with ContentBrowserTestHelper() as content_browser_helper:
            usd_path = get_test_data_path(__name__, "usd/textures")
            await content_browser_helper.drag_and_drop_tree_view(
                usd_path, names=["granite_a_mask.png", "granite_a_mask2.png"], drag_target=drag_target)

        # verify asset
        prim = omni.usd.get_shader_from_material(stage.GetPrimAtPath("/World/Looks/OmniPBR"), True)
        asset_path = prim.GetAttribute("inputs:ao_texture").Get()
        self.assertTrue(prim.IsValid())
        self.assertEqual(asset_path.path, "")
