# 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 weakref
from pathlib import Path
from typing import List

import carb
import omni.ext
import omni.kit.app
import omni.ui as ui
import omni.usd
from pxr import Sdf, UsdGeom

from .attribute_context_menu import AttributeContextMenu
from .control_state_manager import ControlStateManager
from .prim_path_widget import PrimPathWidget
from .prim_selection_payload import PrimSelectionPayload

ICON_PATH = ""
TEST_DATA_PATH = ""


class UsdPropertyWidgets(omni.ext.IExt):
    def __init__(self):
        self._registered = False
        self._examples = None
        self._selection_notifiers = []

    def on_startup(self, ext_id):
        manager = omni.kit.app.get_app().get_extension_manager()
        extension_path = manager.get_extension_path(ext_id)
        global ICON_PATH
        ICON_PATH = Path(extension_path).joinpath("data").joinpath("icons")
        global TEST_DATA_PATH
        TEST_DATA_PATH = Path(extension_path).joinpath("data").joinpath("tests")
        self._selection_notifiers.append(SelectionNotifier())  # default context
        self._usd_preferences = None
        self._hooks = []
        self._attribute_context_menu = AttributeContextMenu()
        self._control_state_manager = ControlStateManager(ICON_PATH)

        self._hooks.append(
            manager.subscribe_to_extension_enable(
                lambda _: self._register_widget(),
                lambda _: self._unregister_widget(),
                ext_name="omni.kit.window.property",
                hook_name="omni.usd listener",
            )
        )

        from .usd_property_widget_builder import UsdPropertiesWidgetBuilder

        UsdPropertiesWidgetBuilder.startup()

        manager = omni.kit.app.get_app().get_extension_manager()
        self._hooks.append(
            manager.subscribe_to_extension_enable(
                on_enable_fn=lambda _: self._register_preferences(),
                on_disable_fn=lambda _: self._unregister_preferences(),
                ext_name="omni.kit.window.preferences",
                hook_name="omni.kit.property.usd omni.kit.window.preferences listener",
            )
        )

    def on_shutdown(self):
        from .usd_property_widget_builder import UsdPropertiesWidgetBuilder

        UsdPropertiesWidgetBuilder.shutdown()

        self._control_state_manager.destory()
        self._control_state_manager = None

        self._attribute_context_menu.destroy()
        self._attribute_context_menu = None

        for notifier in self._selection_notifiers:
            notifier.stop()
        self._selection_notifiers.clear()
        self._hooks = None

        if self._registered:
            self._unregister_widget()

        self._unregister_preferences()

    def _register_preferences(self):
        from omni.kit.window.preferences import register_page

        from .property_preferences_page import PropertyUsdPreferences

        self._usd_preferences = omni.kit.window.preferences.register_page(PropertyUsdPreferences())

    def _unregister_preferences(self):
        if self._usd_preferences:
            import omni.kit.window.preferences

            omni.kit.window.preferences.unregister_page(self._usd_preferences)
            self._usd_preferences = None

    def _register_widget(self):
        try:
            import omni.kit.window.property as p

            from .references_widget import PayloadReferenceWidget
            from .usd_property_widget import RawUsdPropertiesWidget, UsdPropertiesWidget
            from .variants_widget import VariantsWidget

            w = p.get_window()
            if w:
                w.register_widget("prim", "path", PrimPathWidget())
                w.register_widget("prim", "references", PayloadReferenceWidget())
                w.register_widget("prim", "payloads", PayloadReferenceWidget(use_payloads=True))
                w.register_widget("prim", "variants", VariantsWidget())
                w.register_widget(
                    "prim", "attribute", RawUsdPropertiesWidget(title="Raw USD Properties", collapsed=True), False
                )

                # A few examples. Expected to be removed at some point.
                # self._examples = Examples(w)
                for notifier in self._selection_notifiers:
                    notifier.start()
                    notifier._notify_property_window()  # force a refresh
                self._registered = True
        except Exception as exc:
            carb.log_warn(f"error {exc}")

    def _unregister_widget(self):
        try:
            import omni.kit.window.property as p

            w = p.get_window()
            if w:
                for notifier in self._selection_notifiers:
                    notifier.stop()
                w.unregister_widget("prim", "attribute")
                w.unregister_widget("prim", "variants")
                w.unregister_widget("prim", "references")
                w.unregister_widget("prim", "payloads")
                w.unregister_widget("prim", "path")
                if self._examples:
                    self._examples.clean(w)
                    self._examples = None
                self._registered = False
        except Exception as e:
            carb.log_warn(f"Unable to unregister omni.usd.widget: {e}")


class SelectionNotifier:
    def __init__(self, usd_context_id="", property_window_context_id=""):
        self._usd_context = omni.usd.get_context(usd_context_id)
        self._selection = self._usd_context.get_selection()
        self._property_window_context_id = property_window_context_id

    def start(self):
        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop(
            self._on_stage_event, name="omni.usd.widget"
        )

    def stop(self):
        self._stage_event_sub = None

    def _on_stage_event(self, event):
        if event.type == int(omni.usd.StageEventType.SELECTION_CHANGED) or event.type == int(
            omni.usd.StageEventType.CLOSING
        ):
            self._notify_property_window()

    def _notify_property_window(self):
        import omni.kit.window.property as p

        # TODO _property_window_context_id
        w = p.get_window()
        if w and self._usd_context:
            stage = None
            selected_paths = []
            if self._usd_context.get_stage_state() == omni.usd.StageState.OPENED:
                stage = weakref.ref(self._usd_context.get_stage())
                selected_paths = [Sdf.Path(path) for path in self._selection.get_selected_prim_paths()]
            payload = PrimSelectionPayload(stage, selected_paths)
            w.notify("prim", payload)


class Examples: # pragma: no cover
    """
    Examples of showing how to use PropertyWindow APIs.
    """

    def __init__(self, w):
        from omni.kit.window.property.property_scheme_delegate import PropertySchemeDelegate

        from .usd_property_widget import SchemaPropertiesWidget, UsdPropertiesWidget

        # Example of PropertySchemeDelegate
        class MeshWidget(SchemaPropertiesWidget):
            def __init__(self):
                super().__init__("Mesh", UsdGeom.Mesh, False)

            def _group_helper(self, attr, group_prefix):
                if attr.attr_name.startswith(group_prefix):
                    import re

                    attr.override_display_group(group_prefix.capitalize())
                    attr.override_display_name(re.sub(r"(\w)([A-Z])", r"\1 \2", attr.attr_name[len(group_prefix) :]))
                    return True
                return False

            def _filter_props_to_build(self, attrs):
                schema_attr_names = self._schema.GetSchemaAttributeNames(self._include_inherited)
                schema_attr_names.append("material:binding")
                return [attr for attr in attrs if attr.GetName() in schema_attr_names]

            # Example of how to use _customize_props_layout
            def _customize_props_layout(self, attrs):
                auto_arrange = False
                if auto_arrange:
                    for attr in attrs:
                        if not self._group_helper(attr, "crease"):
                            if not self._group_helper(attr, "corner"):
                                if not self._group_helper(attr, "face"):
                                    attr.override_display_group("Other")
                    return attrs
                else:
                    from omni.kit.property.usd.usd_property_widget_builder import UsdPropertiesWidgetBuilder
                    from omni.kit.window.property.templates import (HORIZONTAL_SPACING, LABEL_HEIGHT, LABEL_WIDTH,
                                                                    SimplePropertyWidget)

                    from .custom_layout_helper import CustomLayoutFrame, CustomLayoutGroup, CustomLayoutProperty

                    frame = CustomLayoutFrame(hide_extra=False)
                    with frame:
                        with CustomLayoutGroup("Face"):
                            CustomLayoutProperty("faceVertexCounts", "Vertex Counts")
                            CustomLayoutProperty("faceVertexIndices", "Vertex Indices")
                        with CustomLayoutGroup("Corner"):
                            CustomLayoutProperty("cornerIndices", "Indices")
                            CustomLayoutProperty("cornerSharpnesses", "Sharpnesses")
                        with CustomLayoutGroup("Crease"):
                            CustomLayoutProperty("creaseIndices", "Indices")
                            CustomLayoutProperty("creaseLengths", "Lengths")
                        with CustomLayoutGroup("CustomItems"):
                            with CustomLayoutGroup("Nested"):

                                def build_fn(
                                    stage,
                                    attr_name,
                                    metadata,
                                    property_type,
                                    prim_paths: List[Sdf.Path],
                                    additional_label_kwargs={},
                                    additional_widget_kwarg={},
                                ):
                                    with ui.HStack(spacing=HORIZONTAL_SPACING):
                                        UsdPropertiesWidgetBuilder._create_label(
                                            "Cool Label", additional_label_kwargs={"tooltip": "Cool Tooltip"}
                                        )
                                        ui.Button("Cool Button")
                                    with ui.HStack(spacing=HORIZONTAL_SPACING):
                                        UsdPropertiesWidgetBuilder._create_label("Very Cool Label")
                                        ui.Button("Super Cool Button")

                                CustomLayoutProperty(None, None, build_fn=build_fn)
                    return frame.apply(attrs)

            def get_additional_kwargs(self, ui_attr):
                if ui_attr[0] == "material:binding":
                    from pxr import UsdShade

                    return None, {"target_picker_filter_type_list": [UsdShade.Material], "targets_limit": 1}
                return None, None

        w.register_widget("prim", "mesh", MeshWidget())

    def clean(self, w):
        pass
