# Copyright (c) 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 weakref
from functools import lru_cache
from typing import Callable

import carb
import omni.ui as ui
import omni.usd
from omni.kit.window.property.templates import HORIZONTAL_SPACING, LABEL_HEIGHT, LABEL_WIDTH, SimplePropertyWidget

from .add_attribute_popup import AddAttributePopup
from .context_menu import ContextMenu, ContextMenuEvent
from .prim_selection_payload import PrimSelectionPayload

g_singleton = None


@lru_cache()
def _get_plus_glyph():
    return omni.kit.ui.get_custom_glyph_code("${glyphs}/menu_context.svg")


class Constant:
    def __setattr__(self, name, value):
        raise Exception(f"Can't change Constant.{name}")

    MIXED = "Mixed"
    MIXED_COLOR = 0xFFCC9E61
    LABEL_COLOR = 0xFF9E9E9E
    LABEL_FONT_SIZE = 14
    LABEL_WIDTH = 80
    ADD_BUTTON_SIZE = 52


class ButtonMenuEntry:
    def __init__(
        self,
        path: str,
        glyph: str = None,
        name_fn: Callable = None,
        show_fn: Callable = None,
        enabled_fn: Callable = None,
        onclick_fn: Callable = None,
        add_to_context_menu: bool = False,
    ):
        self.path = path
        self.glyph = glyph
        self.name_fn = name_fn
        self.show_fn = show_fn
        self.enabled_fn = enabled_fn
        self.onclick_fn = onclick_fn
        if add_to_context_menu:
            self.context_menu = omni.kit.context_menu.add_menu(self.get_context_menu(), "ADD", "")
        else:
            self.context_menu = None

    def clean(self):
        self.context_menu = None

    def get_dict(self, payload, name=None):
        data = {}
        if name:
            data = {"name": name}
        else:
            data = {"name": self.path}

        data["glyph"] = self.glyph
        if self.name_fn:
            data["name_fn"] = self.name_fn
        if self.show_fn:
            data["show_fn"] = self.show_fn
        if self.enabled_fn:
            data["enabled_fn"] = self.enabled_fn
        if self.onclick_fn:
            data["onclick_fn"] = lambda o, p=payload: self.onclick_fn(payload=p)

        return data

    def get_context_menu(self):
        menu_sublist = []
        sublist = menu_sublist
        parts = self.path.split("/")
        if len(parts) > 1:
            last_name = parts.pop()
            first_name = parts.pop(0)
            context_root = {'name': { first_name: []}}
            context_item = context_root["name"][first_name]

            while parts:
                name = parts.pop(0)
                context_item.append({'name': {name: []}})
                context_item = context_item[0]["name"][name]

            context_item.append({
                            "name": last_name,
                            "glyph": self.glyph,
                            "name_fn": self.name_fn,
                            "show_fn": self.show_fn,
                            "enabled_fn": self.enabled_fn,
                            "onclick_fn": lambda o, f=weakref.ref(
                                self.onclick_fn
                            ) if self.onclick_fn else None: PrimPathWidget.payload_wrapper(
                                objects=o, onclick_fn_weak=f
                            ),
                        })
            return context_root

        else:
            context_item = {
                "name": self.path,
                "glyph": self.glyph,
                "name_fn": self.name_fn,
                "show_fn": self.show_fn,
                "enabled_fn": self.enabled_fn,
                "onclick_fn": lambda o, f=weakref.ref(
                    self.onclick_fn
                ) if self.onclick_fn else None: PrimPathWidget.payload_wrapper(objects=o, onclick_fn_weak=f),
            }

        return context_item


class PrimPathWidget(SimplePropertyWidget):
    def __init__(self):
        super().__init__(title="Path", collapsable=False)
        self._context_menu = ContextMenu()
        self._path_draw_items = []
        self._button_menu_items = []
        self._path_item_padding = 0.0

        global g_singleton
        g_singleton = self

        self.add_path_item(self._build_prim_name_widget)
        self.add_path_item(self._build_prim_path_widget)
        # self.add_path_item(self._build_prim_stage_widget)
        self.add_path_item(self._build_prim_instanceable_widget)
        self.add_path_item(self._build_prim_large_selection_widget)
        self._add_create_attribute_menu()

    def __del__(self):
        global g_singleton
        g_singleton = None

    def clean(self):
        self._path_draw_items = []
        if self._add_attribute_popup:
            self._add_attribute_popup.destroy()
            self._add_attribute_popup = None
        super().clean()

    def on_new_payload(self, payload):
        if not super().on_new_payload(payload, ignore_large_selection=True):
            return False

        self._anchor_prim = None
        stage = payload.get_stage()
        instance = []
        if stage:
            for prim_path in payload:
                prim = stage.GetPrimAtPath(prim_path)
                if prim:
                    self._anchor_prim = prim
                    instance.append(prim.IsInstanceable())

        self._show_instance = len(set(instance)) == 1
        return payload and len(payload) > 0

    @staticmethod
    def payload_wrapper(objects, onclick_fn_weak):
        onclick_fn = onclick_fn_weak()
        if onclick_fn:
            prim_paths = []
            if "prim_list" in objects:
                prim_paths = [prim.GetPath() for prim in objects["prim_list"]]
            payload = PrimSelectionPayload(weakref.ref(objects["stage"]), prim_paths)
            onclick_fn(payload=payload)

    def build_items(self):
        for item in self._path_draw_items:
            item()

    def _build_prim_name_widget(self):
        from omni.kit.property.usd.usd_property_widget import get_ui_style

        def get_names_as_string(prim_paths):
            paths = ""
            for index, path in enumerate(prim_paths):
                paths += f"{path.name}\n"
            return paths

        prim_paths = self._payload
        selected_info_name = "(nothing selected)"
        stage = self._payload.get_stage()
        tooltip = ""
        if len(prim_paths) > 1:
            selected_info_name = f"({len(prim_paths)} models selected) common attributes shown"
            for index, path in enumerate(prim_paths):
                tooltip += f"{path.name}\n"
                if index > 9:
                    tooltip += f"...."
                    break
        elif len(prim_paths) == 1:
            selected_info_name = f"{prim_paths[0].name}"

        with ui.HStack(height=0):
            weakref_menu = weakref.ref(self._context_menu)
            ui.Spacer(width=8)
            button_width = Constant.ADD_BUTTON_SIZE
            if get_ui_style() == "NvidiaLight":
                button_width = Constant.ADD_BUTTON_SIZE + 25
            add_button = ui.Button(f"{_get_plus_glyph()} Add", width=button_width, height=LABEL_HEIGHT, name="add")
            add_button.set_mouse_pressed_fn(
                lambda x, y, b, m, widget=add_button: self._on_mouse_pressed(b, weakref_menu, widget)
            )
            ui.Spacer(width=(Constant.LABEL_WIDTH + self._path_item_padding) - button_width)
            ui.Spacer(width=8)
            if get_ui_style() == "NvidiaLight":
                ui.Spacer(width=10)
            widget = ui.StringField(
                name="prims_name", height=LABEL_HEIGHT, tooltip=tooltip, tooltip_offset_y=22, enabled=False
            )
            widget.model.set_value(selected_info_name)
            widget.set_mouse_pressed_fn(
                lambda x, y, b, m, model=widget.model: self._build_copy_menu(
                    buttons=b,
                    context_menu=weakref_menu,
                    copy_fn=lambda _: self._copy_to_clipboard(to_copy=get_names_as_string(prim_paths=prim_paths)),
                )
            )

    def _build_prim_path_widget(self):
        from omni.kit.property.usd.usd_property_widget import get_ui_style

        def get_paths_as_string(prim_paths):
            paths = ""
            for index, path in enumerate(prim_paths):
                paths += f"{path}\n"
            return paths

        prim_paths = self._payload
        weakref_menu = weakref.ref(self._context_menu)
        selected_info_name = ""
        tooltip = ""
        style = {}
        if len(prim_paths) == 1:
            selected_info_name = f"{prim_paths[0]}"
        elif len(prim_paths) > 1:
            selected_info_name = Constant.MIXED
            style = {"Field": {"color": Constant.MIXED_COLOR}, "Tooltip": {"color": 0xFF333333}}

            for index, path in enumerate(prim_paths):
                tooltip += f"{path}\n"
                if index > 9:
                    tooltip += f"...."
                    break

        with ui.HStack(height=0):
            ui.Spacer(width=8)
            ui.Label(
                "Prim Path",
                name="path_label",
                width=Constant.LABEL_WIDTH + self._path_item_padding,
                height=LABEL_HEIGHT,
                style={"font_size": Constant.LABEL_FONT_SIZE},
            )
            ui.Spacer(width=8)
            if get_ui_style() == "NvidiaLight":
                ui.Spacer(width=10)
            widget = ui.StringField(
                name="prims_path", height=LABEL_HEIGHT, enabled=False, tooltip=tooltip, tooltip_offset_y=22, style=style
            )
            widget.model.set_value(selected_info_name)
            widget.set_mouse_pressed_fn(
                lambda x, y, b, m: self._build_copy_menu(
                    buttons=b,
                    context_menu=weakref_menu,
                    copy_fn=lambda _: self._copy_to_clipboard(to_copy=get_paths_as_string(prim_paths=prim_paths)),
                )
            )

    def _build_prim_stage_widget(self):
        from omni.kit.property.usd.usd_property_widget import get_ui_style

        stage_path = "unsaved"
        stage = self._payload.get_stage()
        weakref_menu = weakref.ref(self._context_menu)
        if stage and stage.GetRootLayer() and stage.GetRootLayer().realPath:
            stage_path = stage.GetRootLayer().realPath

        with ui.HStack(height=0):
            ui.Spacer(width=8)
            ui.Label(
                "Stage Path",
                name="stage_label",
                width=Constant.LABEL_WIDTH + self._path_item_padding,
                height=LABEL_HEIGHT,
                style={"font_size": Constant.LABEL_FONT_SIZE},
            )
            ui.Spacer(width=8)
            widget = ui.StringField(name="prims_stage", height=LABEL_HEIGHT, enabled=False)
            widget.model.set_value(stage_path)
            widget.set_mouse_pressed_fn(
                lambda x, y, b, m, model=widget.model: self._build_copy_menu(
                    buttons=b,
                    context_menu=weakref_menu,
                    copy_fn=lambda _: self._copy_to_clipboard(to_copy=model.as_string),
                )
            )

    def _build_prim_instanceable_widget(self):
        from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidgetBuilder, get_ui_style

        if not self._anchor_prim or not self._show_instance:
            return

        def on_instanceable_changed(model):
            index = model.as_bool
            stage = self._payload.get_stage()
            for path in self._payload:
                prim = stage.GetPrimAtPath(path)
                if prim:
                    prim.SetInstanceable(model.as_bool)
            self._collapsable_frame.rebuild()

        additional_label_kwargs = {}
        with ui.HStack(spacing=HORIZONTAL_SPACING):
            settings = carb.settings.get_settings()
            left_aligned = settings.get("ext/omni.kit.window.property/checkboxAlignment") == "left"

            if not left_aligned:
                additional_label_kwargs["width"] = 0
            else:
                additional_label_kwargs["width"] = (Constant.LABEL_WIDTH + self._path_item_padding) - 8
            if get_ui_style() == "NvidiaLight":
                additional_label_kwargs["word_wrap"] = False
            ui.Spacer(width=3)
            UsdPropertiesWidgetBuilder._create_label("Instanceable", {}, additional_label_kwargs)
            if not left_aligned:
                ui.Spacer(width=15)
            else:
                ui.Spacer(width=0)
            with ui.VStack(width=10):
                ui.Spacer()
                check_box = ui.CheckBox(width=10, height=0, name="greenCheck")
                check_box.model.set_value(self._anchor_prim.IsInstanceable())
                check_box.model.add_value_changed_fn(on_instanceable_changed)
                ui.Spacer()
            if left_aligned:
                ui.Spacer(width=5)

    def _build_prim_large_selection_widget(self):
        def show_all():
            if self._payload:
                self._payload.set_large_selection_override(True)
                property_window = omni.kit.window.property.get_window()
                if property_window:
                    property_window.request_rebuild()

        if self._payload.is_large_selection():
            from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidgetBuilder

            ui.Separator()
            UsdPropertiesWidgetBuilder._create_label(
                f"You have selected {len(self._payload)} Prims, to preserve fast performance the Property Widget is not showing above the current limit of {self._payload.get_large_selection_count()} Prims. Press the button below to show it anyway but expect it to take some time",
                {},
                additional_label_kwargs={"width": omni.ui.Percent(100), "alignment": ui.Alignment.CENTER},
            )
            ui.Button("Skip Large Selection Protection", clicked_fn=lambda: show_all())

    def _add_create_attribute_menu(self):
        self._add_attribute_popup = AddAttributePopup()

        self.add_button_menu_entry(
            "Attribute",
            show_fn=lambda objects, weak_self=weakref.ref(self): weak_self()._add_attribute_popup.show_fn(objects)
            if weak_self()
            else None,
            onclick_fn=lambda payload, weak_self=weakref.ref(self): weak_self()._add_attribute_popup.click_fn(payload)
            if weak_self()
            else None,
        )

    def _on_mouse_pressed(self, button, context_menu, widget):
        """Called when the user press the mouse button on the item"""
        if button != 0:
            # It's for LMB menu only
            return

        if not context_menu():
            return

        # Show the menu
        xpos = (int)(widget.screen_position_x)
        ypos = (int)(widget.screen_position_y + widget.computed_content_height)
        context_menu().on_mouse_event(ContextMenuEvent(self._payload, self._button_menu_items, xpos, ypos))

    def _build_copy_menu(self, buttons, context_menu, copy_fn):
        if buttons != 1:
            # It's for right mouse button for menu only
            return

        if not context_menu():
            return

        # setup menu
        menu_list = [{"name": "Copy to clipboard", "glyph": "menu_link.svg", "onclick_fn": copy_fn}]

        # show menu
        context_menu().show_context_menu(menu_list=menu_list)

    def _copy_to_clipboard(self, to_copy):
        omni.kit.clipboard.copy(to_copy)

    @staticmethod
    def add_button_menu_entry(
        path: str,
        glyph: str = None,
        name_fn=None,
        show_fn: Callable = None,
        enabled_fn: Callable = None,
        onclick_fn: Callable = None,
        add_to_context_menu: bool = True,
    ):
        item = None
        if g_singleton:
            item = ButtonMenuEntry(path, glyph, name_fn, show_fn, enabled_fn, onclick_fn, add_to_context_menu)
            g_singleton._button_menu_items.append(item)
        else:
            carb.log_error(f"PrimPathWidget not initalized")
        return item

    @staticmethod
    def remove_button_menu_entry(item: ButtonMenuEntry):
        if g_singleton:
            g_singleton._button_menu_items.remove(item)
            item.clean()
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def get_button_menu_entries():
        if g_singleton:
            return g_singleton._button_menu_items
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def add_path_item(draw_fn: Callable):
        if g_singleton:
            g_singleton._path_draw_items.append(draw_fn)
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def remove_path_item(draw_fn: Callable):
        if g_singleton:
            g_singleton._path_draw_items.remove(draw_fn)
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def get_path_items():
        if g_singleton:
            return g_singleton._path_draw_items
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def rebuild():
        if g_singleton:
            if g_singleton._collapsable_frame:
                g_singleton._collapsable_frame.rebuild()
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def set_path_item_padding(padding: float):
        if g_singleton:
            g_singleton._path_item_padding = padding
        else:
            carb.log_error(f"PrimPathWidget not initalized")

    @staticmethod
    def get_path_item_padding(padding: float):
        if g_singleton:
            return g_singleton._path_item_padding
        else:
            carb.log_error(f"PrimPathWidget not initalized")

        return 0.0
