## Copyright (c) 2018-2019, 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 asyncio
import omni.ui as ui
import omni.kit.app
import omni.kit.test
import omni.kit.ui_test as ui_test
from omni.ui.tests.test_base import OmniUiTest
from omni.kit.test_suite.helpers import get_test_data_path, wait_stage_loading
from omni.kit import ui_test

class TestOpenStage(OmniUiTest):
    """Testing omni.usd.core"""

    async def setUp(self):
        await super().setUp()
        await wait_stage_loading()
        self._future_test = None
        self._required_stage_event = -1
        self._stage_event_sub = omni.usd.get_context().get_stage_event_stream().create_subscription_to_pop(self._on_stage_event, name="omni.usd.menu.file")

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

    def _on_stage_event(self, event):
        if self._future_test and int(self._required_stage_event) == int(event.type) and not self._future_test.done():
            self._future_test.set_result(event.type)

    async def reset_stage_event(self, stage_event):
        self._required_stage_event = stage_event
        self._future_test = asyncio.Future()

    async def wait_for_stage_event(self):
        async def wait_for_event():
            await self._future_test

        try:
            await asyncio.wait_for(wait_for_event(), timeout=30.0)
        except asyncio.TimeoutError:
            carb.log_error(f"wait_for_stage_event timeout waiting for {self._required_stage_event}")

        self._future_test = None
        self._required_stage_event = -1

    async def test_open_stage(self):
        def on_load_complete(success):
            if not success:
                if self._future_test:
                    self._future_test.set_result(False)

        usd_context = omni.usd.get_context()

        # test file not found
        stage = omni.usd.get_context().get_stage()
        await self.reset_stage_event(omni.usd.StageEventType.OPENED)
        omni.kit.window.file.open_stage("missing file", callback=on_load_complete)
        await self.wait_for_stage_event()
        stage = omni.usd.get_context().get_stage()
        self.assertEqual(stage, omni.usd.get_context().get_stage())

        await ui_test.human_delay(50)

        # load valid stage
        stage = omni.usd.get_context().get_stage()
        await self.reset_stage_event(omni.usd.StageEventType.OPENED)
        omni.kit.window.file.open_stage(get_test_data_path(__name__, "4Lights.usda"), callback=on_load_complete)
        await self.wait_for_stage_event()
        self.assertNotEqual(stage, omni.usd.get_context().get_stage())

        await ui_test.human_delay(50)

        # create new stage
        await usd_context.new_stage_async()

        await ui_test.human_delay(50)

        # load a material as usd file
        stage = omni.usd.get_context().get_stage()
        await self.reset_stage_event(omni.usd.StageEventType.OPENED)
        success = omni.kit.window.file.open_stage(get_test_data_path(__name__, "anno_material.mdl"), callback=on_load_complete)
        await self.wait_for_stage_event()
        self.assertEqual(stage, omni.usd.get_context().get_stage())

        await ui_test.human_delay(50)
