# Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
import carb.settings
import omni.ext
import omni.kit.app
import omni.client
import asyncio

from typing import Callable
from functools import partial
from .connector import NucleusConnector
from .device_auth import DeviceAuthConnector
from . import NUCLEUS_CONNECTION_SUCCEEDED_EVENT

g_singleton = None


# Any class derived from `omni.ext.IExt` in top level module (defined in `python.modules` of `extension.toml`) will be
# instantiated when extension gets enabled and `on_startup(ext_id)` will be called. Later when extension gets disabled
# on_shutdown() is called.
class NucleusConnectorExtension(omni.ext.IExt):
    """Helper extension for connecting to Nucleus servers"""

    def __init__(self):
        super().__init__()
        self._connector = None

    def on_startup(self, ext_id):
        # Save away this instance as singleton
        global g_singleton
        g_singleton = self

        # OM-98449: Register device auth flow callback if device auth flow is turned on in settings
        self._device_auth_sub = None
        settings = carb.settings.get_settings()
        device_auth = settings.get_as_bool("exts/omni.kit.widget.nucleus_connector/device_auth_flow")
        if device_auth:  # pragma: no cover
            self._connector = DeviceAuthConnector()
            self._device_auth_sub = omni.client.register_device_flow_auth_callback(self._on_device_auth)
        else:
            self._connector = NucleusConnector()
            # Set the callback to invoke when authentication requires opening a web browser to complete sign-in.
            omni.client.set_authentication_message_box_callback(self._on_authenticate)

    def _on_device_auth(self, auth_handle: int, params: omni.client.AuthDeviceFlowParams):
        """
        This callback is invoked when authentication for device auth flow, instead of auth in the web browser.

        Args:
            auth_handle (int): The auth handle that could be used for canceling.
            params (omni.client.AuthDeviceFlowParams): An object containing info to display to user to complete the auth.
                if it is None, it means the auth attempt is finished.
        """
        if self._connector:
            if not params:
                # Auth is done, either succeeded or failed, result needed to be determined from status callback;
                # Here we just close out the dialog.
                self._connector.end_auth_flow(auth_handle)
                return

            self._connector.start_auth_flow(auth_handle, params)

    def _on_authenticate(self, show_alert: bool, host_name: str, auth_handle: int):
        """
        This callback is invoked when the application should show a dialog box letting the user know that a
        browser window has opened and to complete signing in using the web browser.

        Args:
            show_alert (bool): True means to start the authentication cycle and False means to end it
            host_name (str): The server host name that the authentication is for
            auth_handle (int): An integer handle that should be passed to cancel authentication

        """
        if self._connector:
            if show_alert:
                self._connector.start_auth_flow(host_name, auth_handle)
            else:
                self._connector.end_auth_flow(auth_handle)

    def _on_connect_succeeded(self, callback: Callable, name: str, url: str):
        """Called when connection is made successfully; invokes the given callback"""
        # Emit event on successful connection
        event_stream = omni.kit.app.get_app().get_message_bus_event_stream()
        if event_stream:
            event_stream.push(NUCLEUS_CONNECTION_SUCCEEDED_EVENT, payload={"url": url})
        if callback:
            callback(name, url)

    def _on_connect_failed(self, callback: Callable, name: str, url: str):
        """Called when connection fails; invokes the given callback"""
        if callback:
            callback(name, url)

    def connect_with_dialog(self, on_success_fn: Callable = None, on_failed_fn: Callable = None):
        """
        Prompts for server name and Url and proceeds to connect to it.

        Args:
            on_success_fn (Callable): Invoked when successful, on_success_fn(name: str, url: str)
            on_faild_fn (Callable): Invoked when failed, on_faild_fn(name: str, url: str)

        """
        if self._connector:
            self._connector.connect_with_callback(
                on_success_fn=partial(self._on_connect_succeeded, on_success_fn), on_failed_fn=partial(self._on_connect_failed, on_failed_fn))

    def connect(self, name: str, url: str, on_success_fn: Callable = None, on_failed_fn: Callable = None):
        """
        Connects to the named server.

        Args:
            name (str): Name of server
            url (str): Url of server
            on_success_fn (Callable): Invoked when successful, on_success_fn(name: str, url: str)
            on_faild_fn (Callable): Invoked when failed, on_faild_fn(name: str, url: str)

        """
        if self._connector:
            self._connector.connect_with_callback(name=name, url=url,
                on_success_fn=partial(self._on_connect_succeeded, on_success_fn), on_failed_fn=partial(self._on_connect_failed, on_failed_fn))

    def reconnect(self, url: str, on_success_fn: Callable = None, on_failed_fn: Callable = None):
        """
        Reconnects to the named server.

        Args:
            url (str): Url of server
            on_success_fn (Callable): Invoked when successful, on_success_fn(url: str)
            on_faild_fn (Callable): Invoked when failed, on_faild_fn(url: str)

        """
        if self._connector:
            asyncio.ensure_future(self._connector.reconnect_server_async(
                url, partial(self._on_connect_succeeded, on_success_fn), partial(self._on_connect_failed, on_failed_fn)))

    def disconnect(self, url: str):
        """
        Disconnects from the server.

        Args:
            url (str): Url of server

        """
        omni.client.sign_out(url)

    def on_shutdown(self):
        # Clears the auth callback
        omni.client.set_authentication_message_box_callback(None)

        # OM-98449: Clears device auth flow connector and callback
        self._device_auth_sub = None

        # Clean up the connector
        if self._connector:
            self._connector.destroy()
            self._connector = None

        global g_singleton
        g_singleton = None


def get_instance():
    return g_singleton
