__all__ = ["stop_or_show_live_session_widget", "LiveSessionWidgetExtension"]

import asyncio
import os
import carb
import omni.ext
import omni.kit.app
import omni.usd
import omni.client
import omni.ui as ui
import omni.kit.notification_manager as nm
import omni.kit.usd.layers as layers
import omni.kit.clipboard

from typing import Union
from enum import Enum
from omni.kit.widget.prompt import PromptButtonInfo, PromptManager
from .layer_icons import LayerIcons
from .live_session_start_window import LiveSessionStartWindow
from .live_session_end_window import LiveSessionEndWindow
from .utils import is_extension_loaded, join_live_session, is_viewer_only_mode, is_quick_join_enabled
from .join_with_session_link_window import JoinWithSessionLinkWindow
from .share_session_link_window import ShareSessionLinkWindow
from pxr import Sdf
from .live_style import Styles


_extension_instance = None
_debugCollaborationWindow = None


class LiveSessionMenuOptions(Enum):
    QUIT_ONLY = 0
    MERGE_AND_QUIT = 1


class LiveSessionWidgetExtension(omni.ext.IExt):
    def on_startup(self, ext_id):
        global _extension_instance
        _extension_instance = self
        extension_path = omni.kit.app.get_app_interface().get_extension_manager().get_extension_path(ext_id)
        self._settings = None
        LayerIcons.on_startup(extension_path)
        self._live_session_end_window = None
        self._live_session_start_window = None
        self._join_with_session_link_window = None
        self._share_session_link_window = None
        self._live_session_menu = None
        self._app = omni.kit.app.get_app()
        try:
            from omni.kit.window.preferences import register_page
            from .live_session_preferences import LiveSessionPreferences

            self._settings = register_page(LiveSessionPreferences())
        except ImportError: # pragma: no cover
            pass
        Styles.on_startup()

    def on_shutdown(self):
        global _extension_instance
        _extension_instance = None
        self._live_session_menu = None
        if self._live_session_end_window:
            self._live_session_end_window.destroy()
            self._live_session_end_window = None

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

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

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

        if self._settings:
            try:
                from omni.kit.window.preferences import unregister_page

                unregister_page(self._settings)
            except ImportError:
                pass
            self._settings = None

    @staticmethod
    def get_instance():
        global _extension_instance
        return _extension_instance

    def _show_live_session_end_window(self, layers_interface: layers.Layers, layer_identifier):
        if self._live_session_end_window:
            self._live_session_end_window.destroy()

        async def show_dialog():
            self._live_session_end_window = LiveSessionEndWindow(layers_interface, layer_identifier)
            async with self._live_session_end_window:
                pass

        asyncio.ensure_future(show_dialog())

    def _show_join_with_session_link_window(self, layers_interface):
        current_session_link = None

        # Destroys it to center the window.
        if self._join_with_session_link_window:
            current_session_link = self._join_with_session_link_window.current_session_link
            self._join_with_session_link_window.destroy()

        self._join_with_session_link_window = JoinWithSessionLinkWindow(layers_interface)
        self._join_with_session_link_window.visible = True
        if current_session_link:
            self._join_with_session_link_window.current_session_link = current_session_link

    def _show_share_session_link_window(self, layers_interface, layer_identifier):
        # Destroys it to center the window.
        if self._share_session_link_window:
            self._share_session_link_window.destroy()

        self._share_session_link_window = ShareSessionLinkWindow(layers_interface, layer_identifier)
        self._share_session_link_window.visible = True

    def _show_live_session_menu_options(self, layers_interface, layer_identifier, is_stage_session, prim_path=None):
        live_syncing = layers_interface.get_live_syncing()
        current_session = live_syncing.get_current_live_session(layer_identifier)
        self._live_session_menu = ui.Menu("global_live_update")
        is_omniverse_stage = layer_identifier.startswith("omniverse://")
        is_live_prim = prim_path or live_syncing.is_layer_in_live_session(layer_identifier, live_prim_only=True)

        with self._live_session_menu:
            if is_omniverse_stage:
                if current_session:
                    # If layer is in the live session already, showing `leave` and `end and merge`
                    ui.MenuItem(
                        "Leave Session",
                        triggered_fn=lambda *args: self._on_leave_live_session_menu_options(
                            layers_interface, LiveSessionMenuOptions.QUIT_ONLY, layer_identifier
                        ),
                    )

                    if not is_viewer_only_mode():
                        ui.MenuItem(
                            "End and Merge",
                            triggered_fn=lambda *args: self._on_leave_live_session_menu_options(
                                layers_interface, LiveSessionMenuOptions.MERGE_AND_QUIT, layer_identifier
                            ),
                        )
                    ui.Separator()
                else:
                    # If it's not in the live session, show `join` and `create`.
                    ui.MenuItem(
                        "Join Session",
                        triggered_fn=lambda *args: self._show_live_session_start_window(
                            layers_interface, layer_identifier, True, prim_path
                        ),
                    )
                    ui.MenuItem(
                        "Create Session",
                        triggered_fn=lambda *args: self._show_live_session_start_window(
                            layers_interface, layer_identifier, False, prim_path
                        ),
                    )

                # Don't show those options for live prim
                if not is_live_prim:
                    # If layer_identifier is given, it's to show menu options for sublayer.
                    # Show `copy` to copy session link. Otherwise, show share session dialog
                    # to choose session to share.
                    if is_stage_session:
                        ui.Separator()
                        ui.MenuItem(
                            "Share Session Link",
                            triggered_fn=lambda *args: self._show_share_session_link_window(
                                layers_interface, layer_identifier
                            ),
                        )
                    elif current_session:
                        ui.Separator()
                        ui.MenuItem(
                            "Copy Session Link",
                            triggered_fn=lambda *args: self._copy_session_link(layers_interface),
                        )

            # If layer identifier is gven, it will not show `join with session link` menu option.
            if is_stage_session and not is_live_prim:
                ui.MenuItem(
                    "Join With Session Link",
                    triggered_fn=lambda *args: self._show_join_with_session_link_window(layers_interface),
                )

            if is_extension_loaded("omni.kit.collaboration.debug_options"): # pragma: no cover
                ui.Separator()
                ui.MenuItem(
                    "Open Debug Window",
                    triggered_fn=lambda *args: self._on_open_debug_window(),
                )

            if is_omniverse_stage and current_session and is_extension_loaded("omni.timeline.live_session"):
                ui.Separator()
                import omni.timeline.live_session
                timeline_session = omni.timeline.live_session.get_timeline_session()
                if timeline_session is not None:
                    session_window = omni.timeline.live_session.get_session_window()
                    ui.MenuItem(
                        "Synchronize Timeline",
                        triggered_fn=lambda *args: self._on_timeline_sync_triggered(
                            layers_interface, layer_identifier, timeline_session
                        ),
                        checkable=True,
                        checked=timeline_session.is_sync_enabled()
                    )
                    if timeline_session.am_i_owner():
                        ui.MenuItem(
                            "Set Timeline Owner",
                            triggered_fn=lambda *args: self._on_open_timeline_session_window(
                                layers_interface, layer_identifier, timeline_session, session_window
                            ),
                        )
                    elif not timeline_session.am_i_presenter():
                        if not self._requested_timeline_control(timeline_session):
                            ui.MenuItem(
                                "Request Timeline Ownership",
                                triggered_fn=lambda *args: self._on_request_timeline_control(timeline_session),
                            )
                        else:
                            ui.MenuItem(
                                "Withdraw Timeline Ownership Request",
                                triggered_fn=lambda *args: self._on_withdraw_timeline_control(timeline_session),
                            )

        self._live_session_menu.show()

    def _show_live_session_start_window(self, layers_interface, layer_identifier, join_session=None, prim_path=None):
        if self._live_session_start_window:
            self._live_session_start_window.destroy()

        self._live_session_start_window = LiveSessionStartWindow(layers_interface, layer_identifier, prim_path)
        self._live_session_start_window.visible = True
        if join_session is not None:
            self._live_session_start_window.set_focus(join_session)

    def _show_session_start_window_or_menu(
        self, layers_interface: layers.Layers, show_join_options, layer_identifier, is_stage_session, join_session=None, prim_path=None
    ):
        if show_join_options:
            return self._show_live_session_menu_options(
                layers_interface, layer_identifier, is_stage_session, prim_path
            )
        else:
            self._show_live_session_start_window(layers_interface, layer_identifier, join_session, prim_path)

        return None

    def _requested_timeline_control(self, timeline_session) -> bool:
        return timeline_session is not None and\
            timeline_session.live_session_user in timeline_session.get_request_controls()

    def _on_request_timeline_control(self, timeline_session):
        if timeline_session is not None:
            timeline_session.request_control(True)

    def _on_withdraw_timeline_control(self, timeline_session):
        if timeline_session is not None:
            timeline_session.request_control(False)

    def _on_timeline_sync_triggered(
        self, layers_interface: layers.Layers,
        layer_identifier: str,
        timeline_session
    ):
        timeline_session.enable_sync(not timeline_session.is_sync_enabled())

    def _on_open_timeline_session_window(
        self, layers_interface: layers.Layers,
        layer_identifier: str,
        timeline_session,
        session_window,
    ):
        session_window.show()

    def _on_open_debug_window(
        self,
    ): # pragma: no cover
        import omni.kit.collaboration.debug_options
        from omni.kit.collaboration.debug_options import CollaborationWindow
        global _debugCollaborationWindow
        if _debugCollaborationWindow is None :
            _debugCollaborationWindow = CollaborationWindow(width=480, height=320)

        _debugCollaborationWindow.visible = True

    def _on_leave_live_session_menu_options(
        self, layers_interface: layers.Layers,
        options: LiveSessionMenuOptions,
        layer_identifier: str,
    ):
        live_syncing = layers_interface.get_live_syncing()
        current_session = live_syncing.get_current_live_session(layer_identifier)
        if current_session:
            if options == LiveSessionMenuOptions.QUIT_ONLY:
                if self._can_quit_session_directly(live_syncing, layer_identifier):
                    live_syncing.stop_live_session(layer_identifier)
                else:
                    PromptManager.post_simple_prompt(
                        "Leave Session",
                        f"You are about to leave '{current_session.name}' session.",
                        PromptButtonInfo("LEAVE", lambda: live_syncing.stop_live_session(layer_identifier)),
                        PromptButtonInfo("CANCEL")
                    )
            elif options == LiveSessionMenuOptions.MERGE_AND_QUIT:
                layers_state = layers_interface.get_layers_state()
                is_read_only_on_disk = layers_state.is_layer_readonly_on_disk(layer_identifier)
                if current_session.merge_permission and not is_read_only_on_disk:
                    self._show_live_session_end_window(layers_interface, layer_identifier)
                else:
                    if is_read_only_on_disk:
                        owner = layers_state.get_layer_owner(layer_identifier)
                    else:
                        owner = current_session.owner

                    logged_user_name = current_session.logged_user_name
                    if logged_user_name == owner and current_session.merge_permission:
                        message = f"You currently do not have merge privileges as base layer is read-only on disk."\
                                  " Do you want to quit this live session?"
                    else:
                        if is_read_only_on_disk:
                            message = f"You currently do not have merge privileges as base layer is read-only on disk."\
                                      f" Please contact '{owner}' to grant you write access. Do you want to quit this live session?"
                        else:
                            message = f"You currently do not have merge privileges, please contact '{owner}'"\
                                      " to merge any changes from the live session. Do you want to quit this live session?"

                    PromptManager.post_simple_prompt(
                        "Permission Denied",
                        message,
                        PromptButtonInfo("YES", lambda: live_syncing.stop_live_session(layer_identifier)),
                        PromptButtonInfo("NO")
                    )

    def _can_quit_session_directly(self, live_syncing: layers.LiveSyncing, layer_identifier):
        current_live_session = live_syncing.get_current_live_session(layer_identifier)
        if current_live_session:
            if current_live_session.peer_users:
                return False

            layer = Sdf.Find(current_live_session.root)
            if layer and len(layer.rootPrims) > 0:
                return False

        return True

    def _copy_session_link(self, layers_interface):
        live_syncing = layers_interface.get_live_syncing()
        current_session = live_syncing.get_current_live_session()
        if not current_session:
            return

        omni.kit.clipboard.copy(current_session.shared_link)

    def _stop_or_show_live_session_widget_internal(
        self, layers_interface: layers.Layers,
        stop_session_only,
        stop_session_forcely,
        show_join_options,
        layer_identifier,
        quick_join="",
        prim_path=None
    ):
        usd_context = layers_interface.usd_context
        stage = usd_context.get_stage()
        # Skips it if stage is not opened.
        if not stage:
            return None

        # By default, it will join live session of tht root layer.
        if not layer_identifier:
            # If layer identifier is not provided, it's to show global
            # menu that includes share sesison link and join session with link.
            # Otherwise, it will show menu for the sublayer session only.
            is_stage_session = True
            root_layer = stage.GetRootLayer()
            layer_identifier = root_layer.identifier
        else:
            is_stage_session = False

        layer_handle = Sdf.Find(layer_identifier)
        if (
            not layer_handle or
            (not prim_path and not stage.HasLocalLayer(layer_handle)) or
            (prim_path and layer_handle not in stage.GetUsedLayers())
        ):
            nm.post_notification(
                "Only a layer opened in the current stage can start a live-sync session.",
                status=nm.NotificationStatus.WARNING
            )
            return None

        if not layer_identifier.startswith("omniverse://"):
            if is_stage_session and not prim_path:
                return self._show_live_session_menu_options(layers_interface, layer_identifier, is_stage_session)
            else:
                nm.post_notification(
                    "Only a layer in Nucleus can start a live-sync session.",
                    status=nm.NotificationStatus.WARNING
                )

                return None

        live_syncing = layers_interface.get_live_syncing()
        if not live_syncing.is_layer_in_live_session(layer_identifier):
            if not quick_join and not show_join_options:
                quick_join = "Default"
            elif quick_join and show_join_options:
                show_join_options = False

            if quick_join:
                # This runs when the main live button is pressed.
                live_syncing = layers_interface.get_live_syncing()
                quick_session = live_syncing.find_live_session_by_name(layer_identifier, quick_join)
                if is_quick_join_enabled():
                    if not quick_session:
                        quick_session = live_syncing.create_live_session(
                            layer_identifier=layer_identifier,
                            name=quick_join
                        )
                        if not quick_session:
                            nm.post_notification(
                                "Cannot create a Live session on a read only location. Please\n"
                                "save root file in a writable location.",
                                status=nm.NotificationStatus.WARNING
                            )
                            return None

                    join_live_session(
                        layers_interface, layer_identifier, quick_session, prim_path,
                        is_viewer_only_mode()
                    )
                    return None
                else:
                    join_session = True if quick_session is not None else False
                    return self._show_session_start_window_or_menu(
                        layers_interface, False, layer_identifier, False, join_session, prim_path
                    )
            else:
                return self._show_session_start_window_or_menu(
                    layers_interface, show_join_options, layer_identifier, is_stage_session, prim_path
                )
        elif stop_session_forcely:
            live_syncing.stop_live_session(layer_identifier)
        elif stop_session_only:
            self._on_leave_live_session_menu_options(
                layers_interface, LiveSessionMenuOptions.QUIT_ONLY, layer_identifier
            )
        else:
            return self._show_live_session_menu_options(layers_interface, layer_identifier, is_stage_session, prim_path)

        return None


def stop_or_show_live_session_widget(
    usd_context: Union[str, omni.usd.UsdContext] = "",
    stop_session_only: bool = False,
    stop_session_forcely: bool = False,
    show_join_options: bool = True,
    layer_identifier: str = None,
    quick_join: str = False,
    prim_path: Union[str, Sdf.Path] = None,
):
    """Stops current live session or shows widget to join/leave session. If current session is empty, it will show
    the session start dialog. If current session is not empty, it will stop session directly or show stop session menu.

    Args:
        usd_context:            The current usd context to operate on.

        stop_session_only:      If it's true, it will ask user to stop session if session is not empty without showing merge options.
                                If it's false, it will show both leave and merge session options.

        stop_session_forcely:   If it's true, it will stop session forcely even if current sesison is not empty.
                                Otherwise, it will show leave session or both leave and merge options that's dependent on
                                if stop_session_only is true or false.

        show_join_options:      If it's true, it will show menu options.
                                If it's false, it will show session dialog directly.
                                This option is only used when layer is not in a live session. If both quick_join and
                                this param have the same value, it will do quick join.

        layer_identifier:       By default, it will join/stop the live session of the root layer.
                                If layer_identifier is not None, it will join/stop the live session for the specific sublayer.

        quick_join:             Quick Join / Leave a session without bringing up a dialog.
                                This option is only used when layer is not in a live session. If both show_join_options and
                                this param have the same value, this param will be True.

        prim_path:              If prim path is specified, it will join live session for the prim instead of sublayers. Only prim with
                                references or payloads supports to join a live session. Prim path is not needed when it's to
                                stop a live session.

    Returns:
        It will return the menu widgets created if it needs to create options menu. Otherwise, it will return None.
    """

    instance = LiveSessionWidgetExtension.get_instance()
    if not instance: # pragma: no cover
        carb.log_error("omni.kit.widget.live_session_management extension must be enabled.")
        return

    layers_interface = layers.get_layers(usd_context)
    return instance._stop_or_show_live_session_widget_internal(
        layers_interface,
        stop_session_only,
        stop_session_forcely,
        show_join_options,
        layer_identifier,
        quick_join=quick_join,
        prim_path=prim_path
    )
