## 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 os
import omni.kit.app
import omni.usd

from omni.kit.test.async_unittest import AsyncTestCase
from omni.kit import ui_test
from omni.kit.test_suite.helpers import open_stage, get_test_data_path, select_prims, wait_stage_loading, arrange_windows
from omni.kit.window.file_exporter.test_helper import FileExporterTestHelper
from omni.kit.window.file_exporter import get_instance as get_file_exporter


class TestExportSelected(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows()
        await open_stage(get_test_data_path(__name__, "bound_shapes.usda"))
        stage_window = ui_test.find("Stage")
        await stage_window.focus()

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


    async def test_l1_stage_menu_export_selected_single_prim(self):
        stage = omni.usd.get_context().get_stage()
        to_select = ["/World/Cone"]
        test_file = os.path.join(omni.kit.test.get_test_output_path(), "save_single_selected_prim.usd")

        # select prims
        await select_prims(to_select)

        # right click on Cube
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

        # click on context menu item & save
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            await file_export_helper.click_apply_async(filename_url=test_file)
            await wait_stage_loading()

        # load created file
        await open_stage(test_file)
        await wait_stage_loading()

        # verify prims
        stage = omni.usd.get_context().get_stage()
        prims = [prim.GetPath().pathString for prim in stage.TraverseAll() if not omni.usd.is_hidden_type(prim)]
        self.assertEqual(prims, ['/Root', '/Root/Cone'])

        # close stage & delete file
        await omni.usd.get_context().new_stage_async()


    async def test_l1_stage_menu_export_selected_multiple_prims(self):
        stage = omni.usd.get_context().get_stage()
        to_select = ["/World/Cone", "/World/Cube", "/World/Sphere", "/World/Cylinder"]
        test_file = os.path.join(omni.kit.test.get_test_output_path(), "save_multiple_selected_prim.usd")

        # select prims
        await select_prims(to_select)

        # right click on Cube
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

        # click on context menu item & save
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            await file_export_helper.click_apply_async(filename_url=test_file)
            await wait_stage_loading()

        # load created file
        await open_stage(test_file)
        await wait_stage_loading()

        # verify prims
        stage = omni.usd.get_context().get_stage()
        prims = [prim.GetPath().pathString for prim in stage.TraverseAll() if not omni.usd.is_hidden_type(prim)]
        self.assertEqual(prims, ['/Root', '/Root/Cone', '/Root/Cone/Cone', '/Root/Cube', '/Root/Cube/Cube', '/Root/Sphere', '/Root/Sphere/Sphere', '/Root/Cylinder', '/Root/Cylinder/Cylinder'])

        # close stage & delete file
        await omni.usd.get_context().new_stage_async()


    async def test_l1_stage_menu_export_selected_single_material(self):
        stage = omni.usd.get_context().get_stage()
        to_select = ["/World/Looks/OmniPBR"]
        test_file = os.path.join(omni.kit.test.get_test_output_path(), "save_single_selected_material.usd")

        # select prims
        await select_prims(to_select)

        # right click on Cube
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

        # click on context menu item & save
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            await file_export_helper.click_apply_async(filename_url=test_file)
            await wait_stage_loading()

        # load created file
        await open_stage(test_file)
        await wait_stage_loading()

        # verify prims
        stage = omni.usd.get_context().get_stage()
        prims = [prim.GetPath().pathString for prim in stage.TraverseAll() if not omni.usd.is_hidden_type(prim)]
        self.assertEqual(prims, ['/Root', '/Root/OmniPBR', '/Root/OmniPBR/Shader'])

        # close stage & delete file
        await omni.usd.get_context().new_stage_async()


    async def test_l1_stage_menu_export_selected_multiple_materials(self):
        stage = omni.usd.get_context().get_stage()
        to_select = ["/World/Looks/OmniPBR", "/World/Looks/OmniGlass", "/World/Looks/OmniSurface_Plastic"]
        test_file = os.path.join(omni.kit.test.get_test_output_path(), "save_multiple_selected_material.usd")

        # select prims
        await select_prims(to_select)

        # right click on Cube
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

        # click on context menu item & save
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            await file_export_helper.click_apply_async(filename_url=test_file)
            await wait_stage_loading()

        # load created file
        await open_stage(test_file)
        await wait_stage_loading()

        # verify prims
        stage = omni.usd.get_context().get_stage()
        prims = [prim.GetPath().pathString for prim in stage.TraverseAll() if not omni.usd.is_hidden_type(prim)]
        self.assertEqual(prims, ['/Root', '/Root/OmniPBR', '/Root/OmniPBR/Shader', '/Root/OmniGlass', '/Root/OmniGlass/Shader', '/Root/OmniSurface_Plastic', '/Root/OmniSurface_Plastic/Shader'])

        # close stage & delete file
        await omni.usd.get_context().new_stage_async()

    async def test_stage_menu_export_selected_prefill(self):
        """Test that export selected pre-fill prim name and save directory."""
        stage = omni.usd.get_context().get_stage()
        to_select = ["/World/Cone", "/World/Cube", "/World/Sphere", "/World/Cylinder"]
        output_dir = omni.kit.test.get_test_output_path()
        test_file = os.path.join(output_dir, "test_prefill.usd")

        # select prims
        await select_prims(to_select)

        # right click on Cube
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

        # click on context menu item & save
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            # check that prim name is pre-fill with the first prim name
            exporter = get_file_exporter()
            filename = exporter._dialog.get_filename()
            self.assertEqual(filename, "Cone")
            # check that current directory is the same as the stage directory
            dirname = exporter._dialog.get_current_directory()
            self.assertEqual(dirname.rstrip("/").lower(),
                os.path.dirname(stage.GetRootLayer().realPath).replace("\\", "/").lower())
            await file_export_helper.click_apply_async(filename_url=test_file)
            await wait_stage_loading()

        # right click again, now the default directory for save should update to the last save directory
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            exporter = get_file_exporter()
            # check that current directory is the same as the stage directory
            dirname = exporter._dialog.get_current_directory()
            self.assertEqual(dirname.rstrip("/").lower(), output_dir.replace("\\", "/").lower())
            await file_export_helper.click_cancel_async()

        # load another file
        output_dir = get_test_data_path(__name__, "usd")
        await open_stage(os.path.join(output_dir, "cube.usda"))
        await wait_stage_loading()
        stage = omni.usd.get_context().get_stage()

        # check that current directory is the same as the stage directory
        to_select = ["/World/Cube"]
        await select_prims(to_select)

        # right click on Cube
        stage_widget = ui_test.find("Stage//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{to_select[0]}'").right_click()

        # click on context menu item & save
        async with FileExporterTestHelper() as file_export_helper:
            await ui_test.select_context_menu("Save Selected")
            await file_export_helper.wait_for_popup()
            # check that prim name is pre-fill with the selected prim name
            exporter = get_file_exporter()
            filename = exporter._dialog.get_filename()
            self.assertEqual(filename, "Cube")
            # check that current directory is the same as the stage directory
            dirname = exporter._dialog.get_current_directory()
            self.assertEqual(dirname.rstrip("/").lower(),
                os.path.dirname(stage.GetRootLayer().realPath).replace("\\", "/").lower())
            await file_export_helper.click_cancel_async()

        # close stage & delete file
        await omni.usd.get_context().new_stage_async()