import asyncio
import carb
import omni.client
import omni.usd
import omni.kit.usd.layers as layers
import omni.kit.collaboration.presence_layer as pl

from omni.kit.async_engine import run_coroutine

from .manipulator import PeerUserList
from pxr import Sdf, Usd, UsdGeom, Tf, Trace
from typing import Dict, List, Union


SETTING_MULTI_USER_OUTLINE = "/app/viewport/multiUserOutlineMode"
SETTING_SAVED_OUTLINE_COLOR = "/persistent/app/viewport/live_session/saved_outline_color"
SETTING_OUTLINE_COLOR = "/persistent/app/viewport/outline/color"
SETTING_MAXIMUM_SELECTION_ICONS_PER_USER = "/exts/omni/kit/collaboration/selection_outline/maximum_selection_icons_per_user"


class MultiUserSelectionOutlineManager:
    def __init__(self, usd_context: omni.usd.UsdContext) -> None:
        self.__usd_context = usd_context
        self.__settings = carb.settings.get_settings()
        self.__user_last_selected_prims = {}
        self.__prim_selection_queue = {}
        self.__user_selection_group = {}
        self.__selection_group_pool = []
        self.__prim_viewport_widget = {}
        self._delay_refresh_task_or_future = None

    def clear_all_selection(self):
        self.__user_last_selected_prims = {}
        self.__prim_selection_queue = {}
        self.__user_selection_group = {}

        for widget in self.__prim_viewport_widget.values():
            widget.destroy()

        self.__prim_viewport_widget = {}

        if self._delay_refresh_task_or_future:
            self._delay_refresh_task_or_future.cancel()
            self._delay_refresh_task_or_future = None

    def clear_selection(self, user_id):
        """Deregister user from selection group."""

        selection_group = self.__user_selection_group.pop(user_id, None)
        if selection_group:
            self.__update_selection_outline_internal(user_id, selection_group, {})

            # Returns selection group to re-use as it supports 255 at maximum due to limitation of renderer.
            self.__selection_group_pool.append(selection_group)

    def on_usd_objects_changed(self, notice):
        stage = self.__usd_context.get_stage()
        to_remove_paths = []
        to_update_paths = []
        for widget_path, widget in self.__prim_viewport_widget.items():
            prim = stage.GetPrimAtPath(widget_path)
            if not prim or not prim.IsActive() or not prim.IsDefined():
                to_remove_paths.append(widget_path)
            elif notice.ResyncedObject(prim):
                to_update_paths.append(widget_path)
            else:
                changed_info_only_paths = notice.GetChangedInfoOnlyPaths()
                slice_object = Sdf.Path.FindPrefixedRange(changed_info_only_paths, widget_path)
                if not slice_object:
                    continue

                paths = changed_info_only_paths[slice_object]
                if not paths:
                    continue

                for path in paths:
                    if UsdGeom.Xformable.IsTransformationAffectedByAttrNamed(path.name):
                        to_update_paths.append(widget_path)
                        break

        for path in to_remove_paths:
            widget = self.__prim_viewport_widget.pop(path, None)
            widget.destroy()

        for path in to_update_paths:
            widget = self.__prim_viewport_widget.get(path, None)
            widget.on_transform_changed()

    def refresh_selection(self, selection_paths: List[Sdf.Path]):
        """
        Calls this to refresh selection. This is normally called after selection changed, which will
        override the outline of remote users.
        """
        if not selection_paths:
            return

        for path in selection_paths:
            self.__set_prim_outline(0, path)

        async def delay_refresh(paths):
            await omni.kit.app.get_app().next_update_async()
            for path in paths:
                selection_queue = self.__get_selection_queue(path)
                for _, group in selection_queue.items():
                    self.__set_prim_outline(group, path)
                    break

            self._delay_refresh_task_or_future = None

        if self._delay_refresh_task_or_future and not self._delay_refresh_task_or_future.done():
            self._delay_refresh_task_or_future.cancel()
        self._delay_refresh_task_or_future = run_coroutine(delay_refresh(selection_paths))

    def update_selection_outline(self, user_info: layers.LiveSessionUser, selection_paths: List[Sdf.Path]):
        user_id = user_info.user_id
        selection_group = self.__user_selection_group.get(user_id, None)
        if not selection_group:
            # Re-use selection group.
            if self.__selection_group_pool:
                selection_group = self.__selection_group_pool.pop(0)
            else:
                selection_group = self.__usd_context.register_selection_group()

            self.__user_selection_group[user_id] = selection_group
            # RGBA
            color = (
                user_info.user_color[0] / 255.0,
                user_info.user_color[1] / 255.0,
                user_info.user_color[2] / 255.0,
                1.0
            )
            self.__usd_context.set_selection_group_outline_color(selection_group, color)
            self.__usd_context.set_selection_group_shade_color(selection_group, (0.0, 0.0, 0.0, 0.0))

        if selection_group == 0:
            carb.log_error("The maximum number (255) of selection groups is reached.")
            return

        selection_paths = [Sdf.Path(path) for path in selection_paths]
        selection_paths = Sdf.Path.RemoveDescendentPaths(selection_paths)
        stage = self.__usd_context.get_stage()

        # Find all renderable prims
        all_renderable_prims = {}
        for path in selection_paths:
            prim = stage.GetPrimAtPath(path)
            if not prim:
                continue

            renderable_prims = []
            for prim in Usd.PrimRange(prim, Usd.TraverseInstanceProxies(Usd.PrimAllPrimsPredicate)):
                # Camera in Kit has special representation, which needs to be excluded.
                if (
                    not prim.IsActive()
                    or omni.usd.editor.is_hide_in_ui(prim)
                    or omni.usd.editor.is_hide_in_stage_window(prim)
                ):
                    break

                if (
                    not prim.GetPath().name == "CameraModel"
                    and not prim.GetPath().GetParentPath().name == "OmniverseKitViewportCameraMesh"
                    and prim.IsA(UsdGeom.Gprim)
                ):
                    renderable_prims.append(prim.GetPath())

            if renderable_prims:
                all_renderable_prims[path] = renderable_prims

        selection_paths = all_renderable_prims

        # Draw user icon
        maximum_count = self.__settings.get(SETTING_MAXIMUM_SELECTION_ICONS_PER_USER)
        if not maximum_count or maximum_count <= 0:
            maximum_count = 3

        for index, path in enumerate(selection_paths.keys()):
            path = Sdf.Path(path)
            widget = self.__prim_viewport_widget.get(path, None)
            if not widget:
                widget = PeerUserList(self.__usd_context, path)
                self.__prim_viewport_widget[path] = widget

            widget.add_user(user_info)

            # OMFP-2736: Don't show all selections to avoid cluttering the viewport.
            if index + 1 >= maximum_count:
                break

        self.__update_selection_outline_internal(user_id, selection_group, selection_paths)

    def __update_selection_outline_internal(
        self, user_id, selection_group,
        all_renderable_paths: Dict[Sdf.Path, List[Sdf.Path]]
    ):
        last_selected_prim_and_renderable_paths = self.__user_last_selected_prims.get(user_id, {})
        last_selected_prim_paths = set(last_selected_prim_and_renderable_paths.keys())
        selection_paths = set(all_renderable_paths.keys())
        if last_selected_prim_paths == selection_paths:
            return

        if last_selected_prim_paths:
            old_selected_paths = last_selected_prim_paths.difference(selection_paths)
            new_selected_paths = selection_paths.difference(last_selected_prim_paths)
        else:
            old_selected_paths = set()
            new_selected_paths = selection_paths

        # Update selection outline for old paths
        to_clear_paths = []
        for path in old_selected_paths:
            # Remove user icon
            widget = self.__prim_viewport_widget.get(path, None)
            if widget:
                widget.remove_user(user_id)
                if widget.empty_user_list():
                    self.__prim_viewport_widget.pop(path, None)
                    widget.destroy()

            # Remove outline of all renderable prims under the selected path.
            renderable_paths = last_selected_prim_and_renderable_paths.get(path, None)
            for renderable_path in renderable_paths:
                selection_queue = self.__get_selection_queue(renderable_path)
                if selection_queue:
                    # FIFO queue that first user that selects the prim decides the color of selection.
                    group = next(iter(selection_queue.values()))
                    same_group = selection_group == group

                    # Pop this user and outline the prim with next user's selection color.
                    selection_queue.pop(user_id, None)
                    if same_group and selection_queue:
                        for _, group in selection_queue.items():
                            self.__set_prim_outline(group, renderable_path)
                            break
                    elif same_group:
                        to_clear_paths.append(renderable_path)

        # Update selection outline for new paths
        to_select_paths = []
        for path in new_selected_paths:
            renderable_paths = all_renderable_paths.get(path, None)
            if not renderable_paths:
                continue

            for renderable_path in renderable_paths:
                selection_queue = self.__get_selection_queue(renderable_path)
                has_outline = not not selection_queue
                selection_queue[user_id] = selection_group
                if not has_outline:
                    to_select_paths.append(renderable_path)

        if to_clear_paths:
            self.__clear_prim_outline(to_clear_paths)

        if to_select_paths:
            self.__set_prim_outline(selection_group, to_select_paths)

        self.__user_last_selected_prims[user_id] = all_renderable_paths

    def __get_selection_queue(self, prim_path):
        prim_path = Sdf.Path(prim_path)
        if prim_path not in self.__prim_selection_queue:
            # Value is dict to keep order of insertion.
            self.__prim_selection_queue[prim_path] = {}

        return self.__prim_selection_queue[prim_path]

    def __set_prim_outline(self, selection_group, prim_path: Union[Sdf.Path, List[Sdf.Path]]):
        if isinstance(prim_path, list):
            paths = [str(path) for path in prim_path]
        else:
            paths = [str(prim_path)]

        selection = self.__usd_context.get_selection().get_selected_prim_paths()
        for path in paths:
            # If it's locally selected, don't override the outline.
            if path in selection:
                continue

            self.__usd_context.set_selection_group(selection_group, path)

    def __clear_prim_outline(self, prim_path: Union[Sdf.Path, List[Sdf.Path]]):
        self.__set_prim_outline(0, prim_path)


class StageListener:
    def __init__(self, usd_context_name=""):
        self.__usd_context = omni.usd.get_context(usd_context_name)
        self.__presence_layer = pl.get_presence_layer_interface(self.__usd_context)
        self.__live_syncing = layers.get_live_syncing(self.__usd_context)
        self.__layers = layers.get_layers(self.__usd_context)
        self.__layers_event_subscriptions = []
        self.__stage_event_subscription = None
        self.__settings = carb.settings.get_settings()
        self.__selection_paths_initialized = False
        self.__stage_objects_changed = None
        self.__multi_user_selection_manager = MultiUserSelectionOutlineManager(self.__usd_context)
        self.__delayed_task_or_future = None
        self.__changed_user_ids = set()

    def start(self):
        self.__restore_local_outline_color()

        for event in [
            layers.LayerEventType.LIVE_SESSION_STATE_CHANGED,
            layers.LayerEventType.LIVE_SESSION_USER_JOINED,
            layers.LayerEventType.LIVE_SESSION_USER_LEFT,
            pl.PresenceLayerEventType.SELECTIONS_CHANGED,
        ]:
            layers_event_sub = self.__layers.get_event_stream().create_subscription_to_pop_by_type(
                event, self.__on_layers_event, name=f"omni.kit.collaboration.selection_outline {str(event)}"
            )
            self.__layers_event_subscriptions.append(layers_event_sub)

        self.__stage_event_subscription = self.__usd_context.get_stage_event_stream().create_subscription_to_pop(
            self.__on_stage_event, name="omni.kit.collaboration.selection_outline stage event"
        )

        self.__start_subscription()

    def __start_subscription(self):
        self.__settings.set(SETTING_MULTI_USER_OUTLINE, True)
        live_session = self.__live_syncing.get_current_live_session()
        if not live_session:
            return

        for peer_user in live_session.peer_users:
            selection_paths = self.__presence_layer.get_selections(peer_user.user_id)
            self.__multi_user_selection_manager.update_selection_outline(peer_user, selection_paths)

        if not self.__selection_paths_initialized:
            self.__selection_paths_initialized = True
            selection = self.__usd_context.get_selection().get_selected_prim_paths()
            self.__on_selection_changed(selection)

            root_layer_session = self.__live_syncing.get_current_live_session()
            if root_layer_session:
                logged_user = root_layer_session.logged_user
                user_color = logged_user.user_color
                rgba = (user_color[0] / 255.0, user_color[1] / 255.0, user_color[2] / 255.0, 1.0)
                self.__save_and_set_local_outline_color(rgba)

            self.__stage_objects_changed = Tf.Notice.Register(
                Usd.Notice.ObjectsChanged, self.__on_usd_changed, self.__usd_context.get_stage()
            )

    @Trace.TraceFunction
    @carb.profiler.profile
    def __on_usd_changed(self, notice, sender):
        if not sender or sender != self.__usd_context.get_stage():
            return

        self.__multi_user_selection_manager.on_usd_objects_changed(notice)

    def __stop_subscription(self):
        self.__changed_user_ids.clear()

        if self.__delayed_task_or_future and not self.__delayed_task_or_future.done():
            self.__delayed_task_or_future.cancel()
            self.__delayed_task_or_future = None

        self.__restore_local_outline_color()
        self.__selection_paths_initialized = False
        self.__multi_user_selection_manager.clear_all_selection()
        self.__settings.set(SETTING_MULTI_USER_OUTLINE, False)
        if self.__stage_objects_changed:
            self.__stage_objects_changed.Revoke()
            self.__stage_objects_changed = None

    def stop(self):
        self.__stop_subscription()
        self.__layers_event_subscription = None
        self.__stage_event_subscription = None

    def __save_and_set_local_outline_color(self, color):
        old_color = self.__settings.get(SETTING_OUTLINE_COLOR)
        if old_color:
            self.__settings.set(SETTING_SAVED_OUTLINE_COLOR, old_color)
            old_color[-4:] = color
        else:
            # It has maximum 255 group + selection_outline
            old_color = [0] * 256
            old_color[-4:] = color

        color = old_color

        self.__settings.set(SETTING_OUTLINE_COLOR, color)

    def __restore_local_outline_color(self):
        # It's possible that Kit is shutdown abnormally. So it can
        # be used to restore the outline color before live session or abort.
        color = self.__settings.get(SETTING_SAVED_OUTLINE_COLOR)
        if not color:
            return

        self.__settings.set(SETTING_OUTLINE_COLOR, color)
        # Clear saved color
        self.__settings.set(SETTING_SAVED_OUTLINE_COLOR, None)

    def __on_stage_event(self, event):
        if event.type == int(omni.usd.StageEventType.SELECTION_CHANGED):
            if not self.__live_syncing.get_current_live_session():
                return

            selection = self.__usd_context.get_selection().get_selected_prim_paths()
            self.__on_selection_changed(selection)
        elif event.type == int(omni.usd.StageEventType.CLOSING):
            self.__stop_subscription()

    def __on_selection_changed(self, selection_paths):
        self.__multi_user_selection_manager.refresh_selection(selection_paths)

    def __on_layers_event(self, event):
        payload = layers.get_layer_event_payload(event)
        if payload.event_type == layers.LayerEventType.LIVE_SESSION_STATE_CHANGED:
            if not payload.is_layer_influenced(self.__usd_context.get_stage_url()):
                return

            if self.__live_syncing.get_current_live_session():
                self.__start_subscription()
            else:
                self.__stop_subscription()
        elif payload.event_type == layers.LayerEventType.LIVE_SESSION_USER_JOINED:
            if not payload.is_layer_influenced(self.__usd_context.get_stage_url()):
                return

            user_id = payload.user_id
            live_session = self.__live_syncing.get_current_live_session()
            user_info = live_session.get_peer_user_info(user_id)
            selection_paths = self.__presence_layer.get_selections(user_id)
            self.__multi_user_selection_manager.update_selection_outline(user_info, selection_paths)
        elif payload.event_type == layers.LayerEventType.LIVE_SESSION_USER_LEFT:
            if not payload.is_layer_influenced(self.__usd_context.get_stage_url()):
                return

            user_id = payload.user_id
            self.__multi_user_selection_manager.clear_selection(user_id)
        else:
            payload = pl.get_presence_layer_event_payload(event)
            if payload.event_type == pl.PresenceLayerEventType.SELECTIONS_CHANGED:
                self.__changed_user_ids.update(payload.changed_user_ids)
                # Delay the selection updates.
                if not self.__delayed_task_or_future or self.__delayed_task_or_future.done():
                    self.__delayed_task_or_future = run_coroutine(self.__handle_selection_changes())

    async def __handle_selection_changes(self):
        current_live_session = self.__live_syncing.get_current_live_session()
        if not current_live_session:
            return

        for user_id in self.__changed_user_ids:
            user_info = current_live_session.get_peer_user_info(user_id)
            if not user_info:
                continue

            selection_paths = self.__presence_layer.get_selections(user_id)
            self.__multi_user_selection_manager.update_selection_outline(user_info, selection_paths)

        self.__changed_user_ids.clear()
        self.__delayed_task_or_future = None
