# 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.
#

__all__ = [
    "LABEL_WIDTH", "LABEL_WIDTH_LIGHT", "LABEL_HEIGHT", "HORIZONTAL_SPACING",
    "build_frame_header", "SimplePropertyWidget",
    "GroupHeaderContextMenu", "GroupHeaderContextMenuEvent",
    "PropertyWidget",  # omni.kit.property.environment tries to import this from here
]

import asyncio
import functools
import traceback

import carb
import omni.ui as ui
import omni.kit.window.property.managed_frame
import omni.kit.app
from omni.kit.async_engine import run_coroutine
from omni.kit.widget.highlight_label import HighlightLabel
from omni.kit.window.property.property_widget import PropertyWidget
from omni.kit.window.property.style import get_style

from .header_context_menu import GroupHeaderContextMenu, GroupHeaderContextMenuEvent

LABEL_WIDTH = 160
LABEL_WIDTH_LIGHT = 235
LABEL_HEIGHT = 18
HORIZONTAL_SPACING = 4


def handle_exception(func):
    """
    Decorator to print exception in async functions
    """

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except asyncio.CancelledError:
            # We always cancel the task. It's not a problem.
            pass
        except Exception as e:
            carb.log_error(f"Exception when async '{func}'")
            carb.log_error(f"{e}")
            carb.log_error(f"{traceback.format_exc()}")

    return wrapper


def build_frame_header(collapsed, text: str, id: str = None):
    """Custom header for CollapsibleFrame"""
    if id is None:
        id = text

    if collapsed:
        alignment = ui.Alignment.RIGHT_CENTER
        width = 5
        height = 7
    else:
        alignment = ui.Alignment.CENTER_BOTTOM
        width = 7
        height = 5

    header_stack = ui.HStack(spacing=8)
    with header_stack:
        with ui.VStack(width=0):
            ui.Spacer()
            ui.Triangle(
                style_type_name_override="CollapsableFrame.Header", width=width, height=height, alignment=alignment
            )
            ui.Spacer()
        ui.Label(text, style_type_name_override="CollapsableFrame.Header")

    def show_attribute_context_menu(b):
        if b != 1:
            return

        event = GroupHeaderContextMenuEvent(group_id=id, payload=[])
        GroupHeaderContextMenu.on_mouse_event(event)

    header_stack.set_mouse_pressed_fn(lambda x, y, b, _: show_attribute_context_menu(b))


class SimplePropertyWidget(PropertyWidget):
    """
    SimplePropertyWidget provides a simple vertical list of "Label" -> "Value widget" pair layout.
    """

    def __init__(self, title: str, collapsed: bool = False, collapsable: bool = True):
        super().__init__(title)
        self._collapsed_default = collapsed
        self._collapsed = collapsed
        self._collapsable = collapsable
        self._collapsable_frame = None
        self._payload = None
        self._pending_rebuild_task = None
        self._filter_changed_sub = None
        self.__style = get_style()

    def clean(self):
        """
        See PropertyWidget.clean
        """
        self.reset()

        if self._pending_rebuild_task is not None:
            self._pending_rebuild_task.cancel()
        self._pending_rebuild_task = None

        if self._collapsable_frame is not None:
            self._collapsable_frame.set_build_fn(None)
            if self._collapsable:
                self._collapsable_frame.set_collapsed_changed_fn(None)
            self._collapsable_frame = None

    def reset(self):
        """
        See PropertyWidget.reset
        """
        # Do not cancel rebuild task here as we are called from that task through derived build_items()
        pass

    def request_rebuild(self):
        """
        Request the widget to rebuild.
        It will be rebuilt on next frame.
        """
        if self._pending_rebuild_task:
            return
        self._pending_rebuild_task = run_coroutine(self._delayed_rebuild())

    def add_label(self, label: str):
        """
        Add a Label with a highlight text based on current filter
        """
        filter_text = self._filter.name
        HighlightLabel(label, highlight=filter_text, name="label", width=LABEL_WIDTH, height=LABEL_HEIGHT, style=self.__style)

    def add_item(self, label: str, value):
        """
        This function is supposed to be called inside of build_items function.
        Adds a "Label" -> "Value widget" pair item to the widget. "value" will be an uneditable string in a StringField.

        Args:
            label (str): The label text of the entry.
            value: The value to be stringified and displayed in a StringField.
        """
        if not self._filter.matches(label):
            return

        with ui.HStack(spacing=HORIZONTAL_SPACING):
            self.add_label(label)
            ui.StringField(name="models").model.set_value(str(value))
        self._any_item_visible = True

    def add_item_with_model(self, label: str, model, editable: bool = False, identifier: str=None):
        """
        This function is supposed to be called inside of build_items function.
        Adds a "Label" -> "Value widget with model" pair item to the widget. "value" will be an editable string in a
        StringField backed by supplied model.

        Args:
            label (str): The label text of the entry.
            model: The model to be used by the string field.
            editable: If the StringField generated from model should be editable. Default is False.
        """
        if not self._filter.matches(label):
            return

        with ui.HStack(spacing=HORIZONTAL_SPACING):
            self.add_label(label)
            sf = ui.StringField(name="models", model=model, enabled=editable)
            if identifier and sf:
                sf.identifier = identifier
        self._any_item_visible = True

    def on_new_payload(self, payload, ignore_large_selection=False) -> bool:
        """
        See PropertyWidget.on_new_payload
        """
        self._payload = payload

        if (
            not ignore_large_selection
            and payload
            and hasattr(payload, "is_large_selection")
            and payload.is_large_selection()
        ):
            return False

        return True

    def build_items(self):
        """
        When deriving from SimplePropertyWidget, override this function to build your items.
        """
        self.add_item("Hello", "World")

    def build_impl(self):
        """
        See PropertyWidget.build_impl
        """
        if self._filter_changed_sub is None:
            self._filter_changed_sub = self._filter.name_model.add_value_changed_fn(self._on_filter_changed)

        if self._collapsable:
            self._collapsable_frame = ui.CollapsableFrame(
                self._title, build_header_fn=self._build_frame_header, collapsed=self._collapsed_default
            )

            def on_collapsed_changed(collapsed):
                self._collapsed = collapsed

            self._collapsable_frame.set_collapsed_changed_fn(on_collapsed_changed)
            omni.kit.window.property.managed_frame.prep(self._collapsable_frame, type="Property")

        else:
            self._collapsable_frame = ui.Frame(height=0, style={"Frame": {"padding": 0}})
        # We cannot use build_fn because rebuild() won't trigger on invisible frames and toggling visible
        # first causes flicker during property search. So we build explicitly here.
        # request_rebuild() is still deferred to support batching of changes.
        self.request_rebuild()

    def _build_frame_header(self, collapsed, text: str, id: str = None):
        build_frame_header(collapsed, text, id)

    def _build_frame(self):
        import time

        start = time.time()
        with ui.VStack(height=0, spacing=5, name="frame_v_stack"):
            # ui.Spacer(height=0)
            self._any_item_visible = False
            self.build_items()
            if self._filter.name:
                self._collapsable_frame.visible = self._any_item_visible
            else:
                # for compatibility with widgets which don't support filtering
                self._collapsable_frame.visible = True
            ui.Spacer(height=0)
        took = time.time() - start
        if took > 0.2:
            import carb

            carb.log_warn(f"{self.__class__.__name__}.build_items took {took} seconds")

    @handle_exception
    async def _delayed_rebuild(self):
        if self._pending_rebuild_task is not None:
            try:
                if self._collapsable_frame:
                    with self._collapsable_frame:
                        self._build_frame()
                    self._collapsable_frame.rebuild()  # force rebuild of frame header too
            finally:
                self._pending_rebuild_task = None

    def _on_filter_changed(self, m: ui.AbstractValueModel):
        """
        Called when filter changes. Default calls request_rebuild().
        Derived classes can override to optimize by selectively changing property visibility.
        """
        self.request_rebuild()
