import asyncio
import carb
import time
import omni.client
import omni.usd
import omni.kit.app as app
import omni.kit.usd.layers as layers
import omni.kit.viewport.menubar.camera as viewport_camera_menu
import omni.ui as ui
import omni.kit.notification_manager as nm
import omni.kit.collaboration.presence_layer as pl

from omni.kit.viewport.menubar.core import ViewportMenuDelegate
from pxr import Sdf, Usd, UsdGeom, Tf, Trace
from typing import Dict
from functools import partial

from .utils import get_viewport_fps, get_viewport_bound_camera_position, switch_viewport_bound_camera_to_persp

from .local_user_shared_data import LocalUserSharedData
from .peer_user_shared_data import PeerUserSharedData
from .peer_user_menu_item import PeerUserMenuItem, NoPeerUserMenuItem
from .menubar_camera_widget_delegate import (
    MenuBarLiveSessionCameraFollowerButtonDelegate,
    MenuBarLiveSessionCameraFollowerMenuItemDelegate
)


UPDATE_DELAY_FRAME_SETTING = "/exts/omni.kit.collaboration.viewport.camera/visual/delayFrame"


class StageListener:
    def __init__(self, usd_context_name=""):
        self.__usd_context = omni.usd.get_context(usd_context_name)
        self.__live_syncing = layers.get_live_syncing(self.__usd_context)
        self.__layers = layers.get_layers(self.__usd_context)
        self.__app = app.get_app()
        self.__settings = carb.settings.get_settings()

        self.__update_subscription = None
        self.__layers_event_subscriptions = []
        self.__local_stage_objects_changed = None

        self.__peer_users: Dict[str, PeerUserSharedData] = {}
        self.__my_shared_data: LocalUserSharedData = None

        self.__pending_local_bound_camera_changes = False
        self.__pending_changed_user_ids = set()
        self.__delay_dirty_tasks_or_futures = {}

        self.__last_viewport_camera_position = None
        # Key is the id of the viewport
        self.__follow_user_menu_collections = {}
        self.__last_time_in_seconds_to_broadcast_builtin_camera = 0
        self.__menubar_camera_button_delegate = MenuBarLiveSessionCameraFollowerButtonDelegate()
        self.__menubar_camera_menu_item_delegate = MenuBarLiveSessionCameraFollowerMenuItemDelegate()
        self.__entering_or_quitting_follow_mode = False
        self.__follow_user_menu_registered = False

    def start(self):

        for event in [
            layers.LayerEventType.LIVE_SESSION_STATE_CHANGED,
            layers.LayerEventType.LIVE_SESSION_USER_JOINED,
            layers.LayerEventType.LIVE_SESSION_USER_LEFT,
            pl.PresenceLayerEventType.BOUND_CAMERA_CHANGED,
            pl.PresenceLayerEventType.BOUND_CAMERA_PROPERTIES_CHANGED,
            pl.PresenceLayerEventType.BOUND_CAMERA_RESYNCED,
            pl.PresenceLayerEventType.LOCAL_FOLLOW_MODE_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.viewport.camera {str(event)}"
            )
            self.__layers_event_subscriptions.append(layers_event_sub)

        stage_url = self.__usd_context.get_stage_url()
        if not stage_url.startswith("omniverse:"):
            return

        self.__on_session_state_changed()

    def __start_subscription(self):
        current_session = self.__live_syncing.get_current_live_session()
        if not current_session:
            return

        if not current_session.logged_user_name or not current_session.logged_user_id:
            carb.log_error(f"Failed to track users camera as no logged username/id is detected.")
            return

        self.__local_stage_objects_changed = Tf.Notice.Register(
            Usd.Notice.ObjectsChanged, self.__on_local_usd_changed, self.__usd_context.get_stage()
        )

        # OMFP-1986: All users will join session with persp camera.
        switch_viewport_bound_camera_to_persp()
        self.__last_viewport_camera_position = get_viewport_bound_camera_position(self.__usd_context.get_stage())

        self.__my_shared_data = LocalUserSharedData(
            self.__usd_context,
            current_session.logged_user_name,
            current_session.logged_user_id
        )
        self.__my_shared_data.broadcast_local_changes()

        for peer_user in current_session.peer_users:
            self.__track_new_user(peer_user.user_id)

        self.__update_subscription = self.__app.get_update_event_stream().create_subscription_to_pop(
            self.__on_update, name="omni.kit.collaboration.viewport.camera update"
        )

        self.__stage_subscription = self.__usd_context.get_stage_event_stream().create_subscription_to_pop(
            self.__on_stage_event, name="omni.kit.collaboration.viewport.camera stage event"
        )

    def __on_stage_event(self, event):
        if event.type == int(omni.usd.StageEventType.CLOSING):
            self.__stop_subscription()

    def __stop_subscription(self):
        if self.__peer_users:
            for _, camera_info in self.__peer_users.items():
                camera_info.destroy()
            self.__peer_users.clear()

        self.__pending_local_bound_camera_changes = False
        self.__pending_changed_user_ids.clear()

        for task_or_future in self.__delay_dirty_tasks_or_futures.values():
            task_or_future.cancel()
        self.__delay_dirty_tasks_or_futures.clear()

        self.__update_subscription = None
        self.__stage_subscription = None

        if self.__my_shared_data:
            self.__my_shared_data.destroy()
            self.__my_shared_data = None

        if self.__local_stage_objects_changed:
            self.__local_stage_objects_changed.Revoke()
            self.__local_stage_objects_changed = None

        self.__deregister_follow_user_menu()

        self.__last_time_in_seconds_to_broadcast_builtin_camera = 0

    def stop(self):
        self.__stop_subscription()
        self.__layers_event_subscriptions = []

        if self.__menubar_camera_button_delegate:
            self.__menubar_camera_button_delegate.destroy()
            self.__menubar_camera_button_delegate = None

        if self.__menubar_camera_menu_item_delegate:
            self.__menubar_camera_menu_item_delegate.destroy()
            self.__menubar_camera_menu_item_delegate = None

    # Adds a delay to the visual update during translate (only) manipulation
    # It's due to the renderer having a delay of rendering the mesh and the manipulator appears to drift apart.
    # It's only an estimation and may varying from scene/renderer setup.
    @Trace.TraceFunction
    @carb.profiler.profile
    async def __delay_dirty(self, changed_user_ids, update_id, camera_positions_before_delay):
        fps = get_viewport_fps()
        delay_frames = self.__settings.get(UPDATE_DELAY_FRAME_SETTING)
        if fps:
            render_frame_time = 1.0 / fps * delay_frames
            while True:
                dt = await self.__app.next_update_async()
                render_frame_time -= dt
                # break a frame early
                if render_frame_time < dt:
                    break

        # cancel earlier job if a later one catches up (fps suddenly changed?)
        earlier_tasks_or_futures = []
        for key, task_or_future in self.__delay_dirty_tasks_or_futures.items():
            if key < update_id:
                earlier_tasks_or_futures.append(key)
                task_or_future.cancel()
            else:
                break

        for key in earlier_tasks_or_futures:
            self.__delay_dirty_tasks_or_futures.pop(key)

        for user_id in changed_user_ids:
            peer_user = self.__peer_users.get(user_id, None)
            if peer_user:
                position = camera_positions_before_delay.get(user_id, None)
                peer_user.update_camera_widget_label(position)

        self.__delay_dirty_tasks_or_futures.pop(update_id)

    def __is_path_properties_affected(self, notice, path):
        changed_info_only_paths = notice.GetChangedInfoOnlyPaths()
        slice_object = Sdf.Path.FindPrefixedRange(changed_info_only_paths, path)
        if not slice_object:
            return False

        paths = changed_info_only_paths[slice_object]

        return True if paths else False

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

        local_bound_camera_path = self.__my_shared_data.bound_camera_path
        if not local_bound_camera_path:
            return

        bound_camera = stage.GetPrimAtPath(local_bound_camera_path)
        if not bound_camera:
            return

        if self.__pending_local_bound_camera_changes:
            return

        if notice.ResyncedObject(bound_camera) or self.__is_path_properties_affected(notice, local_bound_camera_path):
            self.__pending_local_bound_camera_changes = True
        else:
            # If peer users are bound to non-builtin cameras, and they are changed locally.
            # The camera label needs to be updated, too.
            for user_id, peer_user in self.__peer_users.items():
                if peer_user.is_bound_to_builtin_camera:
                    continue

                camera_prim = peer_user.bound_camera_prim
                if not camera_prim:
                    peer_user.destroy_camera_widget()
                    continue

                # Adds it into resynced last although it only updates camera label since it's bound
                # to non-builtin camera.
                camera_path = camera_prim.GetPath()
                if notice.ResyncedObject(camera_prim) or self.__is_path_properties_affected(notice, camera_path):
                    self.__pending_changed_user_ids.add(user_id)

    def __on_session_state_changed(self):
        if not self.__live_syncing.is_in_live_session():
            self.__stop_subscription()
        else:
            self.__register_follow_user_menu()
            self.__start_subscription()

    def __track_new_user(self, user_id):
        if user_id in self.__peer_users:
            return

        current_session = self.__live_syncing.get_current_live_session()
        if not current_session:
            return

        user_info = current_session.get_peer_user_info(user_id)
        if not user_info:
            return

        peer_user = PeerUserSharedData(
            self.__usd_context,
            user_info,
            self.__my_shared_data
        )
        self.__peer_users[user_id] = peer_user

        peer_user.refresh_camera_widget()

        # Uses cached viewport camera position to avoid gettting viewport camera pos for each camera,
        # which is very slow.
        peer_user.check_and_switch_widget_mode(self.__last_viewport_camera_position)

    def __untrack_user(self, user_id):
        peer_user = self.__peer_users.pop(user_id, None)
        if peer_user:
            peer_user.followed = False
            peer_user.destroy()

    def __on_layers_event(self, event):
        payload = layers.get_layer_event_payload(event)
        if not payload:
            return

        if payload.event_type == layers.LayerEventType.LIVE_SESSION_STATE_CHANGED:
            if not payload.is_layer_influenced(self.__usd_context.get_stage_url()):
                return

            self.__on_session_state_changed()
        elif payload.event_type == layers.LayerEventType.LIVE_SESSION_USER_JOINED:
            if not payload.is_layer_influenced(self.__usd_context.get_stage_url()):
                return

            self.__track_new_user(payload.user_id)
        elif payload.event_type == layers.LayerEventType.LIVE_SESSION_USER_LEFT:
            if not payload.is_layer_influenced(self.__usd_context.get_stage_url()):
                return

            self.__untrack_user(payload.user_id)
        else:
            payload = pl.get_presence_layer_event_payload(event)
            if not payload.event_type:
                return

            if (
                payload.event_type == pl.PresenceLayerEventType.BOUND_CAMERA_CHANGED or
                payload.event_type == pl.PresenceLayerEventType.BOUND_CAMERA_PROPERTIES_CHANGED or
                payload.event_type == pl.PresenceLayerEventType.BOUND_CAMERA_RESYNCED
            ):
                self.__pending_changed_user_ids.update(payload.changed_user_ids)
            elif payload.event_type == pl.PresenceLayerEventType.LOCAL_FOLLOW_MODE_CHANGED:
                if self.__entering_or_quitting_follow_mode:
                    return

                self.__entering_or_quitting_follow_mode = True

                try:
                    following_user_id = self.__my_shared_data.presence_layer.get_following_user_id()
                    if following_user_id and following_user_id == self.__my_shared_data.following_user_id:
                        return

                    # Unfollow previous user
                    current_followed_user = self.__peer_users.get(self.__my_shared_data.following_user_id, None)
                    if current_followed_user:
                        current_followed_user.followed = False
                        self.__pending_changed_user_ids.add(self.__my_shared_data.following_user_id)

                    if following_user_id:
                        self.__follow_user(following_user_id)
                finally:
                    self.__entering_or_quitting_follow_mode = False

    @Trace.TraceFunction
    @carb.profiler.profile
    def __on_update(self, dt):
        if not self.__live_syncing.is_in_live_session():
            return

        current_time = time.monotonic()
        duration = current_time - self.__last_time_in_seconds_to_broadcast_builtin_camera
        # Update camera pos in 30fps for now
        if duration < 0.03:
            return

        self.__last_time_in_seconds_to_broadcast_builtin_camera = current_time

        # FIXME: Delay one frame to try to keep pace ui.scene and geometry rendering,
        # although it's not ideal for only 1 frame.
        all_changed_user_ids = set()
        if self.__pending_changed_user_ids:
            all_changed_user_ids.update(self.__pending_changed_user_ids)
            delay_frame = self.__settings.get(UPDATE_DELAY_FRAME_SETTING) != 0

            camera_positions_before_delay = {}
            for user_id in self.__pending_changed_user_ids:
                peer_user = self.__peer_users.get(user_id, None)
                if not peer_user:
                    continue

                # It will not update camera label when it's to delay frames.
                peer_user.refresh_camera_widget(not delay_frame)
                if delay_frame:
                    camera_positions_before_delay[user_id] = peer_user.widget_position

            self.__pending_changed_user_ids.clear()

            if delay_frame:
                update_id = self.__app.get_update_number()
                try:
                    # run_coroutine only works for Kit 105 or above.
                    from omni.kit.async_engine import run_coroutine
                    task_or_future = run_coroutine(
                        self.__delay_dirty(all_changed_user_ids, update_id, camera_positions_before_delay)
                    )
                except ImportError:
                    # Fallback to asyncio if run_coroutine cannot be imported.
                    task_or_future = asyncio.ensure_future(
                        self.__delay_dirty(all_changed_user_ids, update_id, camera_positions_before_delay)
                    )
                self.__delay_dirty_tasks_or_futures[update_id] = task_or_future

        if self.__my_shared_data and self.__my_shared_data.update_viewport_bound_camera():
            user_id = self.__my_shared_data.following_user_id
            if user_id:
                followed_user = self.__peer_users.get(user_id, None)
                if followed_user:
                    followed_user.followed = False
            self.__pending_local_bound_camera_changes = True
        elif self.__my_shared_data.update_viewport_size():
            self.__pending_local_bound_camera_changes = True

        if self.__pending_local_bound_camera_changes:
            self.__last_viewport_camera_position = get_viewport_bound_camera_position(self.__usd_context.get_stage())
            local_bound_camera_path = self.__my_shared_data.bound_camera_path
            if not local_bound_camera_path:
                return

            for user_id, peer_user in self.__peer_users.items():
                if user_id in all_changed_user_ids:
                    continue

                peer_user.check_and_switch_widget_mode(self.__last_viewport_camera_position)

            try:
                self.__my_shared_data.broadcast_local_changes()
            finally:
                self.__pending_local_bound_camera_changes = False

    def __build_cameras(self, viewport_api, ui_shown: bool):
        """Build the menu with the list of the users"""
        current_session = self.__live_syncing.get_current_live_session()
        if not current_session:
            return

        follow_user_menu_collection = self.__follow_user_menu_collections.get(viewport_api.id, None)
        if not follow_user_menu_collection:
            return

        follow_user_menu_collection.clear()
        if ui_shown:
            with follow_user_menu_collection:
                if current_session.peer_users:
                    for peer_user in current_session.peer_users:
                        item = PeerUserMenuItem(
                            peer_user,
                            lambda user: self.__follow_user(user.user_id)
                        )
                        followed = peer_user.user_id == self.__my_shared_data.following_user_id
                        item.set_checked(followed)
                else:
                    NoPeerUserMenuItem()

    def __follow_user(self, user_id):
        peer_user = self.__peer_users.get(user_id, None)
        if not peer_user:
            return

        following_user_id = peer_user.following_user_id
        following_user = self.__peer_users.get(following_user_id, None)
        if following_user:
            if following_user_id == self.__my_shared_data.user_id:
                user_name = "me"
            else:
                user_name = following_user.user_name

            nm.post_notification(
                f"Cannot follow user {peer_user.user_name} as the user is following {user_name}.",
                status=nm.NotificationStatus.WARNING
            )

            return

        if not peer_user.bound_camera_prim:
            nm.post_notification(
                f"Cannot follow user {peer_user.user_info.user_name} as the user has invalid bound camera.",
                status=nm.NotificationStatus.WARNING
            )
            return

        # Unfollow previous user
        current_followed_user = self.__peer_users.get(self.__my_shared_data.following_user_id, None)
        if current_followed_user:
            current_followed_user.followed = False
            self.__pending_changed_user_ids.add(self.__my_shared_data.following_user_id)

        if peer_user != current_followed_user:
            peer_user.followed = True

    def __build_follow_user_menu(self, viewport_context, root: ui.Menu):
        viewport_api = viewport_context.get("viewport_api")
        follow_user_menu_collection = ui.MenuItemCollection(
            "Follow User",
            checkable=True,
            delegate=ViewportMenuDelegate(),
            shown_changed_fn=partial(self.__build_cameras, viewport_api),
            hide_on_click=False
        )

        self.__follow_user_menu_collections[viewport_api.id] = follow_user_menu_collection

    def __register_follow_user_menu(self):
        def is_follow_enabled():
            settings = carb.settings.get_settings()
            enabled = settings.get(f"/app/liveSession/enableMenuFollowUser")
            if enabled == True or enabled == False:
                return enabled
            return True

        if not is_follow_enabled() or self.__follow_user_menu_registered:
            return

        self.__follow_user_menu_registered = True
        instance = viewport_camera_menu.get_instance()
        instance.register_menu_item(self.__build_follow_user_menu, order=-2)

    def __deregister_follow_user_menu(self):
        if not self.__follow_user_menu_registered:
            return
        self.__follow_user_menu_registered = False
        instance = viewport_camera_menu.get_instance()
        instance.deregister_menu_item(self.__build_follow_user_menu)
        self.__follow_user_menu_collections.clear()
