# 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.
#

from pathlib import Path
from typing import List

import carb.input
import carb.settings
import omni.kit.app
import omni.kit.test
import omni.kit.ui_test as ui_test
import omni.kit.undo
import omni.usd
from omni.kit.window.content_browser.test_helper import ContentBrowserTestHelper
from omni.kit.test_suite.helpers import arrange_windows
from omni.kit.window.file_importer.test_helper import FileImporterTestHelper


TEST_DATA_PATH = Path(
    f"{omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)}/data/tests"
)


ASSIGN_PATH_SETTINGS = "/persistent/app/material/dragDropMaterialPath"
ADD_ASSET_BUTTON_QUERY = "Property//Frame/**/Button[*].identifier=='sdf_asset_array_asset_array.add_asset'"


class TestAssetArray(omni.kit.test.AsyncTestCase):
    async def setUp(self):
        self._context = omni.usd.get_context()
        self._settings = carb.settings.get_settings()
        self._prev_path_setting = self._settings.get(ASSIGN_PATH_SETTINGS)
        self._settings.set(ASSIGN_PATH_SETTINGS, "relative")
        await arrange_windows("Stage", 1)
        async with ContentBrowserTestHelper() as content_browser_helper:
            await content_browser_helper.set_config_menu_settings({'hide_unknown': False, 'hide_thumbnails': True, 'show_details': False, 'show_udim_sequence': False})

    async def tearDown(self):
        self._settings.set(ASSIGN_PATH_SETTINGS, self._prev_path_setting)

    async def _add_asset(self, asset_names: List[str]):
        await self._browse_and_select_asset(ADD_ASSET_BUTTON_QUERY, asset_names)

    async def _browse_and_select_asset(self, browse_button_identifier: str, asset_names: List[str]):
        browse_button = ui_test.find(browse_button_identifier)
        await browse_button.click()

        async with FileImporterTestHelper() as file_importer:
            await file_importer.wait_for_popup()
            dir_url = TEST_DATA_PATH.absolute().resolve().joinpath("usd").joinpath("asset_array")
            await file_importer.select_items_async(str(dir_url), asset_names)
            await file_importer.click_apply_async(filename_url=None)

    async def test_asset_array_widget(self):
        usd_path = TEST_DATA_PATH.absolute().resolve().joinpath("usd").joinpath("asset_array").joinpath("main.usda")
        success, error = await self._context.open_stage_async(str(usd_path))
        self.assertTrue(success, error)

        self._context.get_selection().set_selected_prim_paths(["/Test"], True)

        asset_array_attr = self._context.get_stage().GetAttributeAtPath("/Test.asset_array")

        await ui_test.human_delay()

        # Expand Raw USD Properties section
        raw_properties_frame = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Raw USD Properties'")
        raw_properties_frame.widget.collapsed = False
        await ui_test.human_delay()

        # Add first asset
        await self._add_asset(["dummy0.txt"])
        await ui_test.human_delay()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt"])

        # Add second asset
        await self._add_asset(["dummy1.txt"])
        await ui_test.human_delay()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt", "./dummy1.txt"])

        # Reorder assets
        reorder_grab = ui_test.find(
            "Property//Frame/**/HStack[*].identifier=='sdf_asset_array_asset_array[0].reorder_grab'"
        )
        reorder_target_grab = ui_test.find(
            "Property//Frame/**/HStack[*].identifier=='sdf_asset_array_asset_array[1].reorder_grab'"
        )
        drag_target = reorder_target_grab.position + reorder_target_grab.size
        await reorder_grab.drag_and_drop(drag_target)
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy1.txt", "./dummy0.txt"])

        # Clear the first asset
        clear_asset = ui_test.find("Property//Frame/**/Button[*].identifier=='sdf_clear_asset_asset_array[0]'")
        await clear_asset.click()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["", "./dummy0.txt"])

        # Readd the first asset from file picker
        await self._browse_and_select_asset(
            "Property//Frame/**/Button[*].identifier=='sdf_browse_asset_asset_array[0]'", ["dummy1.txt"]
        )
        await ui_test.human_delay()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy1.txt", "./dummy0.txt"])

        # Remove the first asset
        remove_asset = ui_test.find("Property//Frame/**/Button[*].identifier=='sdf_asset_array_asset_array[0].remove'")
        await remove_asset.click()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt"])

        # Add 2 assets from file picker
        await self._add_asset(["dummy0.txt", "dummy1.txt"])
        await ui_test.human_delay()
        self.assertEqual(
            [asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt", "./dummy0.txt", "./dummy1.txt"]
        )

        # drag drop 2 assets from content window
        await ui_test.find("Content").focus()
        await ui_test.find("Property").focus()

        async with ContentBrowserTestHelper() as content_browser_helper:
            # await content_browser_helper.navigate_to_async(str(usd_path.parent))
            browse_button = ui_test.find(ADD_ASSET_BUTTON_QUERY)
            await content_browser_helper.drag_and_drop_tree_view(
                str(usd_path.parent), ["dummy0.txt", "dummy1.txt"], browse_button.center
            )
        await ui_test.human_delay()
        self.assertEqual(
            [asset_path.path for asset_path in asset_array_attr.Get()],
            ["./dummy0.txt", "./dummy0.txt", "./dummy1.txt", "./dummy0.txt", "./dummy1.txt"],
        )

        # Test undo

        omni.kit.undo.undo()
        self.assertEqual(
            [asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt", "./dummy0.txt", "./dummy1.txt"]
        )

        omni.kit.undo.undo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt"])

        omni.kit.undo.undo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy1.txt", "./dummy0.txt"])

        omni.kit.undo.undo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["", "./dummy0.txt"])

        omni.kit.undo.undo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy1.txt", "./dummy0.txt"])

        omni.kit.undo.undo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt", "./dummy1.txt"])

        omni.kit.undo.undo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt"])

        omni.kit.undo.undo()
        self.assertEqual(asset_array_attr.Get(), None)

        # Test redo
        omni.kit.undo.redo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt"])

        omni.kit.undo.redo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt", "./dummy1.txt"])

        omni.kit.undo.redo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy1.txt", "./dummy0.txt"])

        omni.kit.undo.redo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["", "./dummy0.txt"])

        omni.kit.undo.redo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy1.txt", "./dummy0.txt"])

        omni.kit.undo.redo()
        self.assertEqual([asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt"])

        omni.kit.undo.redo()
        self.assertEqual(
            [asset_path.path for asset_path in asset_array_attr.Get()], ["./dummy0.txt", "./dummy0.txt", "./dummy1.txt"]
        )

        omni.kit.undo.redo()
        self.assertEqual(
            [asset_path.path for asset_path in asset_array_attr.Get()],
            ["./dummy0.txt", "./dummy0.txt", "./dummy1.txt", "./dummy0.txt", "./dummy1.txt"],
        )
