# 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__ = ["WidgetItem", "InspectorTreeModel"]

from typing import Union
import omni.ui as ui
from omni.ui_query import OmniUIQuery


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

    def __init__(self, widget: ui.Widget, path: str):
        super().__init__()
        self.widget = widget
        self.path = path

        self.children = []
        # True when it's necessary to repopulate its children
        self.populated = False

        widget_name = path.split("/")[-1]
        self._widget_name = widget_name

        # Get Type
        self._widget_type = str(type(widget)).split(".")[-1]
        self._widget_type = self._widget_type.split("'")[0]
        self.type_model = ui.SimpleStringModel(self._widget_type)
        self.expanded = False

    def __repr__(self):
        return f"<Omni::UI Widget '{str(type(self.widget))}'>"

    def __str__(self):
        return f"WidgetItem: {self.widget} @ {self.path}"


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

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

        self._root = None
        self._frame = window.frame if window else None
        self._window_name = window.title if window else ""
        self._window = window
        self._frame_item = WidgetItem(self._frame, f"{self._window_name}//Frame")

    def get_item_for_widget(self, widget: ui.Widget) -> Union[WidgetItem, None]:
        def search_for_item(item: WidgetItem, widget: ui.Widget) -> Union[WidgetItem, None]:
            if item.widget == widget:
                return item

            if isinstance(item.widget, ui.Container) or isinstance(item.widget, ui.TreeView):
                for a_child_item in self.get_item_children(item):
                    matched_item = search_for_item(a_child_item, widget)
                    if matched_item:
                        return matched_item

            return None

        return search_for_item(self._frame_item, widget)

    def get_item_with_path_for_widget(self, widget: ui.Widget):
        parents = []

        def search_for_item(item: WidgetItem, widget: ui.Widget):
            if item.widget == widget:
                return item

            if isinstance(item.widget, ui.Container) or isinstance(item.widget, ui.TreeView):
                for a_child_item in self.get_item_children(item):
                    matched_item = search_for_item(a_child_item, widget)
                    if matched_item:
                        parents.append(item)
                        return matched_item

            return None

        item = search_for_item(self._frame_item, widget)
        return item, parents

    def can_item_have_children(self, parentItem: WidgetItem) -> bool:
        if not parentItem:
            return True

        if not parentItem.widget:
            return True

        if issubclass(type(parentItem.widget), ui.Container):
            return True
        elif issubclass(type(parentItem.widget), ui.TreeView):
            return True
        else:
            return False

    def get_item_children(self, item: WidgetItem):
        """Reimplemented from AbstractItemModel"""
        if item is None:
            return [self._frame_item]

        if item.populated:
            return item.children

        widget = item.widget

        # return the children of the Container
        item.children = []
        index = 0
        for w in ui.Inspector.get_children(widget):
            widget_path = OmniUIQuery.get_widget_path(self._window, w)
            if widget_path:
                # TODO: There are widgets that have no paths... why?
                new_item = WidgetItem(w, widget_path)
                item.children.append(new_item)

        item.populated = True
        return item.children

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

    def get_widget_type(self, item):
        return item._widget_type

    def get_widget_path(self, item):
        return item.path

    def get_item_value_model(self, item, column_id):
        """Reimplemented from AbstractItemModel"""
        if item is None:
            return ui.SimpleStringModel("Window")

        if column_id == 0:
            return ui.SimpleStringModel(item._widget_name)

        if column_id == 1:
            style = item.widget.style if item and item.widget else None
            return ui.SimpleBoolModel(not style == None)

        if column_id == 2:
            height = item.widget.computed_height
            return ui.SimpleStringModel("{:10.2f}".format(height))
