# 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.
#

from omni.kit.window.property.templates import SimplePropertyWidget, build_frame_header
from .usd_property_widget_builder import *
from collections import defaultdict, OrderedDict

from .usd_property_widget import *

import asyncio
import carb
import omni.ui as ui
import omni.usd

from typing import List, Dict
from pxr import Usd, Tf, Sdf


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


class UsdAttributesWidget(UsdPropertiesWidget):
    """
    DEPRECATED! KEEP FOR BACKWARD COMPATIBILITY. Use UsdPropertiesWidget instead.

    UsdAttributesWidget provides functionalities to automatically populates UsdAttributes on given prim(s). The UI will
    and models be generated according to UsdAttributes's value type. Multi-prim editing works for shared Attributes
    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, collapsed, multi_edit)

    def build_attribute_item(self, stage, ui_attr: UsdAttributeUiEntry, prim_paths: List[Sdf.Path]):
        return super().build_property_item(stage, ui_attr, prim_paths)

    def _filter_attrs_to_build(self, attrs):
        return super()._filter_props_to_build(attrs)

    def _customize_attrs_layout(self, attrs):
        return super()._customize_props_layout(attrs)

    def _get_shared_attributes_from_selected_prims(self, anchor_prim):
        return super()._get_shared_properties_from_selected_prims(anchor_prim)

    def build_property_item(self, stage, ui_prop: UsdPropertyUiEntry, prim_paths: List[Sdf.Path]):
        return self.build_attribute_item(stage, ui_prop, prim_paths)

    def _filter_props_to_build(self, props):
        return self._filter_attrs_to_build(props)

    def _customize_props_layout(self, props):
        return self._customize_attrs_layout(props)

    def _get_shared_properties_from_selected_prims(self, anchor_prim):
        return self._get_shared_attributes_from_selected_prims(anchor_prim)


class SchemaAttributesWidget(SchemaPropertiesWidget):
    """
    DEPRECATED! KEEP FOR BACKWARD COMPATIBILITY. Use SchemaPropertiesWidget instead.

    SchemaAttributesWidget only filters attributes 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 attributes.
            include_inherited (bool): Whether the filter should include inherited attributes.
        """
        super().__init__(title, schema, include_inherited)

    def build_attribute_item(self, stage, ui_attr: UsdAttributeUiEntry, prim_paths: List[Sdf.Path]):
        return super().build_property_item(stage, ui_attr, prim_paths)

    def _filter_attrs_to_build(self, attrs):
        return super()._filter_props_to_build(attrs)

    def _customize_attrs_layout(self, attrs):
        return super()._customize_props_layout(attrs)

    def _get_shared_attributes_from_selected_prims(self, anchor_prim):
        return super()._get_shared_properties_from_selected_prims(anchor_prim)

    def build_property_item(self, stage, ui_prop: UsdPropertyUiEntry, prim_paths: List[Sdf.Path]):
        return self.build_attribute_item(stage, ui_prop, prim_paths)

    def _filter_props_to_build(self, props):
        return self._filter_attrs_to_build(props)

    def _customize_props_layout(self, props):
        return self._customize_attrs_layout(props)

    def _get_shared_properties_from_selected_prims(self, anchor_prim):
        return self._get_shared_attributes_from_selected_prims(anchor_prim)


class MultiSchemaAttributesWidget(MultiSchemaPropertiesWidget):
    """
    DEPRECATED! KEEP FOR BACKWARD COMPATIBILITY. Use MultiSchemaPropertiesWidget instead.

    MultiSchemaAttributesWidget filters attributes and only show the onces from a given IsA schema or schema subclass list.
    """

    def __init__(self, title: str, schema, schema_subclasses: list, include_list: list = [], exclude_list: list = []):
        """
        Constructor.

        Args:
            title (str): Title of the widgets on the Collapsable Frame.
            schema: The USD IsA schema or applied API schema to filter attributes.
            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
        """
        super().__init__(title, schema, schema_subclasses, include_list, exclude_list)

    def build_attribute_item(self, stage, ui_attr: UsdAttributeUiEntry, prim_paths: List[Sdf.Path]):
        return super().build_property_item(stage, ui_attr, prim_paths)

    def _filter_attrs_to_build(self, attrs):
        return super()._filter_props_to_build(attrs)

    def _customize_attrs_layout(self, attrs):
        return super()._customize_props_layout(attrs)

    def _get_shared_attributes_from_selected_prims(self, anchor_prim):
        return super()._get_shared_properties_from_selected_prims(anchor_prim)

    def build_property_item(self, stage, ui_prop: UsdPropertyUiEntry, prim_paths: List[Sdf.Path]):
        return self.build_attribute_item(stage, ui_prop, prim_paths)

    def _filter_props_to_build(self, props):
        return self._filter_attrs_to_build(props)

    def _customize_props_layout(self, props):
        return self._customize_attrs_layout(props)

    def _get_shared_properties_from_selected_prims(self, anchor_prim):
        return self._get_shared_attributes_from_selected_prims(anchor_prim)
