import os
import asyncio
import omni
import omni.usd
import carb
import traceback

from enum import IntFlag, Enum
from pxr import Sdf, Usd, UsdUtils, UsdShade, UsdLux, Tf
from .omni_client_wrapper import OmniClientWrapper
from .utils import Utils
from .async_utils import aio_open_layer, aio_replace_all, aio_save_layer
from .mdl_parser import MDLParser
from .async_utils import wrap


class CollectorFailureOptions(IntFlag):
    """Options to customize failure options"""

    SILENT = 0  # Silent for all failures except root USD.
    EXTERNAL_USD_REFERENCES = 1  # Throws exception if any external USD file is not found.
    OTHER_EXTERNAL_REFERENCES = 4  # Throws exception if external references other than all above are missing.


class CollectorException(Exception):
    def __init__(self, error: str):
        self._error = error

    def __str__(self):
        return self._error


class CollectorTaskType:
    """Task type"""

    READ_TASK = 0
    WRITE_TASK = 1
    COPY_TASK = 2
    RESOLVE_TASK = 3


class FlatCollectionTextureOptions(Enum):
    """Collection options for textures under 'Flat Collection' mode"""

    BY_MDL = 0  # group textures by MDL
    BY_USD = 1  # group textures by USD
    FLAT = 2  # all textures will be under the same hierarchy, flat


class Collector:
    """Collect stage related asset  to {target_folder}/collected_stagename/}.

    Args:
        usd_path (str): The usd stage to be collected.

        collect_dir (str): The target dir to collect the usd stage to.

        usd_only (bool): Collects usd files only or not. It will ignore all asset types.

        flat_collection (bool): Collects stage without keeping the original dir structure.

        material_only (bool): Collects material and textures only or not. It will ignore all other asset types.

        skip_existing (bool): If files already exist in the target location, don't copy again. Be careful of corrupt files

        texture_option (FlatCollectionTextureOptions): Specifies how textures are grouped in flat collection mode.
            This is to avoid name collision which results in textures overwriting each other in the cases where textures
            for different assets are not uniquely named.

        If both `usd_only` and `material_only` are true, it will collect all.
    """

    def __init__(
        self,
        usd_path: str,
        collect_dir: str,
        usd_only: bool,
        flat_collection: bool,
        material_only: bool = False,
        failure_options=CollectorFailureOptions.SILENT,
        skip_existing: bool = False,
        max_concurrent_tasks=64,
        texture_option=FlatCollectionTextureOptions.BY_MDL,
    ):
        self._usd_path = Utils.normalize_path(usd_path)
        self._collect_dir = Utils.normalize_path(collect_dir)
        if not self._collect_dir.endswith("/"):
            self._collect_dir += "/"
        self._usd_only = usd_only and not material_only
        self._material_only = material_only and not usd_only
        self._flat_collection = flat_collection
        self._texture_option = texture_option
        self._finished = False
        self._cancelled = False
        self.MAX_CONCURRENT_TASKS = max_concurrent_tasks
        self._current_tasks = set([])
        self._finished_callback = None
        self._progress_callback = None
        self._unique_path_name = {}  # Records the count key that appears in the collected paths to generate unique name
        self._source_target_path_mapping = {}  # It records source to target
        self._failure_options = failure_options
        self._skip_existing = skip_existing
        self._layer_store = {}
        self._mdl_parsers = {}
        self._sublayer_offsets = {}
        self._other_asset_paths = set({})
        self._current_progress = 0
        self._total_steps = 0
        self._all_read_paths = set({})

    def destroy(self):
        self._finished_callback = None
        self._progress_callback = None

    def _caculate_flat_target_path(self, path, parent_asset=None):
        target_dir = self._collect_dir
        master_usd = Utils.normalize_path(path) == self._usd_path
        material_only = self._material_only

        file_url = omni.client.break_url(path)
        target_url = omni.client.break_url(target_dir)
        file_path = file_url.path
        file_name = os.path.basename(file_path)

        if omni.usd.is_usd_writable_filetype(path) and master_usd:
            int_dir = ""
        else:
            if not material_only:
                int_dir = "SubUSDs/"
            else:
                int_dir = ""

            if Utils.is_material(file_name):
                int_dir = f"{int_dir}materials/"
            elif not omni.usd.is_usd_writable_filetype(file_name):
                int_dir = f"{int_dir}textures/"
                # OM-52799 add options to group textures by MDL / USD assets to avoid name collision causing overwrites
                if parent_asset:
                    parent_name, _ = os.path.splitext(os.path.basename(parent_asset))
                    # currently not including the file extension in parent folder names, but could add it  in if it is
                    #  more clear to show the ext as a suffix to the folder name
                    int_dir += f"{parent_name}/"

        target_path = target_url.path + int_dir + file_name
        if file_url.query:
            branch_checkpoint = omni.client.get_branch_and_checkpoint_from_query(file_url.query)
            if branch_checkpoint:
                branch, checkpoint = branch_checkpoint
                file_path, ext = os.path.splitext(target_path)
                if not branch:
                    branch = "default"
                target_path = f"{file_path}__{branch}__v{checkpoint}{ext}"

        target_path = omni.client.make_url(
            scheme=target_url.scheme,
            user=target_url.user,
            host=target_url.host,
            port=target_url.port,
            path=target_path,
            query=None,
            fragment=None,
        )
        return Utils.normalize_path(target_path)

    # All paths must be absolute path
    def _calculate_target_path(self, file_path, parent_asset=None):
        file_path = Utils.normalize_path(file_path)

        if self._flat_collection:
            return self._caculate_flat_target_path(file_path, parent_asset=parent_asset)

        target_dir = self._collect_dir
        if not target_dir.endswith("/"):
            target_dir += "/"

        stage_path = self._usd_path
        stage_url = omni.client.break_url(stage_path)
        if not stage_url.scheme:
            stage_path_with_scheme = "file://" + stage_path
        else:
            stage_path_with_scheme = stage_path

        file_url = omni.client.break_url(file_path)
        if not file_url.scheme:
            file_path_with_shceme = "file://" + file_path
        else:
            file_path_with_shceme = file_path

        stage_url = omni.client.break_url(stage_path_with_scheme)
        file_url = omni.client.break_url(file_path_with_shceme)
        target_url = omni.client.break_url(target_dir)

        if stage_url.scheme != file_url.scheme or stage_url.host != file_url.host:
            file_path = file_url.path.lstrip("/")

            if file_url.host:
                target_path = target_url.path + file_url.host + "/" + file_path
            else:
                target_path = target_url.path + file_path
        else:
            stage_url_path = stage_url.path.replace("//", "/")
            file_url_path = file_url.path.replace("//", "/")
            common_path = os.path.commonpath([os.path.dirname(stage_url_path), os.path.dirname(file_url_path)])
            common_path = common_path.replace("\\", "/")
            file_path = file_url_path[len(common_path) :].lstrip("/")
            target_path = target_url.path + file_path

        if file_url.query:
            branch_checkpoint = omni.client.get_branch_and_checkpoint_from_query(file_url.query)
            if branch_checkpoint:
                branch, checkpoint = branch_checkpoint
                file_name, ext = os.path.splitext(target_path)
                if not branch:
                    branch = "default"
                target_path = f"{file_name}__{branch}__v{checkpoint}{ext}"

        target_path = omni.client.make_url(
            scheme=target_url.scheme,
            user=target_url.user,
            host=target_url.host,
            port=target_url.port,
            path=target_path,
            query=None,
            fragment=None,
        )

        return Utils.normalize_path(target_path)

    async def _get_total_steps(self):
        stage_path = self._usd_path
        collect_usd_only = self._usd_only

        carb.log_info(f"Downloading and collecting dependencies: {stage_path}...")

        # Pre-warming to download all USDs and MDLs in parallel.
        target_path = self._calculate_target_path(stage_path)
        self._source_target_path_mapping[stage_path] = target_path

        if target_path and target_path == stage_path:
            self._raise_or_log(
                stage_path, f"Failed to collect layer {stage_path} as it tries to overwrite itself.", True
            )

        await self.__add_read_task(stage_path)

        # Cancelled
        if not await self.wait_all_unfinished_tasks(True):
            return None

        if collect_usd_only:
            total = len(self._layer_store)
        else:
            total_textures = 0
            for parser in self._mdl_parsers.values():
                raw_to_absolute_texture_paths, _ = await parser.parse()
                total_textures += len(raw_to_absolute_texture_paths)

            total_material_assets = len(self._mdl_parsers) + total_textures

            if self._material_only:
                total = len(self._other_asset_paths) + total_material_assets
            else:
                total = len(self._layer_store) + len(self._other_asset_paths) + total_material_assets

        return total

    def _remove_prim_spec(self, layer: Sdf.Layer, prim_spec_path: str):
        prim_spec = layer.GetPrimAtPath(prim_spec_path)
        if not prim_spec:
            return False

        if prim_spec.nameParent:
            name_parent = prim_spec.nameParent
        else:
            name_parent = layer.pseudoRoot

        if not name_parent:
            return False

        name = prim_spec.name
        if name in name_parent.nameChildren:
            del name_parent.nameChildren[name]

    def _remove_all_materials_and_bindings(self, stage_path):
        stage = Usd.Stage.Open(stage_path)
        if not stage:
            return

        to_be_removed = []
        for prim in stage.Traverse():
            if prim.IsA(UsdShade.Material):
                to_be_removed.append(prim)

            binding_api = UsdShade.MaterialBindingAPI(prim)
            if binding_api:
                binding_api.UnbindAllBindings()

            if prim.IsA(UsdLux.DomeLight):
                dome_light = UsdLux.DomeLight(prim)
                dome_light.GetTextureFileAttr().Set("")

        for prim in to_be_removed:
            to_remove_paths = []
            if prim:
                for prim_spec in prim.GetPrimStack():
                    layer = prim_spec.layer
                    to_remove_paths.append((layer, prim_spec.path))

            for item in to_remove_paths:
                self._remove_prim_spec(item[0], item[1])

        stage.Save()

    def __modify_external_references(self, original_layer_identifier, target_layer):
        def modifiy_paths_cb(path):
            absolute_path = Utils.compute_absolute_path(original_layer_identifier, path)
            target_path = self._source_target_path_mapping.get(absolute_path, None)
            if not target_path:
                return path

            return Utils.make_relative_path(target_layer.identifier, target_path)

        try:
            UsdUtils.ModifyAssetPaths(target_layer, modifiy_paths_cb)
        except Exception as e:
            carb.log_error(f"Failed to modify asset paths {target_layer.identifier}:" + str(e))

        LAYER_OMNI_CUSTOM_KEY = "omni_layer"
        LAYER_MUTENESS_CUSTOM_KEY = "muteness"
        custom_data = target_layer.customLayerData
        if LAYER_OMNI_CUSTOM_KEY in custom_data:
            omni_data = custom_data[LAYER_OMNI_CUSTOM_KEY]
            if LAYER_MUTENESS_CUSTOM_KEY in omni_data:
                new_muteness_data = {}
                for path, muted in omni_data[LAYER_MUTENESS_CUSTOM_KEY].items():
                    absolute_path = Utils.compute_absolute_path(original_layer_identifier, path)
                    target_path = self._source_target_path_mapping.get(absolute_path, None)
                    if target_path:
                        target_path = Utils.make_relative_path(target_layer.identifier, target_path)
                        new_muteness_data[target_path] = muted

                if new_muteness_data:
                    omni_data[LAYER_MUTENESS_CUSTOM_KEY] = new_muteness_data
                    custom_data[LAYER_OMNI_CUSTOM_KEY] = omni_data
                    target_layer.customLayerData = custom_data

    def __mapping_source_path(self, source_path, parent_asset=None):
        if source_path in self._source_target_path_mapping:
            return False
        # OM-52799 add options to group textures by MDL / USD assets to avoid name collision causing overwrites
        target_path = self._calculate_target_path(source_path, parent_asset=parent_asset)
        target_path = self._make_sure_unique_path(target_path)
        self._source_target_path_mapping[source_path] = target_path

        return True

    # Gets external references.
    async def __get_external_references(self, original_layer_identifier, target_layer):
        stage_path = self._usd_path

        try:
            sublayer_paths, reference_paths, payload_paths = UsdUtils.ExtractExternalReferences(target_layer.identifier)
        except Exception as e:
            carb.log_error(f"Failed to collect {target_layer.identifier}: " + str(e))
            sublayer_paths = []
            reference_paths = []
            payload_paths = []

        all_usd_paths = set()
        all_mdl_paths = set()
        other_asset_paths = set()

        async def track_path(path):
            absolute_path = Utils.compute_absolute_path(original_layer_identifier, path)
            if absolute_path in self._source_target_path_mapping:
                return

            if Utils.is_material(path):
                existed = await OmniClientWrapper.exists(absolute_path)
                # OM-44975: If it's not existed locally and it cannot be resolved with search paths,
                # just leave the path as it is.
                if not existed:
                    absolute_path = Utils.compute_absolute_path(stage_path, path)
                    # OM-43965: resolve agaist stage path.
                    existed = await OmniClientWrapper.exists(absolute_path)
                    if not existed:
                        return

            mapping_kwargs = {}
            if Utils.is_udim_texture(absolute_path):
                all_asset_paths = await self.__populate_udim_textures(absolute_path)
                for path in all_asset_paths:
                    self.__mapping_source_path(path, **mapping_kwargs)
            else:
                all_asset_paths = [absolute_path]

            if omni.usd.is_usd_writable_filetype(absolute_path):
                all_usd_paths.add(absolute_path)
            elif Utils.is_material(absolute_path):
                all_mdl_paths.add(absolute_path)
            else:
                # should only add parent asset option in mapping if the current path is not USD or MDL
                if self._flat_collection and self._texture_option is FlatCollectionTextureOptions.BY_USD:
                    # do not add the root stage itself as a parent
                    if original_layer_identifier != self._usd_path:
                        mapping_kwargs["parent_asset"] = original_layer_identifier
                # Adds udim mapping but not into copy list.
                if Utils.is_udim_texture(absolute_path):
                    all_asset_paths = await self.__populate_udim_textures(absolute_path)
                    for path in all_asset_paths:
                        self.__mapping_source_path(path, **mapping_kwargs)
                else:
                    all_asset_paths = [absolute_path]
                other_asset_paths.update(all_asset_paths)
            self.__mapping_source_path(absolute_path, **mapping_kwargs)

        async def check_paths(paths):
            tasks = []
            for path in paths:
                future = asyncio.ensure_future(track_path(path))
                tasks.append(future)

            if tasks:
                await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)

        await check_paths(sublayer_paths)
        await check_paths(reference_paths)
        await check_paths(payload_paths)

        return all_usd_paths, all_mdl_paths, other_asset_paths

    async def open_or_create_layer(self, layer_path, clear=True):
        # WA to fix OM-33212
        await omni.client.create_folder_async(os.path.dirname(layer_path))
        layer = await aio_open_layer(layer_path)
        if not layer:
            layer = Sdf.Layer.CreateNew(layer_path)
        elif clear:
            layer.Clear()

        return layer

    async def add_copy_task(self, source, target, skip_if_existed=False):
        if self._skip_existing and await OmniClientWrapper.exists(target):
            carb.log_info(f"Asset file {target} already exists, skipping copy")
        else:
            carb.log_info(f"Adding copy task from {source} to {target}")
            task = asyncio.ensure_future(OmniClientWrapper.copy(source, target, True))
            task.source = source
            task.target = target
            task.task_type = CollectorTaskType.COPY_TASK
            self._current_tasks.add(task)

        if len(self._current_tasks) >= self.MAX_CONCURRENT_TASKS:
            await self.wait_all_unfinished_tasks()

    async def add_write_task(self, target, content):
        carb.log_info(f"Adding write task to {target}")
        task = asyncio.ensure_future(OmniClientWrapper.write(target, content))
        task.task_type = CollectorTaskType.WRITE_TASK
        task.source = None
        task.target = target
        self._current_tasks.add(task)
        if len(self._current_tasks) >= self.MAX_CONCURRENT_TASKS:
            await self.wait_all_unfinished_tasks()

    async def __add_layer_resolve_task(self, source_layer_path, target_layer):
        carb.log_info(f"Adding layer resolve task for {target_layer.identifier}")
        task = asyncio.ensure_future(self.__resolve_layer(source_layer_path, target_layer))
        task.task_type = CollectorTaskType.RESOLVE_TASK
        task.source = target_layer.identifier
        task.target = None
        self._current_tasks.add(task)
        if len(self._current_tasks) >= self.MAX_CONCURRENT_TASKS:
            await self.wait_all_unfinished_tasks()

    async def __resolve_layer(self, source_layer_path, target_layer):
        # Resolve all external assets
        __aio_modify_external_references = wrap(self.__modify_external_references)

        await __aio_modify_external_references(source_layer_path, target_layer)

        # OM-40291: Copy sublayer offsets.
        sublayer_offsets = self._sublayer_offsets.get(source_layer_path, None)
        if sublayer_offsets:
            for i in range(len(sublayer_offsets)):
                target_layer.subLayerOffsets[i] = sublayer_offsets[i]

        return await aio_save_layer(target_layer)

    async def __open_and_analyze_layer(self, source_path, target_path):
        if self._material_only:
            target_layer = await aio_open_layer(source_path)
        else:
            # OM-52207: copy it firstly to keep tags for omniverse path.
            success = await OmniClientWrapper.copy(source_path, target_path, True)
            if not success:
                return False

            try:
                target_layer = await self.open_or_create_layer(target_path, False)
            except Tf.ErrorException:
                # WA for OM-52094: usda in server may be saved as usdc format,
                # which cannot be opened with USD library locally.
                # The solution here is to rename it as .usd.
                target_layer = None
                root, _ = os.path.splitext(target_path)
                renamed_usd = root + ".usd"
                success = await OmniClientWrapper.copy(target_path, renamed_usd)
                if success:
                    await OmniClientWrapper.delete(target_path)
                    target_layer = await self.open_or_create_layer(renamed_usd, False)
                    self._source_target_path_mapping[source_path] = renamed_usd

        if not target_layer:
            return False

        self._layer_store[source_path] = target_layer
        if target_layer.subLayerOffsets:
            self._sublayer_offsets[source_path] = target_layer.subLayerOffsets.copy()

        carb.log_info(f"Collecting dependencies for usd {source_path}...")
        all_usd_paths, all_mdl_paths, all_other_paths = await self.__get_external_references(source_path, target_layer)

        for path in all_usd_paths:
            if self._cancelled:
                break
            await self.__add_read_task(path)

        if self._usd_only:
            self._remove_all_materials_and_bindings(target_layer.identifier)
        else:
            for path in all_mdl_paths:
                if self._cancelled:
                    break
                await self.__add_read_task(path)

            self._other_asset_paths.update(all_other_paths)

        return True

    async def __populate_udim_textures(self, udim_texture_path):
        src_dir = os.path.dirname(udim_texture_path)
        if not src_dir.endswith("/"):
            src_dir = src_dir + "/"

        result, entries = await omni.client.list_async(src_dir)
        if result != omni.client.Result.OK:
            self._raise_or_log(
                udim_texture_path, f"Failed to list UDIM textures {udim_texture_path}, error code: {result}"
            )

        texture_paths = []
        for entry in entries:
            texture_path = Utils.compute_absolute_path(src_dir, entry.relative_path)
            if Utils.is_udim_wildcard_texture(texture_path, udim_texture_path):
                texture_paths.append(texture_path)

        return texture_paths

    async def __open_and_analyze_mdl(self, path, target_path):
        # OM-52207: copy it firstly to keep tags for omniverse path.
        if Utils.is_omniverse_path(path):
            success = await OmniClientWrapper.copy(path, target_path, True)
            if not success:
                return False

        content = await OmniClientWrapper.read(path)
        if not content:
            return False

        mdl_parser = MDLParser(path, content)
        self._mdl_parsers[path] = mdl_parser

        carb.log_info(f"Collecting dependencies for mdl {path}...")
        raw_to_absolute_texture_paths, imports = await mdl_parser.parse()
        mapping_kwargs = {}
        if self._flat_collection and self._texture_option is FlatCollectionTextureOptions.BY_MDL:
            mapping_kwargs["parent_asset"] = path
        for absolute_texture_path in raw_to_absolute_texture_paths.values():
            if not self.__mapping_source_path(absolute_texture_path, **mapping_kwargs):
                continue

            if Utils.is_udim_texture(absolute_texture_path):
                all_asset_paths = await self.__populate_udim_textures(absolute_texture_path)
                for path in all_asset_paths:
                    self.__mapping_source_path(path, **mapping_kwargs)

        for mdl_import in imports:
            if self._cancelled:
                break

            absolute_import_path = mdl_import.package_path
            if not await OmniClientWrapper.exists(absolute_import_path):
                continue

            self.__mapping_source_path(absolute_import_path)
            await self.__add_read_task(absolute_import_path)

        return True

    async def __add_read_task(self, path):
        if path in self._all_read_paths:
            return

        self._all_read_paths.add(path)

        if path in self._layer_store or path in self._mdl_parsers:
            return

        carb.log_info(f"Adding read task to {path}")
        is_usd_type = omni.usd.is_usd_writable_filetype(path)
        is_mdl_path = Utils.is_material(path)
        target_path = self._source_target_path_mapping.get(path)
        if is_usd_type:
            task = asyncio.ensure_future(self.__open_and_analyze_layer(path, target_path))
        elif is_mdl_path and not self._usd_only:
            task = asyncio.ensure_future(self.__open_and_analyze_mdl(path, target_path))
        else:
            return

        task.task_type = CollectorTaskType.READ_TASK
        task.source = path
        task.target = None
        self._current_tasks.add(task)
        if len(self._current_tasks) >= self.MAX_CONCURRENT_TASKS:
            await self.wait_all_unfinished_tasks()

    def __cancel_all_tasks(self):
        try:
            self._progress_callback = None
            for task in self._current_tasks:
                task.cancel()
        except Exception:
            pass
        finally:
            self._current_tasks.clear()

    async def wait_all_unfinished_tasks(self, all_completed=False):
        if len(self._current_tasks) == 0:
            return True

        carb.log_info(f"Waiting for {len(self._current_tasks)} tasks...")

        while len(self._current_tasks) > 0:
            if self._cancelled:
                self.__cancel_all_tasks()
                return False

            try:
                done, _ = await asyncio.wait(self._current_tasks, return_when=asyncio.FIRST_COMPLETED)
                failed_tasks = set()
                for task in done:
                    if task not in self._current_tasks:
                        continue

                    if task.task_type != CollectorTaskType.READ_TASK:
                        self.__report_one_progress()

                    result = task.result()
                    if not result:
                        failed_tasks.add(task)
                        if task.task_type == CollectorTaskType.COPY_TASK:
                            self._raise_or_log(task.source, f"Failed to move from {task.source} to {task.target}.")
                        elif task.task_type == CollectorTaskType.WRITE_TASK:
                            self._raise_or_log(task.target, f"Failed to write {task.target}.")
                        elif task.task_type == CollectorTaskType.READ_TASK:
                            self._raise_or_log(task.source, f"Failed to read {task.source} as it's not existed.")
                        elif task.task_type == CollectorTaskType.RESOLVE_TASK:
                            self._raise_or_log(task.source, f"Failed to resolve layer {task.source}.")

                    self._current_tasks.discard(task)
            except CollectorException as e:
                raise e
            except Exception:
                traceback.print_exc()

            if not all_completed:
                break

            await asyncio.sleep(0.1)

        carb.log_info(f"Waiting for unfinished done, left {len(self._current_tasks)} tasks...")

        return True

    def is_finished(self):
        return self._finished

    def cancel(self):
        self._cancelled = True
        self.__cancel_all_tasks()
        if self._finished_callback:
            self._finished_callback()

    def _make_sure_unique_path(self, path):
        # Generates unique path name
        if Utils.is_udim_texture(path):
            return path

        path = Utils.normalize_path(path)
        while True:
            index = self._unique_path_name.get(path, -1)
            if index == -1:
                self._unique_path_name[path] = 0
                break
            else:
                index += 1
                self._unique_path_name[path] = index
                file_path, ext = os.path.splitext(path)
                path = file_path + "_" + str(index) + ext
        self._unique_path_name[path] = 0

        return path

    def _raise_or_log(self, path, error, force_raise=False):
        is_usd_type = omni.usd.is_usd_writable_filetype(path)
        has_usd_failure_option = self._failure_options & CollectorFailureOptions.EXTERNAL_USD_REFERENCES
        has_other_failure_option = self._failure_options & CollectorFailureOptions.OTHER_EXTERNAL_REFERENCES
        if force_raise or (is_usd_type and has_usd_failure_option) or (not is_usd_type and has_other_failure_option):
            raise CollectorException(error)
        else:
            carb.log_warn(error)

    async def collect(self, progress_callback, finish_callback):
        def __reset():
            self._layer_store = {}
            self._mdl_parsers = {}
            self._sublayer_offsets = {}
            self._other_asset_paths = set()
            self._unique_path_name.clear()
            self._source_target_path_mapping = {}
            self._current_progress = 0
            self._total_steps = 0
            self._progress_callback = None
            self._finished_callback = None
            self._all_read_paths.clear()

        try:
            __reset()
            self._finished = False
            self._finished_callback = finish_callback
            self._progress_callback = progress_callback
            await self._collect_internal()
        finally:
            self._finished = True
            if finish_callback:
                finish_callback()

            __reset()

    def __report_one_progress(self):
        self._current_progress += 1
        if self._current_progress >= self._total_steps:
            self._current_progress = self._total_steps

        if self._progress_callback:
            self._progress_callback(self._current_progress, self._total_steps)

    async def _collect_internal(self):
        carb.log_info(f"Collecting stage {self._usd_path} to {self._collect_dir}...")
        total = await self._get_total_steps()
        if total == 0:
            raise CollectorException(f"Failed to collect {self._usd_path} as it cannot be opened.")
        elif total is None:
            carb.log_info(f"Collecting stage {self._usd_path} to {self._collect_dir} is cancelled.")
            return

        # Add one more step to wait for tasks to be done.
        self._total_steps = total + 1

        carb.log_info(f"Starting to collect: total steps {self._total_steps} found...")

        if not self._material_only:
            for source_layer_path, target_layer in self._layer_store.items():
                carb.log_info(f"Modifying external references for {target_layer.identifier}...")
                if self._cancelled:
                    break

                await self.__add_layer_resolve_task(source_layer_path, target_layer)

        if not self._usd_only:
            for mdl_absolute_path, mdl_parser in self._mdl_parsers.items():
                if self._cancelled:
                    break

                carb.log_info(f"Collecting mdl {mdl_absolute_path}...")
                # Per MDL.
                processed_texture_path = set([])
                mdl_target_path = self._source_target_path_mapping.get(mdl_absolute_path)
                mdl_content = mdl_parser.content
                raw_to_absolute_texture_paths, mdl_imports = await mdl_parser.parse()
                for raw_texture_path, absolute_texture_path in raw_to_absolute_texture_paths.items():
                    if self._cancelled:
                        break

                    if absolute_texture_path in processed_texture_path:
                        continue

                    processed_texture_path.add(absolute_texture_path)

                    texture_target_path = self._source_target_path_mapping.get(absolute_texture_path)
                    relative_path = Utils.make_relative_path(mdl_target_path, texture_target_path)
                    mdl_content = await aio_replace_all(
                        mdl_content, f'texture_2d("{raw_texture_path}', f'texture_2d("{relative_path}'
                    )
                    mdl_content = await aio_replace_all(
                        mdl_content, f'texture_3d("{raw_texture_path}', f'texture_3d("{relative_path}'
                    )
                    mdl_content = await aio_replace_all(
                        mdl_content, f'texture_cube("{raw_texture_path}', f'texture_cube("{relative_path}'
                    )
                    mdl_content = await aio_replace_all(
                        mdl_content, f'texture_ptex("{raw_texture_path}', f'texture_ptex("{relative_path}'
                    )

                    carb.log_info(f"Moving texture file {absolute_texture_path} to {texture_target_path}")
                    await self.add_copy_task(absolute_texture_path, texture_target_path, self._skip_existing)

                # handl mdls that's imported
                for mdl_import in mdl_imports:
                    if self._cancelled:
                        break

                    import_target_path = self._source_target_path_mapping.get(mdl_import.package_path)
                    # It's not mapped so it's not existed.
                    if not import_target_path:
                        continue

                    relative_path = Utils.make_relative_path(mdl_target_path, import_target_path)
                    import_clause = mdl_import.import_clause
                    if relative_path.endswith(".mdl"):
                        relative_path = relative_path[:-4]
                    while relative_path.startswith("./"):
                        relative_path = relative_path[2:]
                    relative_path = relative_path.replace("/", "::")
                    if mdl_import.module_preset:
                        # OM-7604
                        has_prefix = False
                        while relative_path.startswith("..::"):
                            has_prefix = True
                            relative_path = relative_path[4:]
                        if has_prefix and relative_path:
                            relative_path = f"::{relative_path}"
                    import_clause = import_clause.replace(mdl_import.import_package, relative_path)
                    mdl_content = await aio_replace_all(mdl_content, mdl_import.import_clause, import_clause)

                    if self._flat_collection:
                        import_package = mdl_import.import_package
                        while import_package.startswith(".::") or import_package.startswith("..::"):
                            if import_package.startswith(".::"):
                                import_package = import_package[3:]
                            else:
                                import_package = import_package[4:]

                        # OM-37410: It's possible that function is referenced with full module path.
                        if import_package:
                            if not import_package.startswith("::"):
                                import_package = "::" + import_package

                            if not relative_path.startswith("::"):
                                relative_path = f"::{relative_path}"
                            mdl_content = await aio_replace_all(mdl_content, import_package, relative_path)

                await self.add_write_task(mdl_target_path, mdl_content)

        # Copy all external assets except USD and MDLs to target path
        for absolute_asset_path in self._other_asset_paths:
            if self._cancelled:
                break

            asset_target_path = self._source_target_path_mapping.get(absolute_asset_path)
            carb.log_info(f"Moving asset file from {absolute_asset_path} to {asset_target_path}")
            await self.add_copy_task(absolute_asset_path, asset_target_path, self._skip_existing)

        if not await self.wait_all_unfinished_tasks(True):
            carb.log_info("Collecting is cancelled.")
        else:
            carb.log_info("Collecting is finished.")
