import asyncio
import carb
import omni.kit.collaboration.channel_manager as cm

from typing import Dict, List
from .interface_utils import post_notification
from .._omni_kit_usd_layers import LayerEventType


SESSION_MANAGEMENT_VERSION_KEY = "version"
SESSION_MANAGEMENT_VERSION = "1.0"
MESSAGE_GROUP_KEY = "__SESSION_MANAGEMENT__"
MESSAGE_KEY = "message"
MESSAGE_MERGE_STARTED = "MERGE_STARTED"
MESSAGE_MERGE_FINISHED = "MERGE_FINISHED"
MESSAGE_USER_NAME_KEY = "user_name"
MESSAGE_USER_ID_KEY = "user_id"
MESSAGE_SESSION_URL_KEY = "session_url"


class LiveSessionChannelManager:
    def __init__(self, current_session, live_syncing_interface):
        self._base_layer_identifier: str = current_session.base_layer_identifier
        self._session_url = current_session.url
        self._session_name: str = current_session.name
        self._channel_url = current_session.channel_url
        self._channel: cm.Channel = None
        self._channel_subscriber: cm.ChannelSubscriber = None
        self._join_channel_future = None
        self._live_syncing = live_syncing_interface
        self._event_stream = live_syncing_interface._layers_instance.get_event_stream()
    
    def destroy(self):
        self._stop_channel()
        self._live_syncing = None
        self._event_stream = None

    def start_async(self):
        self._join_channel_future = asyncio.ensure_future(self._join_channel_async(self._channel_url))

        return self._join_channel_future

    def stop(self):
        self._stop_channel()
    
    @property
    def stopped(self):
        return self._channel and self._channel.stopped
    
    @property
    def logged_user_name(self):
        return self._channel and self._channel.logged_user_name
    
    @property
    def logged_user_id(self):
        return self._channel and self._channel.logged_user_id
    
    def _send_layer_event(self, event_type: LayerEventType, payload={}):
        payload[MESSAGE_SESSION_URL_KEY] = self._session_url
        if payload:
            self._event_stream.push(int(event_type), 0, payload)
        else:
            self._event_stream.push(int(event_type))
    
    @property
    def peer_users(self) -> List[cm.PeerUser]:
        if self._channel:
            return list(self._channel.peer_users.values())
        
        return []
    
    async def broadcast_merge_started_message_async(self):
        if self._channel:
            message = {
                MESSAGE_GROUP_KEY: {
                    SESSION_MANAGEMENT_VERSION_KEY: SESSION_MANAGEMENT_VERSION,
                    MESSAGE_KEY: MESSAGE_MERGE_STARTED
                }
            }
            await self._channel.send_message_async(message)

    async def broadcast_merge_done_message_async(self):
        if self._channel:
            message = {
                MESSAGE_GROUP_KEY: {
                    SESSION_MANAGEMENT_VERSION_KEY: SESSION_MANAGEMENT_VERSION,
                    MESSAGE_KEY: MESSAGE_MERGE_FINISHED
                }
            }
            await self._channel.send_message_async(message)

    async def _join_channel_async(self, url):
        self._channel = await cm.join_channel_async(url)
        if not self._channel:
            post_notification(f"Failed to join channel {url}. Stopping session '{self._session_name}'...")
            self._live_syncing.stop_live_session(self._base_layer_identifier)

            return False

        self._channel_subscriber = self._channel.add_subscriber(self._on_channel_message)

        return True
    
    def _on_channel_message(self, message: cm.Message):
        info = None
        from_user = message.from_user
        if message.message_type == cm.MessageType.JOIN:
            info = f"User {message.from_user.user_name} has joined the session '{self._session_name}'."
            self._send_layer_event(
                LayerEventType.LIVE_SESSION_USER_JOINED,
                {
                    MESSAGE_USER_NAME_KEY: from_user.user_name,
                    MESSAGE_USER_ID_KEY: from_user.user_id,
                }
            )
        elif message.message_type == cm.MessageType.LEFT:
            info = f"User {message.from_user.user_name} has left the session '{self._session_name}."
            self._send_layer_event(
                LayerEventType.LIVE_SESSION_USER_LEFT,
                {
                    MESSAGE_USER_NAME_KEY: from_user.user_name,
                    MESSAGE_USER_ID_KEY: from_user.user_id,
                }
            )
        elif message.message_type == cm.MessageType.HELLO:
            info = f"User {message.from_user.user_name} has joined the session '{self._session_name}'."
            self._send_layer_event(
                LayerEventType.LIVE_SESSION_USER_JOINED,
                {
                    MESSAGE_USER_NAME_KEY: from_user.user_name,
                    MESSAGE_USER_ID_KEY: from_user.user_id,
                }
            )
        elif message.message_type == cm.MessageType.MESSAGE:
            content = message.content.get(MESSAGE_GROUP_KEY, None) or message.content.get(MESSAGE_GROUP_KEY.lower(), None)
            if not content or not isinstance(content, dict):
                return

            message_type = content.get(MESSAGE_KEY, None) or content.get(MESSAGE_KEY.lower(), None)
            if not message_type:
                return
            
            current_session = self._live_syncing.get_current_live_session(self._base_layer_identifier)
            if not current_session:
                carb.log_error("Invalid state. Live Session is not started but channel is still alive.")
                return

            if message_type.lower() == MESSAGE_MERGE_STARTED.lower():
                if self._live_syncing.is_live_session_merge_notice_muted(self._base_layer_identifier):
                    self._live_syncing.stop_live_session(self._base_layer_identifier)
                else:
                    try:
                        from omni.kit.widget.prompt import PromptButtonInfo, PromptManager
                        PromptManager.post_simple_prompt(
                            "Leave Session",
                            f"`{current_session.owner}' is ending the live session. "
                            "Stop all work and leave the session.",
                            PromptButtonInfo(
                                "LEAVE", lambda: self._live_syncing.stop_live_session(self._base_layer_identifier)
                            ),
                            shortcut_keys=False
                        )
                    except Exception:
                        carb.log_warn("Cannot post prompt as omni.ui is not enabled, stopping live session forcely.")
                        self._live_syncing.stop_live_session(self._base_layer_identifier)
            else:
                return

        # Don't explicitly make layers backend dependent on notification manager
        if info:
            post_notification(info)
            carb.log_info(info)
    
    def _stop_channel(self):
        if self._channel_subscriber:
            self._channel_subscriber.unsubscribe()
            self._channel_subscriber = None
        if self._channel:
            self._channel.stop()
        self._channel = None
        try:
            if self._join_channel_future and not self._join_channel_future.done():
                self._join_channel_future.cancel()
        except Exception:
            pass
