# Copyright (c) 2022, 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.

"""Native streaming extension."""

import carb

import omni.ext
import omni.kit.livestream.bind


class NativeStreamingExtension(omni.ext.IExt):
    """Native streaming extension."""

    def __init__(self) -> None:
        super().__init__()
        self._stream_interface = None

    def on_startup(self) -> None:
        self._kit_livestream = omni.kit.livestream.bind.acquire_livestream_interface()
        self._kit_livestream.startup()
        self._register_streaming_interface()

    def on_shutdown(self) -> None:
        self._unregister_streaming_interface()
        self._kit_livestream.shutdown()
        self._kit_livestream = None

    def _register_streaming_interface(self) -> None:
        """Register the streaming interface for the extension."""
        try:
            from omni.services.streaming.manager import get_stream_manager, StreamManager
            from .streaming_interface import NativeStreamInterface

            self._stream_interface = NativeStreamInterface()
            stream_manager: StreamManager = get_stream_manager()
            stream_manager.register_stream_interface(stream_interface=self._stream_interface)
            stream_manager.enable_stream_interface(stream_interface_id=self._stream_interface.id)
        except ImportError:
            carb.log_info("The Streaming Manager extension was not available when enabling native streaming features.")
        except Exception as exc:
            carb.log_error(f"An error occurred while attempting to register the native streaming interface: {str(exc)}.")

    def _unregister_streaming_interface(self) -> None:
        """Unregister the streaming interface for the extension."""
        if self._stream_interface is None:
            # No stream interface was register when enabling the extension, so there is no need to attempt to unregister
            # it and the function can exit early:
            return

        try:
            from omni.services.streaming.manager import get_stream_manager, StreamManager

            stream_manager: StreamManager = get_stream_manager()
            stream_manager.disable_stream_interface(stream_interface_id=self._stream_interface.id)
            stream_manager.unregister_stream_interface(stream_interface_id=self._stream_interface.id)
        except ImportError:
            carb.log_info("The Streaming Manager extension was not available when disabling native streaming features.")
        except Exception as exc:
            carb.log_error(f"An error occurred while attempting to unregister the native streaming interface: {str(exc)}.")
