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

__all__ = ["StyleComboValueModel", "StyleComboNameValueItem", "StyleComboItemModel", "StylePropertyType", "StyleColorComponentItem", "StyleColorModel", "StyleItem", "StyleGroupItem", "StyleModel"]

from os import remove
from typing import Callable, Union, List, cast
import omni.ui as ui
from enum import Enum
import copy, carb

StyleFloatProperty = [
    "border_radius",
    "border_width",
    "font_size",
    "margin",
    "margin_width",
    "margin_height",
    "padding",
    "padding_width",
    "padding_height",
    "secondary_padding",
    "scrollbar_size",
]

StyleColorProperties = [
    "background_color",
    "background_gradient_color",
    "background_selected_color",
    "border_color",
    "color",
    "selected_color",
    "secondary_color",
    "secondary_selected_color",
    "debug_color",
]

StyleEnumProperty = ["corner_flag", "alignment", "fill_policy", "draw_mode", "stack_direction"]

StyleStringProperty = ["image_url"]


def get_enum_class(property_name: str) -> Union[Callable, None]:
    if property_name == "corner_flag":
        return ui.CornerFlag
    elif property_name == "alignment":
        return ui.Alignment
    elif property_name == "fill_policy":
        return ui.FillPolicy
    elif property_name == "draw_mode":
        return ui.SliderDrawMode
    elif property_name == "stack_direction":
        return ui.Direction
    else:
        return None


class StyleComboValueModel(ui.AbstractValueModel):
    """
    Model to store a pair (label, value of arbitrary type) for use in a ComboBox
    """

    def __init__(self, key, value):
        """
        Args:
            value: the instance of the Style value
        """
        ui.AbstractValueModel.__init__(self)
        self.key = key
        self.value = value

    def __repr__(self):
        return f'"StyleComboValueModel name:{self.key} value:{int(self.value)}"'

    def get_value_as_string(self) -> str:
        """
        this is called to get the label of the combo box item
        """
        return self.key

    def get_style_value(self) -> int:
        """
        we call this to get the value of the combo box item
        """
        return int(self.value)


class StyleComboNameValueItem(ui.AbstractItem):
    def __init__(self, key, value):
        super().__init__()
        self.model = StyleComboValueModel(key, value)

    def __repr__(self):
        return f'"StyleComboNameValueItem {self.model}"'


class StyleComboItemModel(ui.AbstractItemModel):
    """
    Model for a combo box - for each setting we have  a dictionary of key, values
    """

    def __init__(self, class_obj: Callable, default_value: int):
        super().__init__()
        self._class = class_obj
        self._items = []
        self._default_value = default_value

        default_index = 0
        current_index = 0
        for key, value in class_obj.__members__.items():
            if self._default_value == value:
                default_index = current_index
            self._items.append(StyleComboNameValueItem(key, value))
            current_index += 1

        self._current_index = ui.SimpleIntModel(default_index)
        self._current_index.add_value_changed_fn(self._current_index_changed)

    def get_current_value(self):
        return self._items[self._current_index.as_int].model.get_style_value()

    def _current_index_changed(self, model):
        self._item_changed(None)

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

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


class StylePropertyType(Enum):
    ENUM = 0
    COLOR = 1
    FLOAT = 2
    STRING = 3


def get_property_type(property_name: str) -> StylePropertyType:
    if property_name in StyleStringProperty:
        return StylePropertyType.STRING
    elif property_name in StyleEnumProperty:
        return StylePropertyType.ENUM
    elif property_name in StyleColorProperties:
        return StylePropertyType.COLOR
    elif property_name in StyleFloatProperty:
        return StylePropertyType.FLOAT
    else:
        return StylePropertyType.FLOAT


class StyleColorComponentItem(ui.AbstractItem):
    def __init__(self, model: ui.SimpleFloatModel) -> None:
        super().__init__()
        self.model = model


class StyleColorModel(ui.AbstractItemModel):
    """ define a model for a style color, and enable coversion from int32 etc """

    def __init__(self, value: int) -> None:
        super().__init__()
        self._value = value

        # Create root model
        self._root_model = ui.SimpleIntModel()
        self._root_model.add_value_changed_fn(lambda a: self._item_changed(None))

        # convert Value from int
        red = self._value & 255
        green = (self._value >> 8) & 255
        blue = (self._value >> 16) & 255
        alpha = (self._value >> 24) & 255
        rgba_values = [red / 255, green / 255, blue / 255, alpha / 255]

        # Create three models per component
        self._items: List[StyleColorComponentItem] = [
            StyleColorComponentItem(ui.SimpleFloatModel(rgba_values[i])) for i in range(4)
        ]
        for item in self._items:
            item.model.add_value_changed_fn(lambda a, item=item: self._on_value_changed(item))

    def set_value(self, value: int):
        self._value = value

    def set_components_values(self, values: List[float]):
        if len(values) != 4:
            print("Error: List need to be exactly 4 long")
            return

        for i in range(4):
            self._items[i].model.set_value(values[i])

    def set_components_values_as_int(self, values: List[int]):
        if len(values) != 4:
            print("Error: List need to be exactly 4 long")
            return

        for i in range(4):
            self._items[i].model.set_value(values[i] / 255)

    def get_value_as_int8s(self) -> List[int]:

        result: List[int] = []
        for i in range(4):
            result.append(int(round(self._items[i].model.get_value_as_float() * 255)))

        return result

    def get_value_as_floats(self) -> List[float]:

        result: List[float] = []
        for i in range(4):
            result.append(self._items[i].model.get_value_as_float())

        return result

    def get_value_as_int32(self) -> int:
        red, green, blue, alpha = self.get_value_as_int8s()

        # we store then in AABBGGRR
        result = (alpha << 24) + (blue << 16) + (green << 8) + red
        return result

    def _on_value_changed(self, item):
        self._item_changed(item)

    def get_item_children(self, parentItem: ui.AbstractItem) -> List[StyleColorComponentItem]:
        return self._items

    def get_item_value_model(self, item: ui.AbstractItem, column_id: int) -> ui.AbstractValueModel:
        if item is None:
            return self._root_model

        return item.model

    def begin_edit(self, item):
        """
        TODO: if we don't add this override (even without a real implementation) we get crashes
        """
        pass

    def end_edit(self, item):
        """
        TODO: if we don't add this override (even without a real implementation) we get crashes
        """
        pass


class StyleItem(ui.AbstractItem):
    """A single AbstractItemModel item that represents a single prim"""

    def __init__(
        self,
        property_name: str,
        style_type: StylePropertyType,
        values: dict,
        parent_model: "StyleModel",
        parent_item: "StyleGroupItem",
    ):
        super().__init__()

        self.parent_item = parent_item
        self.parent_model = parent_model
        self.property = property_name
        self.style_type = style_type
        self.value = values[property_name]
        self._values = values

        self._model: Union[ui.AbstractValueModel, ui.AbstractItemModel, None] = None

        if self.style_type == StylePropertyType.STRING:
            self._model = ui.SimpleStringModel(str(self.value))
        elif self.style_type == StylePropertyType.ENUM:
            enum_cls = get_enum_class(property_name)
            self._model = StyleComboItemModel(enum_cls, self.value)
        elif self.style_type == StylePropertyType.COLOR:
            self._model = StyleColorModel(int(self.value))
        elif self.style_type == StylePropertyType.FLOAT:
            self._model = ui.SimpleFloatModel(float(self.value))
        else:
            raise ValueError("The Style Type need to be either STRING, ENUM, COLOR or FLOAT")

        if self.style_type in [StylePropertyType.STRING, StylePropertyType.FLOAT]:

            def value_changed(model: ui.AbstractValueModel):
                self._values[self.property] = model.as_float
                self.parent_model.update()

            self._model.add_value_changed_fn(value_changed)
        elif self.style_type == StylePropertyType.COLOR:

            def value_changed(model: StyleColorModel, item: ui.AbstractItem):
                self._values[self.property] = model.get_value_as_int32()
                self.parent_model.update()

            self._model.add_item_changed_fn(value_changed)
        elif self.style_type == StylePropertyType.ENUM:

            def value_changed(model: StyleComboItemModel, item: ui.AbstractItem):
                self._values[self.property] = model.get_current_value()
                self.parent_model.update()

            self._model.add_item_changed_fn(value_changed)

    def delete(self):
        # we simply remove the key from the referenced Value Dict
        del self._values[self.property]

    def get_model(self) -> Union[ui.AbstractValueModel, None]:
        return self._model

    def __repr__(self):
        return f"<StyleItem '{self.property} . {self.style_type}'>"

    def __str__(self):
        return f"<StyleItem '{self.property} . {self.style_type}'>"


class StyleGroupItem(ui.AbstractItem):
    """A single AbstractItemModel item that represents a single prim"""

    def __init__(self, type_name: str, style_dict: dict, parent_model: "StyleModel"):
        super().__init__()

        self.children = []
        self.type_name = type_name
        self._parent_dict = style_dict
        self.style_dict = self._parent_dict[type_name]

        self.parent_model = parent_model

        self._name_model = ui.SimpleStringModel()
        self._name_model.set_value(type_name)

        def value_changed(model: ui.AbstractValueModel):
            new_name = model.as_string
            if new_name in self._parent_dict:
                carb.log_warn(f"Entry with type {new_name} already exits, skipping")
                return

            self._parent_dict[new_name] = self.style_dict
            del self._parent_dict[self.type_name]

            self.type_name = new_name
            self.parent_model.update()

        self._name_model.add_value_changed_fn(value_changed)

    def get_model(self):
        return self._name_model

    def delete_item(self, item: StyleItem):
        # we call delete on the item but also make sure we clear our cach for it
        self.children.remove(item)
        item.delete()

    def delete(self):
        # we simply remove the key from the referenced Value Dict
        del self._parent_dict[self.type_name]

    def __repr__(self):
        return f"<StyleItem '{self.type_name} . {self.style_dict}'>"

    def __str__(self):
        return f"<StyleItem '{self.type_name} . {self.style_dict}'>"


class StyleModel(ui.AbstractItemModel):
    """The item model that watches the stage"""

    def __init__(self, treeview):
        """Flat means the root node has all the children and children of children, etc."""
        super().__init__()

        self._widget = None
        self._style_dict: dict = {}
        self._is_generic = False
        self._root: List[StyleGroupItem] = []
        self._treeview: "StyleTreeView" = treeview

    def _update_style_dict(self, style: dict) -> dict:
        style_dict = copy.copy(style)
        # find if there are "Generic" entry (with no sub dictionary) and move them around
        remove_keys = []
        has_generic_entries = False
        for key, value in style_dict.items():
            if not type(value) == dict:
                has_generic_entries = True
                break

        if has_generic_entries:
            style_dict["All"] = {}

        for key, value in style_dict.items():
            if not type(value) == dict:
                style_dict["All"][key] = value
                remove_keys.append(key)

        for a_key in remove_keys:
            del style_dict[a_key]

        return style_dict

    def set_widget(self, widget: ui.Widget):
        self._widget = widget
        self.update_cached_style()

        self._root: List[StyleGroupItem] = []
        self._item_changed(None)

    def update_cached_style(self):
        if not self._widget.style:
            style_dict = {}
        else:
            style_dict = cast(dict, self._widget.style)
        self._style_dict = self._update_style_dict(style_dict)
        # rebuild
        self._item_changed(None)

    def can_item_have_children(self, parentItem: Union[StyleGroupItem, StyleItem]) -> bool:
        if not parentItem:
            return True

        if isinstance(parentItem, StyleGroupItem):
            return True

        return False

    def get_item_children(self, item: Union[StyleGroupItem, StyleItem, None]) -> Union[None, list]:
        """Reimplemented from AbstractItemModel"""
        if not item:
            if not self._root:
                self._root = []
                for key, value in self._style_dict.items():
                    group_item = StyleGroupItem(key, self._style_dict, self)
                    self._root.append(group_item)
            return self._root

        if isinstance(item, StyleItem):
            return []

        if isinstance(item, StyleGroupItem):
            children = []
            for key, value in item.style_dict.items():
                property_type = get_property_type(key)
                # value is passed by reference and updated internally on change
                style_item = StyleItem(key, property_type, item.style_dict, self, item)
                children.append(style_item)

            item.children = children

            return item.children

        # we should not get here ..
        return []

    def duplicate_item(self, item: StyleGroupItem):
        new_item = item.type_name + "1"
        self._style_dict[new_item] = copy.copy(self._style_dict[item.type_name])

        new_item = StyleGroupItem(new_item, self._style_dict, self)
        self._root.append(new_item)

        self._item_changed(None)
        self.update()

    def delete_item(self, item: Union[StyleGroupItem, StyleItem]):

        if isinstance(item, StyleItem):
            item.parent_item.delete_item(item)
            self._item_changed(item.parent_item)
            self.update()
        elif isinstance(item, StyleGroupItem):
            item.delete()
            self._root.remove(item)
            self._item_changed(None)
            self.update()
        else:
            carb.log_error("Delete_item only support  Union[StyleGroupItem, StyleItem]")

    def update(self):
        style_dict = cast(dict, copy.copy(self._style_dict))

        if "All" in style_dict:
            # update Generic Entry
            for key, value in style_dict["All"].items():
                style_dict[key] = value

            del style_dict["All"]

        self._widget.set_style(style_dict)

    def get_item_value_model_count(self, item):
        """Reimplemented from AbstractItemModel"""
        return 1

    def get_item_value_model(self, item: Union[StyleItem, StyleGroupItem], column_id: int):
        """Reimplemented from AbstractItemModel"""

        if isinstance(item, StyleGroupItem):
            if column_id == 0:
                return item.get_model()

        elif isinstance(item, StyleItem):
            if column_id == 0:
                return ui.SimpleStringModel(item.property)
            if column_id == 1:
                return item.get_model()

    def destroy(self):
        self._widget = None
