## 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 carb
import os
import tempfile
import shutil
import omni.kit.app
import omni.kit.commands
import omni.kit.test

from pathlib import Path
from unittest.mock import Mock, patch
from omni.kit.window.file_exporter.test_helper import FileExporterTestHelper
from .. import FileWindowExtension, get_instance as get_window_file
from .test_base import TestFileBase
from pxr import Usd
import unittest


class TestFileSave(TestFileBase):
    async def setUp(self):
        await super().setUp()

    async def tearDown(self):
        await super().tearDown()

    async def test_file_save(self):
        """Testing that save should write to stage url"""
        omni.usd.get_context().new_stage()
        omni.usd.get_context().set_pending_edit(False)

        test_file = "test_file.usda"
        temp_path = os.path.join(tempfile.gettempdir(), test_file)
        shutil.copyfile(os.path.join(self._usd_path, test_file), temp_path)

        omni.kit.window.file.open_stage(temp_path)
        await self.wait_for_update()

        mock_callback = Mock()
        omni.kit.window.file.save(mock_callback)
        await self.wait_for_update()

        self.assertEqual(omni.usd.get_context().get_stage_url(), temp_path.replace("\\", "/"))
        self.assertTrue(os.path.isfile(temp_path))
        mock_callback.assert_called_once()
        self.remove_file(temp_path)

    async def test_file_save_with_render_setting_changes_only(self):
        usd_context = omni.usd.get_context()
        await usd_context.new_stage_async()
        self.assertFalse(usd_context.has_pending_edit())

        with tempfile.TemporaryDirectory() as tmp_dir_name:
            await omni.kit.app.get_app().next_update_async()

            # save
            tmp_file_path = Path(tmp_dir_name) / "tmp.usda"
            result = usd_context.save_as_stage(str(tmp_file_path))
            self.assertTrue(result)
            self.assertFalse(usd_context.has_pending_edit())

            # save the file with customized key value
            settings = carb.settings.get_settings()
            setting_key_ao = "/rtx/post/aa/op"
            settings.set_int(setting_key_ao, 3)
            await omni.kit.app.get_app().next_update_async()

            # After changing render settings, it will be dirty.
            self.assertTrue(usd_context.has_pending_edit())

            # Saves stage, although no layers are dirty except render settings.
            omni.kit.window.file.save(None)

            # Waiting for save to be done as it's asynchronous.
            for _ in range(10):
                await omni.kit.app.get_app().next_update_async()

            # Making sure that render settings are saved into root layer.
            self.assertFalse(usd_context.has_pending_edit())

            # Reloading stage will reload the settings also.
            await usd_context.reopen_stage_async()
            await omni.kit.app.get_app().next_update_async()
            self.assertFalse(usd_context.has_pending_edit())

            setting_value_ao = settings.get(setting_key_ao)

            self.assertEqual(setting_value_ao, 3)

            # New stage to release temp file.
            await usd_context.new_stage_async()

    async def test_file_save_read_only_file(self):
        """Test saving read-only file calls save_as instead"""
        omni.usd.get_context().new_stage()
        omni.usd.get_context().set_pending_edit(False)

        test_file = "test_file.usda"
        temp_path = os.path.join(tempfile.gettempdir(), test_file)
        shutil.copyfile(os.path.join(self._usd_path, test_file), temp_path)

        omni.kit.window.file.open_stage(temp_path)
        await self.wait_for_update()

        class ListEntry:
            def __init__(self) -> None:
                self.flags = 0

        mock_callback = Mock()
        with patch.object(omni.usd.UsdContext, 'is_new_stage', return_value=False),\
            patch.object(omni.client, "stat_async", return_value=(omni.client.Result.OK, ListEntry())),\
            patch.object(FileWindowExtension, "save_as") as mock_save_as:
                omni.kit.window.file.save(mock_callback)
                await self.wait_for_update()

        mock_save_as.assert_called_once_with(False, mock_callback, allow_skip_sublayers=False)
        self.remove_file(temp_path)

    async def test_file_saveas(self):
        """Testing save as"""
        omni.usd.get_context().new_stage()
        omni.usd.get_context().set_pending_edit(False)

        temp_path = self._temp_path.replace('.usda', '.usd')
        mock_callback = Mock()

        async with FileExporterTestHelper() as file_export_helper:
            with patch.object(omni.usd.UsdContext, 'is_new_stage', return_value=False):
                omni.kit.window.file.save_as(False, mock_callback)
                await self.wait_for_update()

                await file_export_helper.click_apply_async(filename_url=temp_path)
                await self.wait_for_update()

        mock_callback.assert_called_once()
        self.assertTrue(os.path.isfile(temp_path))
        self.remove_file(temp_path)

    async def _test_file_saveas_with_empty_filename(self): # pragma: no cover
        """Testing save as with an empty filename, should not allow saving."""
        omni.usd.get_context().new_stage()
        omni.usd.get_context().set_pending_edit(False)

        temp_path = self._temp_path.replace('.usda', '.usd')
        mock_callback = Mock()

        async with FileExporterTestHelper() as file_export_helper:
            with patch.object(omni.usd.UsdContext, 'is_new_stage', return_value=False):
                omni.kit.window.file.save_as(False, mock_callback)
                await self.wait_for_update()

                await file_export_helper.click_apply_async()
                await self.wait_for_update()

        mock_callback.assert_not_called()
        self.assertFalse(os.path.isfile(temp_path))

    async def _test_file_saveas_flattened_internal(self, include_session_layer):
        """Testing save as flatten"""
        omni.usd.get_context().new_stage()
        omni.usd.get_context().set_pending_edit(False)

        mock_callback = Mock()

        base, ext = os.path.splitext(self._temp_path)
        temp_path = base + str(include_session_layer) + ext
        temp_path = self._temp_path.replace('.usda', '.usd')
        async with FileExporterTestHelper() as file_export_helper:
            with patch.object(omni.usd.UsdContext, 'is_new_stage', return_value=False):
                omni.kit.window.file.save_as(True, mock_callback)
                await self.wait_for_update()

                import omni.kit.ui_test as ui_test
                file_picker = ui_test.find("Save File As...")
                self.assertTrue(file_picker)
                await file_picker.focus()

                include_session_layer_checkbox = file_picker.find(
                    "**/CheckBox[*].identifier=='include_session_layer'"
                )
                self.assertTrue(include_session_layer_checkbox)
                include_session_layer_checkbox.model.set_value(include_session_layer)

                await file_export_helper.click_apply_async(filename_url=temp_path)
                await self.wait_for_update()

        stage = Usd.Stage.Open(temp_path)
        prim = stage.GetPrimAtPath("/OmniverseKit_Persp")
        if include_session_layer:
            self.assertTrue(prim)
        else:
            self.assertFalse(prim)
        stage = None

        mock_callback.assert_called_once()
        self.assertTrue(os.path.isfile(temp_path))
        self.remove_file(temp_path)

    @unittest.skip("A.B. temp disable, omni.physx.bundle extension causes issues")
    async def test_file_saveas_flattened_with_session_layer(self):
        await self._test_file_saveas_flattened_internal(True)

    async def test_file_saveas_flattened_without_session_layer(self):
        await self._test_file_saveas_flattened_internal(False)

    async def test_show_save_layers_failed_prompt(self):
        omni.usd.get_context().new_stage()
        prompt = get_window_file().ui_handler.show_save_layers_failed_prompt()
        await self.wait_for_update()
        prompt.hide()

    async def test_file_saveas_existed_filename(self):
        """Testing save as with a filename which is already existed, should not allow saving."""
        # create an empty stage
        await omni.usd.get_context().new_stage_async()
        omni.usd.get_context().set_pending_edit(False)

        file_name = f"4Lights.usda"
        temp_path = os.path.join(tempfile.gettempdir(), file_name)
        shutil.copyfile(os.path.join(self._usd_path, file_name), temp_path)

        # save the current stage to the temp path where the file name already exists
        async with FileExporterTestHelper() as file_export_helper:
            with patch.object(omni.usd.UsdContext, 'is_new_stage', return_value=False):
                omni.kit.window.file.save_as(False, Mock())
                await self.wait_for_update()

                await file_export_helper.click_apply_async(filename_url=temp_path)
                await self.wait_for_update()

        # check 4Lights.usda is not empty, which means saveas didn't happen
        stage = Usd.Stage.Open(str(temp_path))
        await self.wait_for_update()

        prim = stage.GetPrimAtPath("/Stage/SphereLight_01")
        self.assertTrue(prim)
        self.remove_file(temp_path)

    async def test_invalid_save(self):
        """Test no valid stage save"""
        if omni.usd.get_context().get_stage():
            await omni.usd.get_context().close_stage_async()
        mock_callback = Mock()
        omni.kit.window.file.save(mock_callback)
        mock_callback.assert_not_called()
