# 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

import carb
import omni.kit.commands
import omni.kit.undo
import omni.ui as ui
from pxr import Sdf

from .prim_selection_payload import PrimSelectionPayload


class ValueTypeNameItem(ui.AbstractItem):
    def __init__(self, value_type_name: Sdf.ValueTypeName, value_type_name_str: str):
        super().__init__()
        self.value_type_name = value_type_name
        self.model = ui.SimpleStringModel(value_type_name_str)


class VariabilityItem(ui.AbstractItem):
    def __init__(self, variability: Sdf.Variability, Variability_str: str):
        super().__init__()
        self.variability = variability
        self.model = ui.SimpleStringModel(Variability_str)


class ComboBoxModel(ui.AbstractItemModel):
    def __init__(self):
        super().__init__()

        self._current_index = ui.SimpleIntModel()
        self._current_index.add_value_changed_fn(lambda a: self._item_changed(None))

        self._items = []
        self._populate_items()

    def _populate_items(self):
        ...

    def get_item_children(self, item):
        return self._items

    def get_item_value_model(self, item, column_id):
        if item is None:
            return self._current_index
        return item.model

    def get_selected_item(self):
        return self._items[self._current_index.get_value_as_int()]


class ValueTypeNameModel(ComboBoxModel):
    def _populate_items(self):
        for property in dir(Sdf.ValueTypeNames):
            type_name = getattr(Sdf.ValueTypeNames, property)
            if isinstance(type_name, Sdf.ValueTypeName):
                self._items.append(ValueTypeNameItem(type_name, property))


class VariabilityModel(ComboBoxModel):
    def _populate_items(self):
        self._items = [
            VariabilityItem(Sdf.VariabilityVarying, "Varying"),
            VariabilityItem(Sdf.VariabilityUniform, "Uniform"),
        ]


class AddAttributePopup:
    def __init__(self):
        self._add_attribute_window = ui.Window(
            "Add Attribute...", visible=False, flags=ui.WINDOW_FLAGS_NO_RESIZE, auto_resize=True
        )
        with self._add_attribute_window.frame:
            with ui.VStack(height=0, spacing=5):
                LABEL_WIDTH = 80
                WIDGET_WIDTH = 300
                with ui.HStack(width=0):
                    ui.Label("Name", width=LABEL_WIDTH)
                    self._add_attr_name_field = ui.StringField(width=WIDGET_WIDTH)
                    self._add_attr_name_field.model.set_value("new_attr")

                with ui.HStack(width=0):
                    ui.Label("Type", width=LABEL_WIDTH)

                    self._value_type_name_model = ValueTypeNameModel()
                    ui.ComboBox(self._value_type_name_model, width=WIDGET_WIDTH)

                with ui.HStack(width=0):
                    ui.Label("Custom", width=LABEL_WIDTH)
                    self._custom_checkbox = ui.CheckBox()
                    self._custom_checkbox.model.set_value(True)

                with ui.HStack(width=0):
                    ui.Label("Variability", width=LABEL_WIDTH)

                    self._variability_model = VariabilityModel()
                    ui.ComboBox(self._variability_model, width=WIDGET_WIDTH)

                self._error_msg_label = ui.Label(
                    "",
                    visible=False,
                    style={"color": 0xFF0000FF},
                )

                with ui.HStack():
                    ui.Spacer()

                    def on_add(weak_self):
                        ref_self = weak_self()
                        error_message = ""
                        if ref_self:
                            stage = ref_self._add_attribute_window_payload.get_stage()
                            if stage:
                                selected_item = ref_self._value_type_name_model.get_selected_item()
                                value_type_name = selected_item.value_type_name
                                attr_name = ref_self._add_attr_name_field.model.get_value_as_string()
                                custom = ref_self._custom_checkbox.model.get_value_as_bool()
                                variability = ref_self._variability_model.get_selected_item().variability

                                if Sdf.Path.IsValidNamespacedIdentifier(attr_name):
                                    try:
                                        omni.kit.undo.begin_group()
                                        for path in ref_self._add_attribute_window_payload:
                                            path = Sdf.Path(path).AppendProperty(attr_name)

                                            prop = stage.GetPropertyAtPath(path)
                                            if prop:
                                                error_message = "One or more attribute to be created already exists."
                                            else:
                                                omni.kit.commands.execute(
                                                    "CreateUsdAttributeOnPath",
                                                    attr_path=path,
                                                    attr_type=value_type_name,
                                                    custom=custom,
                                                    variability=variability,
                                                )
                                    finally:
                                        omni.kit.undo.end_group()
                                else:
                                    error_message = f'Invalid identifier "{attr_name}"'

                        if error_message:
                            ref_self._error_msg_label.visible = True
                            ref_self._error_msg_label.text = error_message
                            carb.log_warn(error_message)
                        else:
                            ref_self._add_attribute_window.visible = False

                    ui.Button("Add", clicked_fn=lambda weak_self=weakref.ref(self): on_add(weak_self))

                    def on_cancel(weak_self):
                        ref_self = weak_self()
                        if ref_self:
                            ref_self._add_attribute_window.visible = False

                    ui.Button("Cancel", clicked_fn=lambda weak_self=weakref.ref(self): on_cancel(weak_self))
                    ui.Spacer()

    def show_fn(self, objects: dict):
        if not "prim_list" in objects or not "stage" in objects:
            return False
        stage = objects["stage"]
        if not stage:
            return False
        return len(objects["prim_list"]) >= 1

    def click_fn(self, payload: PrimSelectionPayload):
        self._add_attribute_window.visible = True
        self._add_attribute_window_payload = payload
        self._error_msg_label.visible = False

    def __del__(self):
        self.destroy()

    def destroy(self):
        self._add_attribute_window = None
        self._add_attribute_window_payload = None
        self._error_msg_label = None
        self._add_attr_name_field = None
        self._value_type_name_model = None
