import asyncio
import weakref
import carb
import carb.dictionary
import omni.kit.app
import omni.kit.collaboration.channel_manager as cm
import time

from typing import List, Callable, Awaitable, Dict
from .layer_utils import LayerUtils
from .interface_utils import get_layer_event_payload, post_notification
from .live_session_channel_manager import (
    LiveSessionChannelManager, MESSAGE_SESSION_URL_KEY, MESSAGE_USER_ID_KEY, MESSAGE_USER_NAME_KEY
)

from pxr import Sdf
from .._omni_kit_usd_layers import (
    acquire_live_syncing_interface,
    release_live_syncing_interface,
    ILayersInstance,
    LayerEventType
)


class LiveSession:
    def __init__(self, handle, session_channel, live_syncing):
        self._session_handle = handle
        if session_channel:
            self._session_channel = weakref.ref(session_channel)
        else:
            self._session_channel = None
        self._live_syncing = weakref.ref(live_syncing)

        interface = self._live_syncing()._live_syncing_interface
        layers_instance = self._live_syncing()._layers_instance
        session_handle = self._session_handle
        self._session_url = interface.get_live_session_url(layers_instance, session_handle)
        self._session_name = interface.get_live_session_name(layers_instance, session_handle)
        self._session_root = interface.get_live_session_root_identifier(layers_instance, session_handle)
        self._merge_permission = interface.permission_to_merge_session_changes(layers_instance, session_handle)
        self._owner = interface.get_live_session_owner(layers_instance, session_handle)
        self._channel_url = interface.get_live_session_channel_url(layers_instance, session_handle)
        self._base_layer_identifier = interface.get_live_session_base_layer_identifier(layers_instance, session_handle)
        self._logged_user_name = interface.get_logged_in_user_name_for_layer(layers_instance, self._base_layer_identifier)
        self._logged_user_id = interface.get_logged_in_user_id_for_layer(layers_instance, self._base_layer_identifier)

    @property
    def valid(self):
        if (
            not self._session_handle or
            not self._session_channel or
            not self._session_channel() or
            not self._live_syncing()
        ):
            return False
        
        interface = self._live_syncing()._live_syncing_interface
        layers_instance = self._live_syncing()._layers_instance
        return interface.is_valid_live_session(layers_instance, self._session_handle)

    @property
    def channel_url(self):
        return self._channel_url
    
    @property
    def base_layer_identifier(self):
        return self._base_layer_identifier

    @property
    def name(self):
        return self._session_name
    
    @property
    def url(self):
        return self._session_url
    
    @property
    def owner(self):
        return self._owner

    @property
    def merge_permission(self):
        return self._merge_permission

    @property
    def root(self):
        return self._session_root
    
    @property
    def peer_users(self) -> List[cm.PeerUser]:
        if not self.valid:
            return []

        return list(self._session_channel().peer_users)
    
    @property
    def logged_user_name(self):
        """User name that joins this session."""
        return self._logged_user_name
    
    @property
    def logged_user_id(self):
        """
        User id that joins this session. User id can be used to uniquely identify
        the user in a live session.
        """

        return self._logged_user_id
    
    def __str__(self) -> str:
        return f"<root = {self._session_root}, name = {self._session_name}, url = {self._session_url}>"


class LiveSyncing:
    def __init__(self, layers_instance: ILayersInstance, usd_context, layers_state) -> None:
        self._layers_state = layers_state
        self._layers_instance = layers_instance
        self._usd_context = usd_context
        self._live_syncing_interface = acquire_live_syncing_interface()
        self._dictionary = carb.dictionary.get_dictionary()
        self._layers_event_sub = self._layers_instance.get_event_stream().create_subscription_to_pop(
            self._on_layer_event, name="omni.kit.usd.layers.LiveSyncing"
        )
        self._update_subscription = omni.kit.app.get_app().get_update_event_stream().create_subscription_to_pop(
            self._on_update, name="omni.kit.usd.layers.LiveSyncing Update"
        )
        self._stage_event_sub = self._usd_context.get_stage_event_stream().create_subscription_to_pop(
            self._on_stage_event, name="omni.kit.usd.layers Stage events"
        )
        self._session_channels = {}
        self._merging_live_session = False
        self._all_outdated_layers = set([])
        self._outdate_notification = None
        self._last_outdate_notification_time = 0
        self._joining_session = False
        self._silent_leave_layers = set([])
    
    @property
    def usd_context(self):
        return self._usd_context
    
    def _on_stage_event(self, event):
        if event.type == int(omni.usd.StageEventType.CLOSING):
            self._all_outdated_layers.clear()
            self._last_outdate_notification_time = 0
            self._silent_leave_layers.clear()
            if self._outdate_notification:
                self._outdate_notification.dismiss()
                self._outdate_notification = None
    
    def _on_update(self, dt):
        if self._outdate_notification or not self._all_outdated_layers:
            return

        now = time.time()
        # It will accumulate all changes in 2s.
        if now - self._last_outdate_notification_time < 2:
            return
        
        self._last_outdate_notification_time = now

        try:
            import omni.kit.notification_manager as nm 

            def reload_all_outdated_layers():
                outdated_layers, self._all_outdated_layers = self._all_outdated_layers, set([])
                for outdate_layer in outdated_layers:
                    layer = Sdf.Find(outdate_layer)
                    if layer:
                        layer.Reload()

                self._outdate_notification = None

            self._outdate_notification = nm.post_notification(
                "Base USD files have been changed, please fetch changes.",
                hide_after_timeout=False,
                button_infos=[
                    nm.NotificationButtonInfo("FETCH", on_complete=reload_all_outdated_layers),
                    nm.NotificationButtonInfo("CANCEL", on_complete=None),
                ]
            )
        except Exception:
            pass
    
    def _on_layer_event(self, event: carb.events.IEvent):
        payload = get_layer_event_payload(event)
        if not payload:
            return

        if payload.event_type == LayerEventType.LIVE_SESSION_STATE_CHANGED:
            if self._joining_session or self._merging_live_session:
                return

            for layer_identifier in payload.identifiers_or_spec_paths:
                current_live_session = self.get_current_live_session(layer_identifier)
                if not current_live_session:
                    session_channel = self._session_channels.pop(layer_identifier, None)
                    if session_channel:
                        session_channel.destroy()
                elif current_live_session and not self._session_channels.get(layer_identifier, None):
                    session_channel = LiveSessionChannelManager(current_live_session, self)
                    session_channel.start_async()
                    self._session_channels[layer_identifier] = session_channel
        elif payload.event_type == LayerEventType.OUTDATE_STATE_CHANGED:
            if not self._all_outdated_layers:
                # Delays the notification
                self._last_outdate_notification_time = time.time()
            outdated_layer_identifiers = payload.identifiers_or_spec_paths
            for layer_identifier in outdated_layer_identifiers:
                if self._layers_state.is_layer_outdated(layer_identifier):
                    self._all_outdated_layers.add(layer_identifier)
                else:
                    self._all_outdated_layers.discard(layer_identifier)
                    # Dismiss the notification if all layers are not outdated anymore.
                    if not self._all_outdated_layers and self._outdate_notification:
                        self._outdate_notification.dismiss()
                        self._outdate_notification = None

    def destroy(self):
        self._stage_event_sub = None
        self._layers_event_sub = None
        self._layers_instance = None
        self._outdate_notification = None
        self._update_subscription = None
        self._layers_state = None
        self._silent_leave_layers.clear()
        for session_channel in self._session_channels.values():
            session_channel.destroy()
        self._session_channels.clear()
        release_live_syncing_interface(self._live_syncing_interface)
    
    def _to_live_session(self, session_handle, base_layer_identifier):
        session_channel = self._session_channels.get(base_layer_identifier, None)
        live_session = LiveSession(session_handle, session_channel, self)

        return live_session
    
    def get_all_live_sessions(self, layer_identifier: str = None) -> List[LiveSession]:
        """
        Gets all existing live sessions for a specific layer.
        
        Args:
            layer_identifier (str): Layer identifier. If it's empty, it will be root layer by default.
        
        Returns:
            A list of live sessions.
        """
        
        if not self.usd_context.get_stage():
            return []

        if not layer_identifier:
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        total_live_sessions = self._live_syncing_interface.get_total_live_sessions(self._layers_instance, layer_identifier)
        all_live_sessions = []
        for i in range(total_live_sessions):
            session_handle = self._live_syncing_interface.get_live_session_at_index(self._layers_instance, layer_identifier, i)
            all_live_sessions.append(self._to_live_session(session_handle, layer_identifier))
        
        return all_live_sessions
    
    def create_live_session(self, name: str, layer_identifier: str = None):
        """
        Creates a named live session.

        Args:
            name (str): Name of the session. Currently, name should be unique across all live sessions.

            layer_identifier (str): The sublayer identifier to create live session. If it's not provided,
                                    it will be root layer by default.
        """
        
        if not self.usd_context.get_stage():
            return None

        # It's to be compatible with old interface. The second param
        # of old interface is to enable auto_authoring, which is deprecated for now. 
        if not layer_identifier or isinstance(layer_identifier, bool):
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        session_handle = self._live_syncing_interface.create_live_session(self._layers_instance, layer_identifier, name)
        if not session_handle:
            return None
            
        live_session = self._to_live_session(session_handle, layer_identifier)

        return live_session
    
    def join_live_session(self, live_session: LiveSession):
        """Joins the live session."""

        self._joining_session = True
        session_handle = live_session._session_handle
        success = self._live_syncing_interface.join_live_session(self._layers_instance, session_handle)
        self._joining_session = False

        if success:
            session_channel = LiveSessionChannelManager(live_session, self)
            session_channel.start_async()
            self._session_channels[live_session.base_layer_identifier] = session_channel

        return success
    
    def join_live_session_by_url(self, layer_identifier: str, live_session_url: str, create_if_not_existed: bool=False):
        """
        Joins the live session specified by the live session URL.

        Args:
            layer_identifier (str): The base layer of the live session.

            live_session_url (str): The unique URL of the live session. The live session must be created against
                                    with the base layer, otherwise, it will fail to join.
            
            create_if_not_existed (bool): Creates the live session if it's not existed.
        """

        self._joining_session = True
        success = self._live_syncing_interface.join_live_session_by_url(self._layers_instance, layer_identifier, live_session_url, create_if_not_existed)
        self._joining_session = False
        
        if success:
            live_session = self.get_current_live_session(layer_identifier)
            session_channel = LiveSessionChannelManager(live_session, self)
            session_channel.start_async()
            self._session_channels[live_session.base_layer_identifier] = session_channel

        return success
    
    def stop_live_session(self, layer_identifier: str = None):
        """Stops the live session for specific layer."""
        
        if not self.usd_context.get_stage():
            return

        if not layer_identifier or isinstance(layer_identifier, bool):
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        self._joining_session = True
        self._live_syncing_interface.stop_live_session(self._layers_instance, layer_identifier)
        self._joining_session = False

        session_channel = self._session_channels.pop(layer_identifier, None)
        if session_channel:
            session_channel.destroy()
        
    def stop_all_live_sessions(self):
        """Stops all the live sessions that are enabled for the current stage."""

        self._joining_session = True
        self._live_syncing_interface.stop_all_live_sessions(self._layers_instance)
        self._joining_session = False

        for session_channel in self._session_channels.values():
            session_channel.destroy()
        self._session_channels.clear()
    
    def is_stage_in_live_session(self):
        return self._live_syncing_interface.is_stage_in_live_session(self._layers_instance)
    
    def is_layer_in_live_session(self, layer_identifier):
        return self._live_syncing_interface.is_layer_in_live_session(self._layers_instance, layer_identifier)
    
    def is_in_live_session(self):
        """[DEPRECATED] If root layer is in a live session."""
        
        if not self.usd_context.get_stage():
            return False
            
        layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        return self.is_layer_in_live_session(layer_identifier)
    
    def is_live_session_layer(self, layer_identifier: str):
        return self._live_syncing_interface.is_live_session_layer(self._layers_instance, layer_identifier)
    
    def get_current_live_session_layers(self, layer_identifier: str = None):
        """
        [DEPRECATED] Returns the live session layers attached to this live session of specifc base layer.
        """

        current_live_session = self.get_current_live_session(layer_identifier)
        if current_live_session:
            return [current_live_session.root]

        return []
    
    def get_current_live_session_peer_users(self, layer_identifier: str = None) -> List[cm.PeerUser]:
        """
        [DEPRECATED] Returns the live session layers attached to this live session of specifc base layer.

        Args:
            layer_identifier (str): It's root layer if it's not provided.
        """

        current_live_session = self.get_current_live_session(layer_identifier)
        if current_live_session:
            return current_live_session.peer_users
        
        return []

    def get_live_session_for_live_layer(self, live_layer_identifier: str):
        """
        Gets the live session that the live layer is attached to.

        Args:
            live_layer_identifier (str): The .live layer that's in the live session.
        """

        session_handle = self._live_syncing_interface.get_live_session_for_live_layer(
            self._layers_instance, live_layer_identifier
        )
        if not session_handle:
            return None
        
        base_layer_identifier = self._live_syncing_interface.get_live_session_base_layer_identifier(
            self._layers_instance, session_handle
        )
        return self._to_live_session(session_handle, base_layer_identifier)
    
    def get_current_live_session(self, layer_identifier: str = None) -> LiveSession:
        """
        Gets the current live session of the base layer joined.

        Args:
            layer_identifier (str): Base layer identifier. It's root layer if it's not provided.
        """ 

        if not self.usd_context.get_stage():
            return None

        if not layer_identifier:
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        session_handle = self._live_syncing_interface.get_current_live_session(self._layers_instance, layer_identifier)
        if not session_handle:
            return None
        
        return self._to_live_session(session_handle, layer_identifier)
    
    def get_all_current_live_sessions(self) -> List[LiveSession]:
        """Gets all live sessions that are joined for the current stage."""

        layer_identifiers = LayerUtils.get_all_sublayers(self.usd_context.get_stage(), True, True, False)
        all_joined_sessions = []
        for layer_identifier in layer_identifiers:
            if layer_identifier.endswith(".live"):
                continue

            current_session = self.get_current_live_session(layer_identifier)
            if current_session:
                all_joined_sessions.append(current_session)
        
        return all_joined_sessions
    
    def get_live_session_by_url(self, session_url) -> LiveSession:
        """
        Gets the live session by the url. It can only get the live session belongs to one of the sublayer
        in the current stage.
        """

        session_handle = self._live_syncing_interface.get_live_session_by_url(self._layers_instance, session_url)
        base_layer_identifier = self._live_syncing_interface.get_live_session_base_layer_identifier(
            self._layers_instance, session_handle
        )
        
        return self._to_live_session(session_handle, base_layer_identifier)
    
    @property
    def permission_to_merge_current_session(self, layer_identifier: str = None):
        """
        [DEPRECATED] Checkes to see if it can merge the current live session of specified layer.

        Args:
            layer_identifier (str): It's root layer if it's not provided.
        """

        current_session = self.get_current_live_session(layer_identifier)
        if not current_session:
            return False
        
        return current_session.merge_permission

    def merge_live_session_changes(self, layer_identifier: str, stop_session: bool):
        return self._live_syncing_interface.merge_live_session_changes(
            self._layers_instance, layer_identifier, stop_session
        )

    def merge_live_session_changes_to_specific_layer(
        self, layer_identifier: str,
        target_layer_identifier: str,
        stop_session: bool, clear_target_layer: bool
    ):
        return self._live_syncing_interface.merge_live_session_changes_to_specific_layer(
            self._layers_instance,
            layer_identifier,
            target_layer_identifier,
            stop_session,
            clear_target_layer
        )
    
    def merge_changes_to_base_layers(self, stop_session: bool):
        """
        [DEPRECATED] Merges changes of root layer's live session to its base layer.
        """
        
        if not self.usd_context.get_stage():
            return False
       
        layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier
        return self.merge_live_session_changes(layer_identifier, stop_session)

    def merge_changes_to_specific_layer(self, target_layer_identifier: str, stop_session: bool, clear_target_layer: bool):
        """
        [DEPRECATED] Merges changes of root layer's live session to specified layer.
        """
        
        if not self.usd_context.get_stage():
            return False
        
        base_layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier
        return self.merge_live_session_changes_to_specific_layer(
            base_layer_identifier, target_layer_identifier, stop_session, clear_target_layer
        )

    async def broadcast_merge_started_message_async(self, layer_identifier: str = None):
        """
        Broadcasts merge started message to other peer clients. This is normally
        called before merging the live session to its base layer.
        """
        
        if not self.usd_context.get_stage():
            return

        if not layer_identifier:
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        session_channel = self._session_channels.get(layer_identifier, None)
        if session_channel:
            await session_channel.broadcast_merge_started_message_async()

    async def broadcast_merge_done_message_async(self, destroy=True, layer_identifier: str = None):
        """
        Broadcasts merge finished message to other peer clients. This is normally
        called after merging the live session to its base layer.

        "Args:
            destroy (bool): If it needs to destroy the session channel after merging.

            layer_identifier (str): The base layer of the live session.
        """

        if not layer_identifier:
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        session_channel = self._session_channels.get(layer_identifier, None)
        if session_channel:
            await session_channel.broadcast_merge_done_message_async()
            if destroy:
                session_channel.destroy()
                self._session_channels.pop(layer_identifier, None)

    async def merge_and_stop_live_session_async(
        self, target_layer: str = None, comment="",
        pre_merge: Callable[[LiveSession], Awaitable] = None,
        post_merge: Callable[[bool], Awaitable] = None,
        layer_identifier: str = None
    ):
        """Saves chagnes of live session to their base layer if it's in live mode.
        
        Args:
            target_layer (str): Target layer to merge all live changes to.

            comment (str): The checkpoint comments.

            pre_merge (Callable[[LiveSession], Awaitable]): It will be called before merge.
            
            post_merge (Callable[[bool, str], Awaitable]): It will be called after merge is done. The first
                                                           param means if it's successful, and the second includes
                                                           the error message if it's not.
            
            layer_identifier (str): The base layer of the live session. It's root layer if it's not provided.
        """
        
        if not self.usd_context.get_stage():
            return False
        
        if not layer_identifier:
            layer_identifier = self.usd_context.get_stage().GetRootLayer().identifier

        current_session = self.get_current_live_session(layer_identifier)
        if pre_merge:
            await pre_merge(current_session)

        error = None
        success = True
        try:
            import omni.usd_resolver
            
            if not self.is_layer_in_live_session(layer_identifier):
                error = "Cannot merge live layers since it's not in live session."
                success = False
                return False
            
            if not current_session.merge_permission:
                error = f"Cannot merge live layers since only owner {current_session.owner} can merge them."
                success = False
                return False
            
            if self._merging_live_session:
                error = "Current session is already in progress of merging."
                success = False
                return False
            
            self._merging_live_session = True

            def report_checkpoint_errors(failed_layers):
                nonlocal error
                nonlocal success
                success = False
                if len(failed_layers) == 1:
                    error = f"Failed to create checkpoints for {failed_layers[0]}."
                elif len(failed_layers) > 1:
                    error = "Failed to create checkpoints for several layers. Please check console for more details."
           
            all_live_layers = [current_session.root]
            _, failed_layers = await LayerUtils.create_checkpoint_async(all_live_layers, comment, True)
            if failed_layers:
                report_checkpoint_errors(failed_layers)
                return False
            
            comment = f"Merge changes from live {current_session.name}: {comment}"
            omni.usd_resolver.set_checkpoint_message(comment)
            # Clear checkpoint message to ensure comment is not used in future file operations.

            if not target_layer:
                success = self.merge_live_session_changes(layer_identifier, True)
                if success:
                    dirty_layers = LayerUtils.get_dirty_layers(self.usd_context.get_stage(), True, True)
                    self._save_layers(dirty_layers)
            else:
                layer = Sdf.Find(target_layer)
                success = self.merge_live_session_changes_to_specific_layer(layer_identifier, target_layer, True, True)
                # It means it's existing layer, create a checkpoint firstly then save.
                if layer and success:
                    self._save_layers([target_layer])
            
            if not success:
                error = "Failed to merge live changes due to permission issues."
        finally:
            omni.usd_resolver.set_checkpoint_message("")
            
            if not success:
                carb.log_error(error)
                post_notification(error, False)
            
            if post_merge:
                await post_merge(success)
            
            if not self.is_layer_in_live_session(layer_identifier):
                session_channel = self._session_channels.pop(layer_identifier, None)
                if session_channel:
                    session_channel.destroy()
            
            self._merging_live_session = False

            return success
    
    def mute_live_session_merge_notice(self, layer_identifier):
        """
        By default, when session owner merges the session, it will
        notify peer clients to leave the session by showing a prompt.
        This is used to disable the prompt and stop session directly
        so it will not be disruptive especially for applications that
        don't have layer window and manages all sublayer live sessions
        with scripting and don't want to receive prompt for sublayer sessions.
        """

        self._silent_leave_layers.add(layer_identifier)
    
    def unmute_live_session_merge_notice(self, layer_identifier):
        self._silent_leave_layers.discard(layer_identifier)
    
    def is_live_session_merge_notice_muted(self, layer_identifier):
        return layer_identifier in self._silent_leave_layers
    
    def _save_layers(self, layer_identifiers):
        success = True
        for layer_identifier in layer_identifiers:
            layer = Sdf.Find(layer_identifier)
            if not layer:
                carb.log_error(f"Failed to save layer {layer_identifier} as it cannot be found.")
                continue

            if not layer.Save():
                carb.log_error(f"Failed to save layer {layer_identifier} because of permission issue.")
                success = False
        
        return success

 
