import os
import asyncio
import carb
import omni.kit.test
import omni.usd
import omni.client
import omni.kit.commands

from pathlib import Path
from omni.kit.tool.collect import get_instance
from omni.kit.tool.collect.collector import (
    Collector,
    CollectorException,
    CollectorFailureOptions,
    FlatCollectionTextureOptions,
)
from omni.kit.tool.collect.utils import Utils
from pxr import Usd


# NOTE: those tests belong to omni.kit.tool.collect extension.
class TestCollect(omni.kit.test.AsyncTestCase):
    def list_folder(self, folder_path):
        all_file_names = []
        result, entry = omni.client.stat(folder_path)
        if result == omni.client.Result.OK and entry.flags & omni.client.ItemFlags.CAN_HAVE_CHILDREN:
            is_folder = True
        else:
            is_folder = False

        if not is_folder:
            all_file_names = [os.path.basename(folder_path)]
        else:
            folder_queue = [folder_path]
            while len(folder_queue) > 0:
                folder = folder_queue.pop(0)
                (result, entries) = omni.client.list(folder)
                if result != omni.client.Result.OK:
                    break
                folders = set((e.relative_path for e in entries if e.flags & omni.client.ItemFlags.CAN_HAVE_CHILDREN))
                for f in folders:
                    folder_queue.append(f"{folder}/{f}")
                files = set((e.relative_path for e in entries if not e.flags & omni.client.ItemFlags.CAN_HAVE_CHILDREN))
                for file in files:
                    all_file_names.append(os.path.basename(file))
        return all_file_names

    def get_test_dir(self):
        token = carb.tokens.get_tokens_interface()
        data_dir = token.resolve("${data}")
        if not data_dir.endswith("/"):
            data_dir += "/"

        data_dir = Utils.normalize_path(data_dir)

        return f"{data_dir}collect_tool_tests"

    async def setUp(self):
        pass

    async def tearDown(self):
        await omni.client.delete_async(self.get_test_dir())

    async def __test_internal(
        self,
        stage_name,
        root_usd,
        usd_only,
        material_only,
        flat_collection,
        texture_option=FlatCollectionTextureOptions.FLAT,
    ):
        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        extension_path = Path(extension_path)
        test_data_path = extension_path.joinpath("data")
        test_stage_dir = str(test_data_path.joinpath("test_stages").joinpath(stage_name))
        test_root_usd = test_stage_dir + "/" + root_usd
        collected_stage_dir = self.get_test_dir() + f"/collected_{stage_name}"

        collector = Collector(
            test_root_usd, collected_stage_dir, usd_only, flat_collection, material_only, texture_option=texture_option
        )
        await collector.collect(None, None)

        before = self.list_folder(test_stage_dir)
        after = self.list_folder(collected_stage_dir)
        self.assertTrue(len(after) > 0)

        if usd_only:
            before_filtered = []
            for f in before:
                if omni.usd.is_usd_writable_filetype(f):
                    before_filtered.append(f)
            before = before_filtered

        if material_only:
            before_filtered = []
            for f in before:
                if not omni.usd.is_usd_writable_filetype(f):
                    before_filtered.append(f)
            before = before_filtered

        self.assertEqual(set(before), set(after))

    async def test_collect_with_usd_and_material(self):
        await self.__test_internal("normal", "FullScene.usd", False, False, False)

    async def test_collect_with_udim_textures(self):
        await self.__test_internal("udim", "SM_Hood_A1_1.usd", False, False, False)

    async def test_collect_without_material(self):
        await self.__test_internal("normal", "FullScene.usd", False, True, False)

    async def test_collect_without_usd(self):
        await self.__test_internal("normal", "FullScene.usd", True, False, False)

    async def test_flatten_collection(self):
        await self.__test_internal("normal", "FullScene.usd", False, False, True)
        await self.__test_internal("layer_offsets", "root.usd", False, False, True)

    async def test_flatten_collection_texture_options(self):
        stage_name = "texture_options"
        texture_dir = self.get_test_dir() + f"/collected_{stage_name}/SubUSDs/textures/"
        mdl_name = "Contour1_Surface_0"
        mdl_texture_name = "Contour1_Surface_1.bmp"
        usd_preview_texture_name = "another_texture.bmp"

        # test group textures by MDL
        await self.__test_internal(
            stage_name, "test_scene.usd", False, False, True, texture_option=FlatCollectionTextureOptions.BY_MDL
        )
        # mdl texture should be under the mdl named folder, while the usd preview texture should be directly under
        #  textures dir
        self.assertTrue(os.path.exists(os.path.join(texture_dir, mdl_name, mdl_texture_name)))
        self.assertTrue(os.path.exists(os.path.join(texture_dir, usd_preview_texture_name)))
        await omni.client.delete_async(self.get_test_dir())

        # test group textures by USD
        await self.__test_internal(
            stage_name, "test_scene.usd", False, False, True, texture_option=FlatCollectionTextureOptions.BY_USD
        )
        # usd preview texture should be under the usd asset named folder, while the mdl texture should be directly under
        #  textures dir
        self.assertTrue(os.path.exists(os.path.join(texture_dir, "assetB", usd_preview_texture_name)))
        self.assertTrue(os.path.exists(os.path.join(texture_dir, mdl_texture_name)))
        await omni.client.delete_async(self.get_test_dir())

        # test flat
        await self.__test_internal(
            stage_name, "test_scene.usd", False, False, True, texture_option=FlatCollectionTextureOptions.FLAT
        )
        # all textures should be directly under textures dir
        self.assertEqual(set([mdl_texture_name, usd_preview_texture_name]), set(os.listdir(texture_dir)))

    async def test_layer_offsets_collect(self):
        current_path = Path(__file__).parent
        test_data_path = current_path.parent.parent.parent.parent.parent.joinpath("data")
        test_stage_dir = str(test_data_path.joinpath("test_stages").joinpath("layer_offsets"))
        test_root_usd = test_stage_dir + "/root.usd"
        collected_stage_dir = self.get_test_dir() + f"/collected_layer_offfsets"
        collector = Collector(test_root_usd, collected_stage_dir, False, False, False)
        await collector.collect(None, None)

        before = self.list_folder(test_stage_dir)
        after = self.list_folder(collected_stage_dir)
        self.assertEqual(set(before), set(after))

        after_stage_usd = collected_stage_dir + "/root.usd"
        before_stage = Usd.Stage.Open(test_root_usd)
        after_stage = Usd.Stage.Open(after_stage_usd)
        before_root = before_stage.GetRootLayer()
        after_root = after_stage.GetRootLayer()
        before_sublayers = before_root.subLayerPaths
        after_sublayers = after_root.subLayerPaths
        self.assertTrue(len(before_sublayers) > 0)
        self.assertEqual(len(before_sublayers), len(after_sublayers))
        self.assertEqual(len(before_sublayers), len(after_sublayers))

        for i in range(len(before_sublayers)):
            before_sublayer = before_root.ComputeAbsolutePath(before_sublayers[i])
            after_sublayer = after_root.ComputeAbsolutePath(after_sublayers[i])
            self.assertNotEqual(before_sublayer, after_sublayer)

        self.assertEqual(before_root.subLayerOffsets, after_root.subLayerOffsets)

    async def test_collect_failure_with_exception(self):
        with self.assertRaises(CollectorException):
            collector = Collector("invalid_source.usd", self.get_test_dir(), False, False, False)
            await collector.collect(None, None)

        with self.assertRaises(CollectorException):
            current_path = Path(__file__).parent
            test_data_path = current_path.parent.parent.parent.parent.parent.joinpath("data")
            test_stage_path = str(test_data_path.joinpath("test_stages").joinpath("normal/FullScene.usd"))
            collector = Collector(test_stage_path, "omniverse://invalid_path", False, False, False)
            await collector.collect(None, None)

        usd_error_option = CollectorFailureOptions.EXTERNAL_USD_REFERENCES
        other_error_option = CollectorFailureOptions.OTHER_EXTERNAL_REFERENCES
        all_option = CollectorFailureOptions.EXTERNAL_USD_REFERENCES | CollectorFailureOptions.OTHER_EXTERNAL_REFERENCES
        for option in [usd_error_option, other_error_option, all_option]:
            with self.assertRaises(CollectorException):
                current_path = Path(__file__).parent
                test_data_path = current_path.parent.parent.parent.parent.parent.joinpath("data")
                test_stage_path = str(test_data_path.joinpath("test_stages/normal").joinpath("FullScene.usd"))
                collected_stage_dir = self.get_test_dir() + "/collected_normal"
                collector = Collector(test_stage_path, collected_stage_dir, False, False, False, failure_options=option)
                await collector.collect(None, None)

    async def test_utils(self):
        self.assertTrue(Utils.is_udim_texture("test.<UDIM>.png"))
        self.assertTrue(Utils.is_udim_texture("test_<UDIM>_suffix.png"))
        self.assertTrue(Utils.is_udim_texture("test.%3cUDIM%3e.png"))
        self.assertFalse(Utils.is_udim_texture(""))
        self.assertFalse(Utils.is_udim_texture("random"))
        self.assertTrue(Utils.is_udim_wildcard_texture("test_0001_suffix.png", "test_<UDIM>_suffix.png"))
        self.assertFalse(Utils.is_udim_wildcard_texture("test_0001_another_suffix.png", "test_<UDIM>_suffix.png"))
        self.assertTrue(
            Utils.is_udim_wildcard_texture(
                "omniverse://fake-server/base_path/test_0001_suffix.png",
                "omniverse://fake-server/base_path/test_<UDIM>_suffix.png",
            )
        )
        self.assertFalse(
            Utils.is_udim_wildcard_texture(
                "omniverse://fake-server/base_path/test_0001_another_suffix.png",
                "omniverse://fake-server/base_path/test_<UDIM>_suffix.png",
            )
        )

    async def test_path_calculation(self):
        # Tests for https://nvidia-omniverse.atlassian.net/browse/OM-34746

        current_path = Path(__file__).parent
        test_data_path = current_path.parent.parent.parent.parent.parent.joinpath("data")
        test_stage_path = str(test_data_path.joinpath("test_stages").joinpath("normal/FullScene.usd"))
        collector = Collector(test_stage_path, self.get_test_dir(), False, False, False)
        path = collector._calculate_target_path("http://test_server/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/test_server/testfile.png")

        path = collector._calculate_target_path(str(test_data_path) + "/deep_folder/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/deep_folder/testfile.png")

        path = collector._calculate_target_path("omniverse://test_server/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/test_server/testfile.png")

        # Flat collect
        collector._flat_collection = True
        collector._material_only = False
        path = collector._calculate_target_path("omniverse://test_server/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/SubUSDs/textures/testfile.png")

        path = collector._calculate_target_path("http://test_server/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/SubUSDs/textures/testfile.png")

        path = collector._calculate_target_path("http://test_server/testfile.usd")
        self.assertEqual(path, self.get_test_dir() + "/SubUSDs/testfile.usd")

        # Materials only with flat collect
        collector._flat_collection = True
        collector._material_only = True
        path = collector._calculate_target_path("omniverse://test_server/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/textures/testfile.png")

        path = collector._calculate_target_path("http://test_server/testfile.png")
        self.assertEqual(path, self.get_test_dir() + "/textures/testfile.png")

        # Master USD with flat collect
        collector._flat_collection = True
        collector._material_only = False
        path = collector._calculate_target_path(test_stage_path)
        self.assertEqual(path, self.get_test_dir() + "/FullScene.usd")

    async def __wait(self, frames=2):
        for _ in range(frames):
            await omni.kit.app.get_app().next_update_async()

    async def __test_ui_internal(self, stage_name, root_usd, usd_only, material_only, flat_collection):
        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        extension_path = Path(extension_path)
        test_data_path = extension_path.joinpath("data")
        test_stage_dir = str(test_data_path.joinpath("test_stages").joinpath(stage_name))
        test_root_usd = test_stage_dir + "/" + root_usd
        collected_stage_dir = self.get_test_dir() + f"/collected_{stage_name}"

        done = False

        def on_finish():
            nonlocal done
            done = True

        collect_extension = get_instance()
        collect_extension.collect(test_root_usd, on_finish)
        await self.__wait()

        from omni.kit import ui_test

        collect_window = ui_test.find("Collection Options")
        await collect_window.focus()

        usd_only_checkbox = collect_window.find("**/CheckBox[*].identifier=='usd_only_checkbox'")
        material_only_checkbox = collect_window.find("**/CheckBox[*].identifier=='material_only_checkbox'")
        flat_collection_checkbox = collect_window.find("**/CheckBox[*].identifier=='flat_collection_checkbox'")
        folder_button = collect_window.find("**/Button[*].identifier=='folder_button'")
        path_field = collect_window.find("**/StringField[*].identifier=='collect_path'")
        start_button = collect_window.find("**/Button[*].text=='Start'")
        cancel_button = collect_window.find("**/Button[*].text=='Cancel'")
        option_combo = collect_window.find("**/ComboBox[*].identifier=='texture_option_combo'")
        self.assertTrue(usd_only_checkbox)
        self.assertTrue(material_only_checkbox)
        self.assertTrue(flat_collection_checkbox)
        self.assertTrue(folder_button)
        self.assertTrue(path_field)
        self.assertTrue(start_button)
        self.assertTrue(cancel_button)
        self.assertTrue(option_combo)

        usd_only_checkbox.model.set_value(False)
        material_only_checkbox.model.set_value(False)
        flat_collection_checkbox.model.set_value(False)

        if usd_only:
            await usd_only_checkbox.click()

        if material_only:
            await material_only_checkbox.click()

        if flat_collection:
            self.assertFalse(option_combo.widget.visible)
            await flat_collection_checkbox.click()

            # should show texture options combo here
            self.assertTrue(option_combo.widget.visible)

        await folder_button.click()

        file_picker = ui_test.find("Select Collect Destination")
        self.assertTrue(file_picker)
        await file_picker.focus()

        open_button = file_picker.find("**/Button[*].text=='Open'")
        cancel_button = file_picker.find("**/Button[*].text=='Cancel'")
        self.assertTrue(open_button)
        self.assertTrue(cancel_button)
        await cancel_button.click()

        path_field.model.set_value(collected_stage_dir)
        await start_button.click()

        await self.__wait()

        progress_window = ui_test.find("Collecting")
        self.assertTrue(progress_window)
        while not done:
            await self.__wait()

        self.assertFalse(progress_window.window.visible)

        before = self.list_folder(test_stage_dir)
        after = self.list_folder(collected_stage_dir)
        self.assertTrue(len(after) > 0)

        if usd_only:
            before_filtered = []
            for f in before:
                if omni.usd.is_usd_writable_filetype(f):
                    before_filtered.append(f)
            before = before_filtered

        if material_only:
            before_filtered = []
            for f in before:
                if not omni.usd.is_usd_writable_filetype(f):
                    before_filtered.append(f)
            before = before_filtered

        self.assertEqual(set(before), set(after))
        self.assertTrue(done)

    async def test_ui_collect_with_usd_and_material(self):
        await self.__test_ui_internal("normal", "FullScene.usd", False, False, False)

    async def test_ui_collect_without_material(self):
        await self.__test_ui_internal("normal", "FullScene.usd", False, True, False)

    async def test_ui_collect_without_usd(self):
        await self.__test_ui_internal("normal", "FullScene.usd", True, False, False)
    
    async def test_om_55150(self):
        await self.__test_ui_internal("OM_55150", "collect_donut.usd", False, False, False)

    async def test_ui_flatten_collection(self):
        await self.__test_ui_internal("normal", "FullScene.usd", False, False, True)

    async def test_file_menu(self):
        await omni.usd.get_context().new_stage_async()

        from omni.kit import ui_test

        await ui_test.menu_click("File/Collect As...")

        # It cannot collect anonymous stage
        collect_window = ui_test.find("Collection Options")
        self.assertFalse(collect_window)

        extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
        extension_path = Path(extension_path)
        test_data_path = extension_path.joinpath("data")
        test_stage_dir = str(test_data_path.joinpath("test_stages").joinpath("udim"))
        test_root_usd = test_stage_dir + "/SM_Hood_A1_1.usd"

        await omni.usd.get_context().open_stage_async(test_root_usd)
        await ui_test.menu_click("File/Collect As...")

        # It cannot collect anonymous stage
        collect_window = ui_test.find("Collection Options")
        self.assertFalse(collect_window)
