# Copyright (c) 2020, 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 traceback
from collections import defaultdict
from typing import Any, DefaultDict, Dict, List, Sequence, Set, Tuple

import carb
import carb.events
import carb.profiler
import omni.kit.app
import omni.kit.context_menu
import omni.ui as ui
import omni.usd
from omni.kit.window.property.templates import SimplePropertyWidget, build_frame_header
from pxr import Sdf, Tf, Trace, Usd

from .usd_model_base import UsdBase
from .usd_property_widget_builder import *
from .message_bus_events import ADDITIONAL_CHANGED_PATH_EVENT_TYPE

# Clipboard to store copied group properties
__properties_to_copy: Dict[Sdf.Path, Any] = {}


def get_group_properties_clipboard():
    return __properties_to_copy


def set_group_properties_clipboard(properties_to_copy: Dict[Sdf.Path, Any]):
    global __properties_to_copy
    __properties_to_copy = properties_to_copy


class UsdPropertyUiEntry:
    def __init__(
        self,
        prop_name: str,
        display_group: str,
        metadata,
        property_type,
        build_fn=None,
        display_group_collapsed: bool = False,
        prim_paths: List[Sdf.Path] = None,
    ):
        """
        Constructor.

        Args:
            prop_name: name of the Usd Property. This is not the display name.
            display_group: group of the Usd Property when displayed on UI.
            metadata: metadata associated with the Usd Property.
            property_type: type of the property. Either Usd.Property or Usd.Relationship.
            build_fn: a custom build function to build the UI. If not None the default builder will not be used.
            display_group_collapsed: if the display group should be collapsed. Group only collapses when ALL its contents request such.
            prim_paths: to override what prim paths this property will be built upon. Leave it to None to use default (currently selected paths, or last selected path if multi-edit is off).
        """
        self.prop_name = prop_name
        self.display_group = display_group
        self.display_group_collapsed = display_group_collapsed
        self.metadata = metadata
        self.property_type = property_type
        self.prim_paths = prim_paths
        self.build_fn = build_fn

    def add_custom_metadata(self, key: str, value):
        """
        If value is not None, add it to the custom data of the metadata using the specified key.  Otherwise, remove that
        key from the custom data.

        Args:
            key: the key that should containt the custom metadata value
            value: the value that should be added to the custom metadata if not None
        """
        custom_data = self.metadata.get(Sdf.PrimSpec.CustomDataKey, {})
        if value is not None:
            custom_data[key] = value
        elif key in custom_data:
            del custom_data[key]
        self.metadata[Sdf.PrimSpec.CustomDataKey] = custom_data

    def override_display_group(self, display_group: str, collapsed: bool = False):
        """
        Overrides the display group of the property. It only affects UI and DOES NOT write back DisplayGroup metadata to USD.

        Args:
            display_group: new display group to override to.
            collapsed: if the display group should be collapsed. Group only collapses when ALL its contents request such.
        """
        self.display_group = display_group
        self.display_group_collapsed = collapsed

    def override_display_name(self, display_name: str):
        """
        Overrides the display name of the property. It only affects UI and DOES NOT write back DisplayName metadata to USD.
        Args:
            display_group: new display group to override to.
        """
        self.metadata[Sdf.PropertySpec.DisplayNameKey] = display_name

    # for backward compatibility
    def __getitem__(self, key):
        if key == 0:
            return self.prop_name
        elif key == 1:
            return self.display_group
        elif key == 2:
            return self.metadata
        return None

    @property
    def attr_name(self):
        return self.prop_name

    @attr_name.setter
    def attr_name(self, value):
        self.prop_name = value

    def get_nested_display_groups(self):
        if len(self.display_group) == 0:
            return []
        # Per USD documentation nested display groups are separated by colon
        return self.display_group.split(":")

    def __eq__(self, other):
        return (
            type(self) == type(other)
            and self.prop_name == other.prop_name
            and self.display_group == other.display_group
            and self._compare_metadata(self.metadata, other.metadata)
            and self.property_type == other.property_type
            and self.prim_paths == other.prim_paths
        )

    def _compare_metadata(self, meta1, meta2) -> bool:
        ignored_metadata = {"default", "colorSpace"}
        for key, value in meta1.items():
            if key not in ignored_metadata and (key not in meta2 or meta2[key] != value):
                return False
        for key, value in meta2.items():
            if key not in ignored_metadata and (key not in meta1 or meta1[key] != value):
                return False
        return True


class UiDisplayGroup:
    def __init__(self, name):
        self.name = name
        self.sub_groups = DefaultDict()
        self.props = []


class UsdPropertiesWidget(SimplePropertyWidget):
    """
    UsdPropertiesWidget provides functionalities to automatically populates UsdProperties on given prim(s). The UI will
    and models be generated according to UsdProperties's value type. Multi-prim editing works for shared Properties
    between all selected prims if instantiated with multi_edit = True.
    """

    def __init__(self, title: str, collapsed: bool, multi_edit: bool = True):
        """
        Constructor.

        Args:
            title: title of the widget.
            collapsed: whether the collapsable frame should be collapsed for this widget.
            multi_edit: whether multi-editing is supported.
                If False, properties will only be collected from the last selected prim.
                If True, shared properties among all selected prims will be collected.
        """
        super().__init__(title=title, collapsed=collapsed)
        self._multi_edit = multi_edit
        self._models = defaultdict(list)
        self._message_bus = omni.kit.app.get_app().get_message_bus_event_stream()
        self._bus_sub = None
        self._listener = None
        self._pending_dirty_task = None
        self._pending_dirty_paths = set()
        self._group_menu_entries = []

    def clean(self):
        """
        See PropertyWidget.clean
        """
        self.reset_models()
        super().clean()

    def reset(self):
        """
        See PropertyWidget.reset
        """
        self.reset_models()
        super().reset()

    def reset_models(self):
        # models can be shared among multiple prims. Only clean once!
        unique_models = set()
        if self._models is not None:
            for models in self._models.values():
                for model in models:
                    unique_models.add(model)
        for model in unique_models:
            model.clean()
        self._models = defaultdict(list)
        if self._listener:
            self._listener.Revoke()
        self._listener = None
        self._bus_sub = None
        if self._pending_dirty_task is not None:
            self._pending_dirty_task.cancel()
            self._pending_dirty_task = None
        self._pending_dirty_paths.clear()
        for entry in self._group_menu_entries:
            entry.release()
        self._group_menu_entries.clear()

    def get_additional_kwargs(self, ui_prop: UsdPropertyUiEntry):
        """
        Override this function if you want to supply additional arguments when building the label or ui widget.
        """
        additional_label_kwargs = None
        additional_widget_kwargs = None
        return additional_label_kwargs, additional_widget_kwargs

    def build_property_item(self, stage, ui_prop: UsdPropertyUiEntry, prim_paths: List[Sdf.Path]):
        """
        Override this function to customize property building.
        """

        # override prim paths to build if UsdPropertyUiEntry specifies one
        if ui_prop.prim_paths:
            prim_paths = ui_prop.prim_paths

        build_fn = ui_prop.build_fn if ui_prop.build_fn else UsdPropertiesWidgetBuilder.build
        additional_label_kwargs, additional_widget_kwargs = self.get_additional_kwargs(ui_prop)
        models = build_fn(
            stage,
            ui_prop.prop_name,
            ui_prop.metadata,
            ui_prop.property_type,
            prim_paths,
            additional_label_kwargs,
            additional_widget_kwargs,
        )
        if models:
            if not isinstance(models, list):
                models = [models]
            for model in models:
                for prim_path in prim_paths:
                    self._models[prim_path.AppendProperty(ui_prop.prop_name)].append(model)

    def build_nested_group_frames(self, stage, display_group: UiDisplayGroup):
        """
        Override this function to build group frames differently.
        """
        if self._multi_edit:
            prim_paths = self._payload.get_paths()
        else:
            prim_paths = self._payload[-1:]

        def build_props(props):
            collapse_frame = len(props) > 0
            # we need to build those property in 2 different possible locations
            for prop in props:
                self.build_property_item(stage, prop, prim_paths)
                collapse_frame &= prop.display_group_collapsed

            return collapse_frame

        def get_sub_props(group: UiDisplayGroup, sub_prop_list: list):
            if len(group.sub_groups) > 0:
                for name, sub_group in group.sub_groups.items():
                    for sub_prop in sub_group.props:
                        sub_prop_list.append(sub_prop)
                    get_sub_props(sub_group, sub_prop_list)

        def build_nested(display_group: UiDisplayGroup, level: int, prefix: str):
            # Only create a collapsable frame if the group is not "" (for root level group)
            if len(display_group.name) > 0:
                id = prefix + ":" + display_group.name
                frame = ui.CollapsableFrame(
                    title=display_group.name,
                    build_header_fn=lambda collapsed, text, id=id: self._build_frame_header(collapsed, text, id),
                    name="subFrame",
                )
            else:
                id = prefix
                frame = ui.Frame(name="subFrame")

            prop_list = []
            sub_props = get_sub_props(display_group, prop_list)

            with frame:
                with ui.VStack(height=0, spacing=5, name="frame_v_stack"):
                    for name, sub_group in display_group.sub_groups.items():
                        build_nested(sub_group, level + 1, id)

                    # Only do "Extra Properties" for root level group
                    if len(display_group.sub_groups) > 0 and len(display_group.props) > 0 and level == 0:
                        extra_group_name = "Extra Properties"
                        extra_group_id = prefix + ":" + extra_group_name
                        with ui.CollapsableFrame(
                            title=extra_group_name,
                            build_header_fn=lambda collapsed, text, id=extra_group_id: self._build_frame_header(
                                collapsed, text, id
                            ),
                            name="subFrame",
                        ):
                            with ui.VStack(height=0, spacing=5, name="frame_v_stack"):
                                build_props(display_group.props)

                            self._build_header_context_menu(
                                group_name=extra_group_name, group_id=extra_group_id, props=sub_props
                            )
                    else:
                        collapse = build_props(display_group.props)
                        if collapse and isinstance(frame, ui.CollapsableFrame):
                            frame.collapsed = collapse

                # if level is 0, this is the root level group, and we use the self._title for its name
                self._build_header_context_menu(
                    group_name=display_group.name if level > 0 else self._title, group_id=id, props=sub_props
                )

        build_nested(display_group, 0, self._title)

        # if we have subFrame we need the main frame to assume the groupFrame styling
        if len(display_group.sub_groups) > 0:
            # here we reach into the Parent class frame
            self._collapsable_frame.name = "groupFrame"

    def build_items(self):
        """
        See SimplePropertyWidget.build_items
        """
        self.reset()

        if len(self._payload) == 0:
            return

        last_prim = self._get_prim(self._payload[-1])
        if not last_prim:
            return

        stage = last_prim.GetStage()
        if not stage:
            return

        self._listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, self._on_usd_changed, stage)
        self._bus_sub = self._message_bus.create_subscription_to_pop_by_type(
            ADDITIONAL_CHANGED_PATH_EVENT_TYPE, self._on_bus_event
        )

        shared_props = self._get_shared_properties_from_selected_prims(last_prim)
        if not shared_props:
            return

        shared_props = self._customize_props_layout(shared_props)
        grouped_props = UiDisplayGroup("")

        for prop in shared_props:
            nested_groups = prop.get_nested_display_groups()
            sub_group = grouped_props
            for sub_group_name in nested_groups:
                sub_group = sub_group.sub_groups.setdefault(sub_group_name, UiDisplayGroup(sub_group_name))
            sub_group.props.append(prop)

        self.build_nested_group_frames(stage, grouped_props)

    def _build_header_context_menu(self, group_name: str, group_id: str, props: List[UsdPropertyUiEntry] = None):
        """
        Override this function to build the context menu when right click on a Collapsable group header

        Args:
            group_name: Display name of the group. If it's not a subgroup, it's the title of the widget.
            group_id: A unique identifier for group context menu.
            props: Properties under this group. It contains all properties in its subgroups as well.
        """
        self._build_group_builtin_header_context_menu(group_name, group_id, props)
        self._build_group_additional_header_context_menu(group_name, group_id, props)

    def _build_group_builtin_header_context_menu(
        self, group_name: str, group_id: str, props: List[UsdPropertyUiEntry] = None
    ):
        from .usd_attribute_model import GfVecAttributeSingleChannelModel

        prop_names: Set[str] = set()
        if props:
            for prop in props:
                prop_names.add(prop.prop_name)

        def can_copy(object):
            # Only support single selection copy oer OM-20206
            return len(self._payload) == 1

        def on_copy(object):
            visited_models = set()

            properties_to_copy: Dict[Sdf.Path, Any] = dict()
            for models in self._models.values():
                for model in models:
                    if model in visited_models:
                        continue

                    visited_models.add(model)

                    # Skip "Mixed"
                    if model.is_ambiguous():
                        continue

                    paths = model.get_property_paths()
                    if paths:
                        # Copy from the last path
                        # In theory if the value is not mixed all paths should have same value, so which one to pick doesn't matter
                        last_path = paths[-1]

                        # Only copy from the properties from this group
                        if props is not None and last_path.name not in prop_names:
                            continue

                        if issubclass(type(model), UsdBase):
                            # No need to copy single channel model. Each vector attribute also has a GfVecAttributeModel
                            if isinstance(model, GfVecAttributeSingleChannelModel):
                                continue

                            properties_to_copy[paths[-1]] = model.get_value()

                        elif isinstance(model, RelationshipEditWidget):
                            properties_to_copy[paths[-1]] = model.get_targets().copy()

            if properties_to_copy:
                set_group_properties_clipboard(properties_to_copy)

        menu = {
            "name": f'Copy All Property Values in "{group_name}"',
            "show_fn": lambda object: True,
            "enabled_fn": can_copy,
            "onclick_fn": on_copy,
        }

        self._group_menu_entries.append(self._register_header_context_menu_entry(menu, group_id))

        def on_paste(object):
            properties_to_copy = get_group_properties_clipboard()
            if not properties_to_copy:
                return

            unique_model_prim_paths: Set[Sdf.Path] = set()
            for prop_path in self._models:
                unique_model_prim_paths.add(prop_path.GetPrimPath())

            with omni.kit.undo.group():
                try:
                    for path, value in properties_to_copy.items():
                        for prim_path in unique_model_prim_paths:

                            # Only paste to the properties in this group
                            if props is not None and path.name not in prop_names:
                                continue

                            paste_to_model_path = prim_path.AppendProperty(path.name)
                            models = self._models.get(paste_to_model_path, [])
                            for model in models:
                                if isinstance(model, GfVecAttributeSingleChannelModel):
                                    continue
                                else:
                                    model.set_value(value)
                except Exception as e:
                    carb.log_warn(traceback.format_exc())

        menu = {
            "name": f'Paste All Property Values to "{group_name}"',
            "show_fn": lambda object: True,
            "enabled_fn": lambda object: get_group_properties_clipboard(),
            "onclick_fn": on_paste,
        }
        self._group_menu_entries.append(self._register_header_context_menu_entry(menu, group_id))

        def can_reset(object):
            visited_models = set()

            for models in self._models.values():
                for model in models:
                    if model in visited_models:
                        continue

                    visited_models.add(model)

                    paths = model.get_property_paths()
                    if paths:

                        last_path = paths[-1]

                        # Only reset from the properties from this group
                        if props is not None and last_path.name not in prop_names:
                            continue

                        if issubclass(type(model), UsdBase):
                            if model.is_different_from_default():
                                return True
            return False

        def on_reset(object):
            visited_models = set()

            for models in self._models.values():
                for model in models:
                    if model in visited_models:
                        continue

                    visited_models.add(model)

                    paths = model.get_property_paths()
                    if paths:

                        last_path = paths[-1]

                        # Only reset from the properties from this group
                        if props is not None and last_path.name not in prop_names:
                            continue

                        if issubclass(type(model), UsdBase):
                            model.set_default()

        menu = {
            "name": f'Reset All Property Values in "{group_name}"',
            "show_fn": lambda object: True,
            "enabled_fn": can_reset,
            "onclick_fn": on_reset,
        }

        self._group_menu_entries.append(self._register_header_context_menu_entry(menu, group_id))

    def _build_group_additional_header_context_menu(
        self, group_name: str, group_id: str, props: List[UsdPropertyUiEntry] = None
    ):
        """
        Override this function to build the additional context menu to Kit's built-in ones when right click on a Collapsable group header

        Args:
            group_name: Display name of the group. If it's not a subgroup, it's the title of the widget.
            group_id: A unique identifier for group context menu.
            props: Properties under this group. It contains all properties in its subgroups as well.
        """

        ...

    def _register_header_context_menu_entry(self, menu: Dict, group_id: str):
        """
        Registers a menu entry to Collapsable group header

        Args:
            menu: The menu entry to be registered.
            group_id: A unique identifier for group context menu.

        Return:
            The subscription object of the menu entry to be kept alive during menu's life span.
        """
        return omni.kit.context_menu.add_menu(menu, "group_context_menu." + group_id, "omni.kit.window.property")

    def _filter_props_to_build(self, props):
        """
        When deriving from UsdPropertiesWidget, override this function to filter properties to build.
        Args:
            props: List of Usd.Property on a selected prim.
        """
        return [prop for prop in props if not prop.IsHidden()]

    def _customize_props_layout(self, props):
        """
        When deriving from UsdPropertiesWidget, override this function to reorder/regroup properties to build.
        To reorder the properties display order, reorder entries in props list.
        To override display group or name, call prop.override_display_group or prop.override_display_name respectively.
        If you want to hide/add certain property, remove/add them to the list.

        NOTE: All above changes won't go back to USD, they're pure UI overrides.

        Args:
            props: List of Tuple(property_name, property_group, metadata)

        Example:

            for prop in props:
                # Change display group:
                prop.override_display_group("New Display Group")

                # Change display name (you can change other metadata, it won't be write back to USD, only affect UI):
                prop.override_display_name("New Display Name")

            # add additional "property" that doesn't exist.
            props.append(UsdPropertyUiEntry("PlaceHolder", "Group", { Sdf.PrimSpec.TypeNameKey: "bool"}, Usd.Property))
        """

        return props

    @Trace.TraceFunction
    def _on_usd_changed(self, notice, stage):
        carb.profiler.begin(1, "UsdPropertyWidget._on_usd_changed")
        try:
            if stage != self._payload.get_stage():
                return

            if not self._collapsable_frame:
                return

            if len(self._payload) == 0:
                return

            # Widget is pending rebuild, no need to check for dirty
            if self._pending_rebuild_task is not None:
                return

            dirty_paths = set()

            for path in notice.GetResyncedPaths():
                if path in self._payload:
                    self.request_rebuild()
                    return
                elif path.GetPrimPath() in self._payload:
                    prop = stage.GetPropertyAtPath(path)
                    # If prop is added or removed, rebuild frame
                    # TODO only check against the properties this widget cares about
                    if (not prop.IsValid()) != (self._models.get(path) is None):
                        self.request_rebuild()
                        return
                    # else trigger existing model to reload the value
                    else:
                        dirty_paths.add(path)

            for path in notice.GetChangedInfoOnlyPaths():
                dirty_paths.add(path)

            self._pending_dirty_paths.update(dirty_paths)

            if self._pending_dirty_task is None:
                self._pending_dirty_task = asyncio.ensure_future(self._delayed_dirty_handler())
        finally:
            carb.profiler.end(1)

    def _on_bus_event(self, event: carb.events.IEvent):
        # TODO from C++?
        # stage = event.payload["stage"]
        # if stage != self._payload.get_stage():
        #     return

        if not self._collapsable_frame:
            return

        if len(self._payload) == 0:
            return

        # Widget is pending rebuild, no need to check for dirty
        if self._pending_rebuild_task is not None:
            return

        path = event.payload["path"]
        self._pending_dirty_paths.add(Sdf.Path(path))
        if self._pending_dirty_task is None:
            self._pending_dirty_task = asyncio.ensure_future(self._delayed_dirty_handler())

    def _get_prim(self, prim_path):
        if prim_path:
            stage = self._payload.get_stage()
            if stage:
                return stage.GetPrimAtPath(prim_path)
        return None

    def _get_shared_properties_from_selected_prims(self, anchor_prim):
        shared_props_dict = None

        if self._multi_edit:
            prim_paths = self._payload.get_paths()
        else:
            prim_paths = self._payload[-1:]

        for prim_path in prim_paths:
            prim = self._get_prim(prim_path)
            if not prim:
                continue

            props = self._filter_props_to_build(prim.GetProperties())
            prop_dict = {}
            for prop in props:
                if prop.IsHidden():
                    continue
                prop_dict[prop.GetName()] = UsdPropertyUiEntry(
                    prop.GetName(), prop.GetDisplayGroup(), prop.GetAllMetadata(), type(prop)
                )

            if shared_props_dict is None:
                shared_props_dict = prop_dict
            else:
                # Find intersection of the dicts
                intersect_shared_props = {}
                for prop_name, prop_info in shared_props_dict.items():
                    if prop_dict.get(prop_name) == prop_info:
                        intersect_shared_props[prop_name] = prop_info

                if len(intersect_shared_props) == 0:
                    # No intersection, nothing to build
                    # early return
                    return

                shared_props_dict = intersect_shared_props

        shared_props = list(shared_props_dict.values())
        # Sort properties to PropertyOrder using the last selected object
        order = anchor_prim.GetPropertyOrder()
        shared_prop_order = []
        shared_prop_unordered = []
        for prop in shared_props:
            if prop[0] in order:
                shared_prop_order.append(prop[0])
            else:
                shared_prop_unordered.append(prop[0])

        shared_prop_order.extend(shared_prop_unordered)
        sorted(shared_props, key=lambda x: shared_prop_order.index(x[0]))
        return shared_props

    def request_rebuild(self):
        # If widget is going to be rebuilt, _pending_dirty_task does not need to run.
        if self._pending_dirty_task is not None:
            self._pending_dirty_task.cancel()
        self._pending_dirty_task = None
        self._pending_dirty_paths.clear()
        super().request_rebuild()

    async def _delayed_dirty_handler(self):
        while True:
            # Do not refresh UI until visible/uncollapsed
            if not self._collapsed:
                break

            await omni.kit.app.get_app().next_update_async()

        # clear the pending dirty tasks BEFORE dirting model.
        # dirtied model may trigger additional USD notice that needs to be scheduled.
        self._pending_dirty_task = None

        if self._pending_dirty_paths:
            # Make a copy of the paths. It may change if USD edits are made during iteration
            pending_dirty_paths = self._pending_dirty_paths.copy()
            self._pending_dirty_paths.clear()

            carb.profiler.begin(1, "UsdPropertyWidget._delayed_dirty_handler")
            # multiple path can share the same model. Only dirty once!
            dirtied_models = set()
            for path in pending_dirty_paths:
                models = self._models.get(path)
                if models:
                    for model in models:
                        if model not in dirtied_models:
                            model._set_dirty()
                            dirtied_models.add(model)

            carb.profiler.end(1)


class SchemaPropertiesWidget(UsdPropertiesWidget):
    """
    SchemaPropertiesWidget only filters properties and only show the onces from a given IsA schema or applied API schema.
    """

    def __init__(self, title: str, schema, include_inherited: bool):
        """
        Constructor.

        Args:
            title (str): Title of the widgets on the Collapsable Frame.
            schema: The USD IsA schema or applied API schema to filter properties.
            include_inherited (bool): Whether the filter should include inherited properties.
        """
        super().__init__(title, collapsed=False)
        self._title = title
        self._schema = schema
        self._include_inherited = include_inherited

    def on_new_payload(self, payload):
        """
        See PropertyWidget.on_new_payload
        """
        if not super().on_new_payload(payload):
            return False

        if not self._payload or len(self._payload) == 0:
            return False

        for prim_path in self._payload:
            prim = self._get_prim(prim_path)

            if not prim:
                return False

            is_api_schema = Usd.SchemaRegistry().IsAppliedAPISchema(self._schema)
            if not (is_api_schema and prim.HasAPI(self._schema) or not is_api_schema and prim.IsA(self._schema)):
                return False

        return True

    def _filter_props_to_build(self, props):
        """
        See UsdPropertiesWidget._filter_props_to_build
        """
        if len(props) == 0:
            return props

        if Usd.SchemaRegistry().IsMultipleApplyAPISchema(self._schema):
            prim = props[0].GetPrim()
            schema_instances = set()
            schema_type_name = Usd.SchemaRegistry().GetSchemaTypeName(self._schema)
            for schema in prim.GetAppliedSchemas():
                if schema.startswith(schema_type_name):
                    schema_instances.add(schema[len(schema_type_name) + 1 :])
            filtered_props = []
            api_path_func_name = f"Is{schema_type_name}Path"
            api_path_func = getattr(self._schema, api_path_func_name)

            # self._schema.GetSchemaAttributeNames caches the query result in a static variable and any new instance
            # token passed into it won't change the cached property names. This can potentially cause problems as other
            # code calling same function with different instance name will get wrong result.
            #
            # There's any other function IsSchemaPropertyBaseName but it's not implemented on all applied schemas. (not
            # implemented in a few PhysicsSchema, for example.
            #
            # if include_inherited is True, it returns SchemaTypeName:BaseName for properties, otherwise it only returns
            # BaseName.
            schema_attr_names = self._schema.GetSchemaAttributeNames(self._include_inherited, "")

            for prop in props:
                if prop.IsHidden():
                    continue
                prop_path = prop.GetPath().pathString
                for instance_name in schema_instances:
                    instance_seg = prop_path.find(":" + instance_name + ":")
                    if instance_seg != -1:
                        api_path = prop_path[0 : instance_seg + 1 + len(instance_name)]
                        if api_path_func(api_path):
                            base_name = prop_path[instance_seg + 1 + len(instance_name) + 1 :]
                            if base_name in schema_attr_names:
                                filtered_props.append(prop)
                                break
            return filtered_props
        else:
            schema_attr_names = self._schema.GetSchemaAttributeNames(self._include_inherited)
            return [prop for prop in props if prop.GetName() in schema_attr_names]


class MultiSchemaPropertiesWidget(UsdPropertiesWidget):
    """
    MultiSchemaPropertiesWidget filters properties and only show the onces from a given IsA schema or schema subclass list.
    """

    __known_api_schemas = set()

    def __init__(
        self,
        title: str,
        schema,
        schema_subclasses: list,
        include_list: list = [],
        exclude_list: list = [],
        api_schemas: Sequence[str] = None,
        group_api_schemas: bool = False,
    ):
        """
        Constructor.

        Args:
            title (str): Title of the widgets on the Collapsable Frame.
            schema: The USD IsA schema or applied API schema to filter properties.
            schema_subclasses (list): list of subclasses
            include_list (list): list of additional schema named to add
            exclude_list (list): list of additional schema named to remove
            api_schemas (sequence): a sequence of AppliedAPI schema names that this widget handles
            group_api_schemas (bool): whether to create default groupings for any AppliedSchemas on the Usd.Prim
        """
        super().__init__(title=title, collapsed=False)
        self._title = title
        self._schema = schema

        # create schema_attr_names
        self._schema_attr_base = schema.GetSchemaAttributeNames(False)
        for subclass in schema_subclasses:
            self._schema_attr_base += subclass.GetSchemaAttributeNames(False)
        self._schema_attr_base += include_list
        self._schema_attr_base = set(self._schema_attr_base) - set(exclude_list)
        # Setup the defaults for handling applied API schemas
        self._applied_schemas = {}
        self._schema_attr_names = None
        self._group_api_schemas = group_api_schemas
        # Save any custom Applied Schemas and mark them to be ignored when building default widget
        self._custom_api_schemas = api_schemas
        if self._custom_api_schemas:
            MultiSchemaPropertiesWidget.__known_api_schemas.update(self._custom_api_schemas)

    def __del__(self):
        # Mark any custom Applied Schemas to start being handled by the defaut widget
        if self._custom_api_schemas:
            MultiSchemaPropertiesWidget.__known_api_schemas.difference_update(self._custom_api_schemas)

    def clean(self):
        """
        See PropertyWidget.clean
        """
        self._applied_schemas = {}
        self._schema_attr_names = None
        super().clean()

    def on_new_payload(self, payload):
        """
        See PropertyWidget.on_new_payload
        """
        if not super().on_new_payload(payload):
            return False

        if not self._payload or len(self._payload) == 0:
            return False

        used = []
        schema_reg = Usd.SchemaRegistry()
        self._applied_schemas = {}
        # Build out our _schema_attr_names variable to include properties from the base schema and any applied schemas
        self._schema_attr_names = self._schema_attr_base

        for prim_path in self._payload:
            prim = self._get_prim(prim_path)
            if not prim or not prim.IsA(self._schema):
                return False
            # If the base-schema has requested auto-grouping of all Applied Schemas, handle that grouping now
            if self._group_api_schemas:
                # XXX: Should this be delayed until _customize_props_layout ?
                for api_schema in prim.GetAppliedSchemas():
                    # Ignore any API schemas that are already registered as custom widgets
                    if api_schema in MultiSchemaPropertiesWidget.__known_api_schemas:
                        continue
                    # Skip over any API schemas that USD doesn't actually know about
                    prim_def = schema_reg.FindAppliedAPIPrimDefinition(api_schema)
                    if not prim_def:
                        continue
                    api_prop_names = prim_def.GetPropertyNames()
                    self._schema_attr_names = self._schema_attr_names.union(api_prop_names)
                    for api_prop in api_prop_names:
                        display_group = prim_def.GetPropertyMetadata(api_prop, "displayGroup")
                        prop_grouping = self._applied_schemas.setdefault(api_schema, {}).setdefault(display_group, [])
                        prop_grouping.append((api_prop, prim_def.GetPropertyMetadata(api_prop, "displayName")))

            used += self._filter_props_to_build(prim.GetProperties())

        return used

    def _filter_props_to_build(self, props):
        """
        See UsdPropertiesWidget._filter_props_to_build
        """
        return [prop for prop in props if prop.GetName() in self._schema_attr_names and not prop.IsHidden()]

    def _customize_props_layout(self, attrs):
        # If no applied schemas, just use the base class' layout.
        if not self._applied_schemas:
            return super()._customize_props_layout(attrs)

        from omni.kit.property.usd.custom_layout_helper import (
            CustomLayoutFrame,
            CustomLayoutGroup,
            CustomLayoutProperty,
        )

        # We can't really escape the parent group, so all default/base and applied-schemas will be common to a group
        frame = CustomLayoutFrame(hide_extra=False)
        with frame:
            # Add all base properties at the top
            base_attrs = [attr for attr in attrs if attr.prop_name in self._schema_attr_base]
            if base_attrs:
                with CustomLayoutGroup(self._title):
                    for attr in base_attrs:
                        CustomLayoutProperty(attr.prop_name, attr.display_group)
                        attr.override_display_group(self._title)

        # Now create a master-group for each applied schema, and possibly sub-groups for it's properties
        # Here's where we may want to actually escape the parent and create a totally new group
        with frame:
            for api_schema, api_schema_groups in self._applied_schemas.items():
                with CustomLayoutGroup(api_schema):
                    for prop_group, props in api_schema_groups.items():
                        with CustomLayoutGroup(prop_group):
                            for prop in props:
                                CustomLayoutProperty(*prop)

        return frame.apply(attrs)


class RawUsdPropertiesWidget(UsdPropertiesWidget):

    MULTI_SELECTION_LIMIT_SETTING_PATH = "/persistent/exts/omni.kit.property.usd/raw_widget_multi_selection_limit"
    MULTI_SELECTION_LIMIT_DO_NOT_ASK_SETTING_PATH = (
        "/exts/omni.kit.property.usd/multi_selection_limit_do_not_ask"  # session only, not persistent!
    )

    def __init__(self, title: str, collapsed: bool, multi_edit: bool = True):
        super().__init__(title=title, collapsed=collapsed, multi_edit=multi_edit)
        self._settings = carb.settings.get_settings()
        self._settings.set_default(RawUsdPropertiesWidget.MULTI_SELECTION_LIMIT_DO_NOT_ASK_SETTING_PATH, False)

        self._skip_multi_selection_protection = False

    def on_new_payload(self, payload):
        if not super().on_new_payload(payload):
            return False

        for prim_path in self._payload:
            if not prim_path.IsPrimPath():
                return False

        self._skip_multi_selection_protection = False

        return bool(self._payload and len(self._payload) > 0)

    def build_impl(self):
        # rebuild frame when frame is opened
        super().build_impl()
        self._collapsable_frame.set_collapsed_changed_fn(self._on_collapsed_changed)

    def build_items(self):
        # only show raw items if frame is open
        if self._collapsable_frame and not self._collapsable_frame.collapsed:
            if not self._multi_selection_protected():
                super().build_items()

    def _on_collapsed_changed(self, collapsed):
        if not collapsed:
            self.request_rebuild()

    def _multi_selection_protected(self):
        if self._no_multi_selection_protection_this_session():
            return False

        def show_all(do_not_ask: bool):
            self._settings.set(RawUsdPropertiesWidget.MULTI_SELECTION_LIMIT_DO_NOT_ASK_SETTING_PATH, do_not_ask)
            self._skip_multi_selection_protection = True
            self.request_rebuild()

        multi_select_limit = self._settings.get(RawUsdPropertiesWidget.MULTI_SELECTION_LIMIT_SETTING_PATH)

        if multi_select_limit and len(self._payload) > multi_select_limit and not self._skip_multi_selection_protection:
            ui.Separator()
            ui.Label(
                f"You have selected {len(self._payload)} Prims, to preserve fast performance the Raw Usd Properties Widget is not showing above the current limit of {multi_select_limit} Prims. Press the button below to show it anyway but expect potential performance penalty.",
                width=omni.ui.Percent(100),
                alignment=ui.Alignment.CENTER,
                name="label",
                word_wrap=True,
            )
            button = ui.Button("Skip Multi Selection Protection")
            with ui.HStack(width=0):
                checkbox = ui.CheckBox()
                ui.Spacer(width=5)
                ui.Label("Do not ask again for current session.", name="label")
            button.set_clicked_fn(lambda: show_all(checkbox.model.get_value_as_bool()))
            return True

        return False

    def _no_multi_selection_protection_this_session(self):
        return self._settings.get(RawUsdPropertiesWidget.MULTI_SELECTION_LIMIT_DO_NOT_ASK_SETTING_PATH)
