# 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.
#

from omni.kit.widget.settings import SettingType
import omni.ui as ui
import carb
import omni.kit.app

from omni.rtx.window.settings.rtx_settings_stack import RTXSettingsStack
from omni.rtx.window.settings.settings_collection_frame import SettingsCollectionFrame
from omni.rtx.window.settings.rendersettingsdefaults import SettingsDefaultHandler

class ClusterCompositingSettingsFrame(SettingsCollectionFrame):
    def _build_ui(self):

        compositingModeItems = {
            "All": "COMPOSITING_ALL",
            "Local Only": "COMPOSITING_LOCAL_ONLY",
            "Remote Only": "COMPOSITING_REMOTE_ONLY",
            "None": "NO_COMPOSITING"
        }
        self._add_setting_combo(
            "Compositing Mode",
            "/rtx/index/compositingMode",
            compositingModeItems,
            tooltip="Control distribution of the compositing workload",
        )

        if not self._settings.get("/nvindex/showAdvancedSettings"):
            return

        self._add_setting(
            SettingType.BOOL,
            "Transfer Span Alpha Channel",
            "/rtx/index/spanAlphaChannel",
            tooltip="Enable transfer of the alpha channel for horizontal span image data",
        )

        self._add_setting(
            SettingType.BOOL,
            "Span Image Encoding",
            "/rtx/index/spanImageEncoding",
            tooltip="Use special image encoding to improve compression for horizontal span image data",
        )

        self._add_setting(
            SettingType.BOOL,
            "Tile Image Encoding",
            "/rtx/index/tileImageEncoding",
            tooltip="Use special image encoding to improve compression for intermediate tile image data",
        )

        self._add_setting(
            SettingType.INT,
            "Span Compression Level",
            "/rtx/index/spanCompressionLevel",
            0,
            9,
            tooltip="Compression level for transfer of horizontal span image data",
        )

        self._add_setting(
            SettingType.INT,
            "Tile Compression Level",
            "/rtx/index/tileCompressionLevel",
            0,
            9,
            tooltip="Compression level for transfer of intermediate tiles image data",
        )

        self._add_setting(
            SettingType.INT,
            "Span Compression Mode",
            "/rtx/index/spanCompressionMode",
            0,
            9,
            tooltip="Compression mode for transfer of horizontal span image data",
        )

        self._add_setting(
            SettingType.INT,
            "Tile Compression Mode",
            "/rtx/index/tileCompressionMode",
            0,
            9,
            tooltip="Compression mode for transfer of intermediate tiles image data",
        )

        self._add_setting(
            SettingType.INT,
            "Span Compression Threads",
            "/rtx/index/spanCompressionThreads",
            0,
            64,
            tooltip="Number of threads to use for compressing horizontal span image data (0=default)",
        )

        self._add_setting(
            SettingType.INT,
            "Tile Compression Threads",
            "/rtx/index/tileCompressionThreads",
            0,
            64,
            tooltip="Number of threads to use for compressing intermediate tiles image data (0=default)",
        )

        if self._settings.get("/nvindex/compositeRenderingMode"):
            self._add_setting(
                SettingType.INT,
                "Depth Buffer Compression Level",
                "/rtx/index/depthBufferCompressionLevel",
                0,
                9,
                tooltip="Compression level for transfer of depth buffer data",
            )

            self._add_setting(
                SettingType.INT,
                "Depth Buffer Compression Mode",
                "/rtx/index/depthBufferCompressionMode",
                0,
                9,
                tooltip="Compression mode for transfer of depth buffer data",
            )

            self._add_setting(
                SettingType.INT,
                "Depth Buffer Compression Threads",
                "/rtx/index/depthBufferCompressionThreads",
                0,
                64,
                tooltip="Number of threads to use for compressing depth buffer data (0=default)",
            )


class ClusterNetworkSettingsFrame(SettingsCollectionFrame):
    def _build_ui(self):

        network_enabled_path = "/nvindex/state/networkEnabled"
        cluster_message_path = "/nvindex/state/clusterMessage"
        self._change_cb_network_enabled = omni.kit.app.SettingChangeSubscription(network_enabled_path, self._on_change)
        self._change_cb_cluster_message = omni.kit.app.SettingChangeSubscription(cluster_message_path, self._on_change)

        restart_msg = "changing it requires restarting IndeX"

        networkModeItem = {
            "UDP": "udp",
            "TCP": "tcp",
            "OFF": "off",
        }
        self._add_setting_combo(
            "Network Mode",
            "/nvindex/networkMode",
            networkModeItem,
            tooltip="Mode for IndeX cluster mode, " + restart_msg,
        )

        self._add_setting(
            SettingType.STRING,
            "Discovery Address",
            "/nvindex/discoveryAddress",
            tooltip="Discovery address for cluster mode, " + restart_msg,
        )

        self._add_setting(
            SettingType.STRING,
            "Multicast Address",
            "/nvindex/multicastAddress",
            tooltip="Multicast address for cluster mode, " + restart_msg
        )

        self._add_setting(
            SettingType.STRING,
            "Cluster Interface Address",
            "/nvindex/clusterInterfaceAddress",
            tooltip="Network interface to use cluster mode, " + restart_msg
        )

        self._add_setting(
            SettingType.INT,
            "Minimum Cluster Size",
            "/nvindex/minimumClusterSize",
            tooltip="Minimum size of the cluster to start rendering"
        )

        serviceModeItems = {
            "Rendering + Compositing ": "rendering_and_compositing",
            "Rendering only": "rendering",
            "Compositing only": "compositing",
            "None": "none",
        }
        self._add_setting_combo(
            "Service Mode",
            "/nvindex/clusterServiceMode",
            serviceModeItems,
            tooltip="Service mode of the current host, use 'None' render remotely only, " + restart_msg,
        )

        self._add_setting(
            SettingType.STRING,
            "Cluster Identifier",
            "/nvindex/clusterIdentifier",
            tooltip="Optional identifier, must be the same on all hosts, " + restart_msg,
        )

        def do_restart_index():
            self._settings.set("/nvindex/restartIndex", True)

        ui.Button("Restart IndeX", tooltip="Restart the NVIDIA IndeX session, reloading all data", clicked_fn=do_restart_index)

        if self._settings.get(network_enabled_path):
            ui.Label("Cluster State:", alignment=ui.Alignment.CENTER)
            ui.Label(self._settings.get(cluster_message_path), word_wrap=True)

    def destroy(self):
        self._change_cb = None
        self._change_cb_network_enabled = None
        self._change_cb_cluster_message = None
        super().destroy()


class IndeXClusterSettingStack(RTXSettingsStack):
    def __init__(self) -> None:
        self._stack = ui.VStack(spacing=10)
        self._settings = carb.settings.get_settings()
        with self._stack:
            ClusterCompositingSettingsFrame("Cluster Compositing Settings", parent=self)
            ClusterNetworkSettingsFrame("Network Settings", parent=self)

            ui.Spacer()
