## 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.test
import asyncio
import os
import carb
import omni.usd
import omni.kit.app
from omni.kit.async_engine import run_coroutine
import concurrent.futures

from typing import List
from omni.kit.test.async_unittest import AsyncTestCase
from omni.kit import ui_test
from pxr import Tf, Usd, Sdf
from omni.kit.test_suite.helpers import (
    open_stage,
    get_test_data_path,
    wait_stage_loading,
    delete_prim_path_children,
    arrange_windows
    )
from omni.kit.window.content_browser.test_helper import ContentBrowserTestHelper


class DragDropUsdViewportItem(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        carb.settings.get_settings().set("/persistent/app/stage/dragDropImport", "reference")
        await arrange_windows("Stage", 512)
        await open_stage(get_test_data_path(__name__, "bound_shapes.usda"))
        await delete_prim_path_children("/World")
        self._usd_path = get_test_data_path(__name__, "shapes").replace("\\", "/")

    # After running each test
    async def tearDown(self):
        await wait_stage_loading()
        carb.settings.get_settings().set("/persistent/app/stage/dragDropImport", "reference")

    def assertPathsEqual(self, path_a: str, path_b: str):
        # Make drive comparison case insensitive on Windows
        drive_a, path_a = os.path.splitdrive(path_a)
        drive_b, path_b = os.path.splitdrive(path_b)
        self.assertEqual(path_a, path_b)
        self.assertEqual(drive_a.lower(), drive_b.lower())

    async def wait_for_import(self, stage, prim_paths, drag_drop_fn):
        # create future
        future_test = asyncio.Future()

        all_expected_prim_paths = set(prim_paths)
        def on_objects_changed(notice, sender, future_test):
            for p in notice.GetResyncedPaths():
                if p.pathString in all_expected_prim_paths:
                    all_expected_prim_paths.discard(p.pathString)

                    if not all_expected_prim_paths:
                        async def future_test_complete(future_test):
                            await omni.kit.app.get_app().next_update_async()
                            future_test.set_result(True)

                        if not self._test_complete:
                            self._test_complete = run_coroutine(future_test_complete(future_test))
                        break

        async def wait_for_event(future_test):
            await future_test

        # create listener
        self._test_complete = None
        listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, lambda n, s, : on_objects_changed(n, s, future_test), stage)

        # do drag/drop
        if asyncio.isfuture(drag_drop_fn()):
            await drag_drop_fn()
        else:
            concurrent.futures.wait(drag_drop_fn())

        # wait for Tf.Notice event
        try:
            await asyncio.wait_for(wait_for_event(future_test), timeout=30.0)
        except asyncio.TimeoutError:
            carb.log_error(f"wait_for_import timeout")

        # release listener
        listener.Revoke()

    def _verify(self, stage, prim_path, usd, reference):
        prim = stage.GetPrimAtPath(prim_path)
        self.assertIsNotNone(prim)
        payloads = omni.usd.get_composed_payloads_from_prim(prim)
        references = omni.usd.get_composed_references_from_prim(prim)

        if reference:
            items = references
            self.assertEqual(payloads, [])
            self.assertEqual(len(references), 1)
        else:
            items = payloads
            self.assertEqual(len(payloads), 1)
            self.assertEqual(references, [])

        for (ref, layer) in items:
            # unlike stage_window this ref.assetPath is not absoloute
            abs_path = get_test_data_path(__name__, f"{ref.assetPath}").replace("\\", "/")
            self.assertPathsEqual(abs_path, usd)

    async def _drag_drop_usd_items(self, usd_path: str, usd_names: List[str], reference: bool = False):
        await ui_test.find("Content").focus()
        viewport_window = ui_test.find("Viewport")
        await viewport_window.focus()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()
        await wait_stage_loading()

        # set dragDropImport
        import_method = "reference" if reference else "payload"
        carb.settings.get_settings().set("/persistent/app/stage/dragDropImport", import_method)

        all_expected_prim_paths = []
        for usd_name in usd_names:
            prim_name, _ = os.path.splitext(os.path.basename(usd_name))
            all_expected_prim_paths.append(f"/World/{prim_name}")

        # drag/drop
        content_browser_helper = ContentBrowserTestHelper()
        await self.wait_for_import(stage, all_expected_prim_paths,
            lambda: run_coroutine(content_browser_helper.drag_and_drop_tree_view(
                usd_path, names=usd_names, drag_target=viewport_window.center
            ))
        )
        await wait_stage_loading()

        # verify
        for usd_name, prim_path in zip(usd_names, all_expected_prim_paths):
            self._verify(stage, prim_path, f"{usd_path}/{usd_name}", reference)

    async def test_l1_drag_drop_usd_stage_item_reference(self):
        await self._drag_drop_usd_items(self._usd_path, ["basic_sphere.usda"], reference=True)

    async def test_l1_drag_drop_usd_viewport_item_payload(self):
        await self._drag_drop_usd_items(self._usd_path, ["basic_cone.usda"], reference=False)

    async def test_l1_drag_drop_multiple_usd_stage_items_reference(self):
        usd_names = ["basic_cube.usda", "basic_cone.usda", "basic_sphere.usda"]
        await self._drag_drop_usd_items(self._usd_path, usd_names, reference=True)

    async def test_l1_drag_drop_multiple_usd_viewport_items_payload(self):
        usd_names = ["basic_cube.usda", "basic_cone.usda", "basic_sphere.usda"]
        await self._drag_drop_usd_items(self._usd_path, usd_names, reference=False)
