# Copyright (c) 2018-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.
#
__all__ = ["DefaultSelectionWatch"]

from pxr import Sdf, Trace
import omni.usd
from typing import Optional


class DefaultSelectionWatch(object):
    """
    The object that update selection in TreeView when the scene selection is
    changed and updated scene selection when TreeView selection is changed.
    """

    def __init__(self, tree_view=None, usd_context=None):
        self._usd_context = usd_context or omni.usd.get_context()
        self._selection = None
        self._in_selection = False
        self._tree_view = None
        self._selection = self._usd_context.get_selection()
        self._events = self._usd_context.get_stage_event_stream()
        self._stage_event_sub = self._events.create_subscription_to_pop(
            self._on_stage_event, name="Stage Window Selection Update"
        )

        self.set_tree_view(tree_view)

        self.__filter_string: Optional[str] = None
        # When True, SelectionWatch should consider filtering
        self.__filter_checking: bool = False

    def destroy(self):
        self._usd_context = None
        self._selection = None
        self._stage_event_sub = None
        self._events = None
        if self._tree_view:
            self._tree_view.set_selection_changed_fn(None)
            self._tree_view = None

    def set_tree_view(self, tree_view):
        """Replace TreeView that should show the selection"""
        if self._tree_view != tree_view:
            if self._tree_view:
                self._tree_view.set_selection_changed_fn(None)

            self._tree_view = tree_view
            self._tree_view.set_selection_changed_fn(self._on_widget_selection_changed)
            self._last_selected_prim_paths = None
            self._on_kit_selection_changed()

    def set_filtering(self, filter_string: Optional[str]):
        if filter_string:
            self.__filter_string = filter_string.lower()
        else:
            self.__filter_string = filter_string

    def enable_filtering_checking(self, enable: bool):
        """
        When `enable` is True, SelectionWatch should consider filtering when
        changing Kit's selection.
        """
        self.__filter_checking = enable

    def _on_stage_event(self, event):
        """Called by stage_event_stream"""
        if event.type == int(omni.usd.StageEventType.SELECTION_CHANGED):
            self._on_kit_selection_changed()

    @Trace.TraceFunction
    def _on_kit_selection_changed(self):
        """Send the selection from Kit to TreeView"""
        if not self._tree_view or self._in_selection:
            return

        # Make sure it's a new selection. It happens that omni.usd sends the same selection twice. No sorting because
        # the order of selection is important.
        prim_paths = self._selection.get_selected_prim_paths()
        if prim_paths == self._last_selected_prim_paths:
            return

        self._last_selected_prim_paths = [Sdf.Path(prim_path) for prim_path in prim_paths]

        # Pump the changes to the model because to select something, TreeView should be updated.
        self._tree_view.model.update_dirty()

        selection = []
        for path in prim_paths:
            # Get the selected item and its parents. Expand all the parents of the new selection.
            full_chain = self._tree_view.model.find_full_chain(path)
            # When the new object is created, omni.usd sends the selection is changed before the object appears in the
            # stage and it means we can't select it. In this way it's better to return because we don't have the item
            # in the model yet.
            # TODO: Use UsdNotice to track if the object is created.
            if not full_chain or full_chain[-1].path != path:
                continue
            if full_chain:
                for item in full_chain[:-1]:
                    self._tree_view.set_expanded(item, True, False)
                # Save the last item in the chain. It's the selected item
                selection.append(full_chain[-1])

        # Send all of this to TreeView.
        self._in_selection = True
        self._tree_view.selection = selection
        self._in_selection = False

    @Trace.TraceFunction
    def _on_widget_selection_changed(self, selection):
        """Send the selection from TreeView to Kit"""
        if self._in_selection:
            return

        prim_paths = [item.path for item in selection if item]

        # Filter selection
        if self.__filter_string or self.__filter_checking:
            # Check if the selected prims are filtered and re-select filtered items only if necessary.
            filtered_paths = [item.path for item in selection if item and item.filtered]
            if filtered_paths != prim_paths:
                filtered_selection = [item for item in selection if item and item.path in filtered_paths]
                self._tree_view.selection = filtered_selection
                return

        if prim_paths == self._last_selected_prim_paths:
            return

        # Send the selection to Kit
        self._in_selection = True
        old_paths = [sdf_path.pathString for sdf_path in self._last_selected_prim_paths]
        new_paths = [sdf_path.pathString for sdf_path in prim_paths]
        omni.kit.commands.execute("SelectPrims", old_selected_paths=old_paths, new_selected_paths=new_paths, expand_in_stage=True)
        self._in_selection = False
        self._last_selected_prim_paths = prim_paths
