"""
* Copyright (c) 2021, 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 pathlib import Path
from typing import Union

import carb
import omni.graph.core as og
import omni.kit.notification_manager as nm
import omni.ui as ui
from pxr import Sdf


# ======================================================================
def get_icons_dir() -> str:
    """Returns the location of the directory containing the icons used by these windows"""
    return Path(__file__).parent.parent.parent.parent.parent.joinpath("icons")


# ======================================================================
class DestructibleButton(ui.Button):
    """Class that enhances the standard button with a destroy method that removes callbacks"""

    def destroy(self):
        """Called when the button is being destroyed"""
        self.set_clicked_fn(None)


# ======================================================================
class DestructibleImage(ui.Image):
    """Class that enhances the standard image with a destroy method that removes callbacks"""

    def destroy(self):
        """Called when the image is being destroyed"""
        self.set_mouse_pressed_fn(None)


INSTANCE_NODES_WHITELIST = {
    # Safe to use despite state information in the node
    "omni.graph.action.OnTick",
    "omni.graph.core.ReadVariable",
    "omni.graph.core.WriteVariable",
    "omni.graph.instancing.ReadGraphVariable",
    "omni.graph.instancing.WriteGraphVariable",
    "omni.graph.nodes.FindPrims",
    "omni.graph.nodes.ReadPrimAttribute",
    "omni.graph.nodes.WritePrimAttribute",
    "omni.graph.nodes.ReadPrimBundle",
    "omni.graph.ui.OnNewFrame",
    "omni.graph.ui.PrintText",
    # The following temporal nodes are ignored because there is no way to
    # set them up with instances at the moment. Any uses of them use explicit paths,
    # not the graph target node
    "omni.graph.nodes.MoveToTarget",
    "omni.graph.nodes.MoveToTransform",
    "omni.graph.nodes.RotateToOrientation",
    "omni.graph.nodes.RotateToTarget",
    "omni.graph.nodes.ScaleToSize",
    "omni.graph.nodes.TranslateToLocation",
    "omni.graph.nodes.TranslateToTarget",
    # visualization nodes are safe - internal state is for caching
    "omni.graph.visualization.nodes.DrawLine",
    "omni.graph.visualization.nodes.DrawLabel",
    "omni.graph.visualization.nodes.DrawScreenSpaceText",
}


def validate_graph_for_instancing(graph_path: Union[str, Sdf.Path]):
    """Validate the node types used in a graph are okay to be used with instancing,
    and post a warning if they are not"""
    graph = og.get_graph_by_path(str(graph_path))
    if not graph:
        return

    problematic_nodes = set()
    nodes = graph.get_nodes()
    for node in nodes:
        node_type = node.get_node_type()
        if node_type and node_type.has_state() and node_type.get_node_type() not in INSTANCE_NODES_WHITELIST:
            problematic_nodes.add(node_type.get_node_type())

    if len(problematic_nodes) == 0:
        return

    newline = "\n"
    node_list = newline.join(sorted(problematic_nodes))

    warning_string = (
        f"The OmniGraph at path {graph_path} contains nodes that may not execute as"
        f" expected with multiple instances{newline}{node_list}"
    )
    nm.post_notification(warning_string, status=nm.NotificationStatus.WARNING)
    carb.log_warn(warning_string)


class Prompt:
    def __init__(
        self,
        title,
        text,
        ok_button_text="OK",
        cancel_button_text=None,
        middle_button_text=None,
        ok_button_fn=None,
        cancel_button_fn=None,
        middle_button_fn=None,
        modal=False,
    ):
        self._title = title
        self._text = text
        self._cancel_button_text = cancel_button_text
        self._cancel_button_fn = cancel_button_fn
        self._ok_button_fn = ok_button_fn
        self._ok_button_text = ok_button_text
        self._middle_button_text = middle_button_text
        self._middle_button_fn = middle_button_fn
        self._modal = modal
        self._build_ui()

    def __del__(self):
        self._cancel_button_fn = None
        self._ok_button_fn = None

    def __enter__(self):
        self._window.show()
        return self

    def __exit__(self, exit_type, value, trace):
        self._window.hide()

    def show(self):
        self._window.visible = True

    def hide(self):
        self._window.visible = False

    def is_visible(self):
        return self._window.visible

    def set_text(self, text):
        self._text_label.text = text

    def set_confirm_fn(self, on_ok_button_clicked):
        self._ok_button_fn = on_ok_button_clicked

    def set_cancel_fn(self, on_cancel_button_clicked):
        self._cancel_button_fn = on_cancel_button_clicked

    def set_middle_button_fn(self, on_middle_button_clicked):
        self._middle_button_fn = on_middle_button_clicked

    def _on_ok_button_fn(self):
        self.hide()
        if self._ok_button_fn:
            self._ok_button_fn()

    def _on_cancel_button_fn(self):
        self.hide()
        if self._cancel_button_fn:
            self._cancel_button_fn()

    def _on_middle_button_fn(self):
        self.hide()
        if self._middle_button_fn:
            self._middle_button_fn()

    def _build_ui(self):
        self._window = ui.Window(self._title, visible=False, height=0, dockPreference=ui.DockPreference.DISABLED)
        self._window.flags = (
            ui.WINDOW_FLAGS_NO_COLLAPSE
            | ui.WINDOW_FLAGS_NO_RESIZE
            | ui.WINDOW_FLAGS_NO_SCROLLBAR
            | ui.WINDOW_FLAGS_NO_RESIZE
            | ui.WINDOW_FLAGS_NO_MOVE
        )

        if self._modal:
            self._window.flags = self._window.flags | ui.WINDOW_FLAGS_MODAL

        with self._window.frame:
            with ui.VStack(height=0):
                ui.Spacer(width=0, height=10)
                with ui.HStack(height=0):
                    ui.Spacer()
                    self._text_label = ui.Label(self._text, word_wrap=True, width=self._window.width - 80, height=0)
                    ui.Spacer()
                ui.Spacer(width=0, height=10)
                with ui.HStack(height=0):
                    ui.Spacer(height=0)
                    if self._ok_button_text:
                        ok_button = ui.Button(self._ok_button_text, width=60, height=0)
                        ok_button.set_clicked_fn(self._on_ok_button_fn)
                    if self._middle_button_text:
                        middle_button = ui.Button(self._middle_button_text, width=60, height=0)
                        middle_button.set_clicked_fn(self._on_middle_button_fn)
                    if self._cancel_button_text:
                        cancel_button = ui.Button(self._cancel_button_text, width=60, height=0)
                        cancel_button.set_clicked_fn(self._on_cancel_button_fn)
                    ui.Spacer(height=0)
                ui.Spacer(width=0, height=10)
