# Copyright (c) 2022, 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 __future__ import annotations

from typing import TYPE_CHECKING, Dict, Set

import carb.events
import carb.settings
import omni.usd
from pxr import Sdf

if TYPE_CHECKING:
    from .manipulator_base import ManipulatorBase, ManipulatorOrderManager


class ManipulatorSelector:
    def __init__(self, order_manager: ManipulatorOrderManager, usd_context_name: str):
        self._order_manager = order_manager

        self._usd_context_name = usd_context_name
        self._manipulators: Dict[str, Set[ManipulatorBase]] = {}
        self._usd_context = omni.usd.get_context(usd_context_name)

        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop_by_type(
            int(omni.usd.StageEventType.SELECTION_CHANGED),
            self._on_stage_selection_event,
            name="ManipulatorSelector stage event",
        )

        self._order_sub = self._order_manager.subscribe_to_orders_changed(self._on_orders_changed)
        self._refresh()

    def destroy(self):
        self._stage_event_sub = None
        if self._order_sub:
            self._order_manager.unsubscribe_to_orders_changed(self._order_sub)
            self._order_sub = None

    def __del__(self):
        self.destroy()

    def register_manipulator_instance(self, name: str, manipulator: ManipulatorBase):
        needs_sort = False
        if name not in self._manipulators:
            self._manipulators[name] = set()
            needs_sort = True

        self._manipulators[name].add(manipulator)

        if needs_sort:
            self._sort()

        self._refresh()

    def unregister_manipulator_instance(self, name: str, manipulator: ManipulatorBase):
        manipulator_set = self._manipulators.get(name, set())
        manipulator_set.remove(manipulator)

        self._refresh()

    def _sort(self) -> bool:
        orders_dict = self._order_manager.orders_dict

        # sort by order
        sorted_manipulators = dict(sorted(self._manipulators.items(), key=lambda item: orders_dict.get(item[0], 0)))

        # compare keys to check order difference (direct dicts compare are equal if their content is same but in
        # different order)
        if sorted_manipulators != self._manipulators or list(sorted_manipulators.keys()) != list(
            self._manipulators.keys()
        ):
            self._manipulators = sorted_manipulators
            return True
        return False

    def _refresh(self):
        if self._usd_context.get_stage_state() == omni.usd.StageState.OPENED:
            self._on_selection_changed()

    def _on_stage_selection_event(self, event: carb.events.IEvent):
        self._on_selection_changed()

    def _on_selection_changed(self):
        selection = self._usd_context.get_selection().get_selected_prim_paths()
        selection_sdf = [Sdf.Path(path) for path in selection]
        for _, manipulators in self._manipulators.items():
            handled = False
            for manipulator in manipulators:
                if manipulator.on_selection_changed(self._usd_context.get_stage(), selection_sdf):
                    manipulator.enabled = True
                    handled = True
                else:
                    manipulator.enabled = False
            if handled:
                # Set selection_sdf to None to signal subsequent manipulator the selection has been handled
                # This is different from being empty []
                selection_sdf = None

    def _on_orders_changed(self):
        if self._sort():
            self._refresh()
