# Copyright (c) 2021, 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__ = ["LiveColumnDelegate"]

from omni.kit.widget.stage.abstract_stage_column_delegate import AbstractStageColumnDelegate, StageColumnItem
from omni.kit.widget.stage import StageItem, StageModel
from typing import List
from ..models import LiveModel
from ..icons import Icons
from enum import Enum

import omni.ui as ui


class LiveColumnSortPolicy(Enum):
    DEFAULT = 0
    LIVE_TO_OFFLINE = 1
    OFFLINE_TO_LIVE = 2


class LiveColumnDelegate(AbstractStageColumnDelegate):
    """The column delegate that represents the Live column"""

    def __init__(self):
        super().__init__()

        self.__stage_item_models = {}
        from omni.ui import color as cl

        ITEM_DARK = 0x7723211F
        ICON_DARK = 0x338A8777
        LIVE_GREEN = cl("#76b800")
        LIVE_SEL = cl("#90e203")

        self.WIDGET_STYLE = {
            "Button.Image::live": {"image_url": Icons().get("lightning"), "color": ICON_DARK},
            "Button.Image::live:checked": {"color": LIVE_GREEN},
            "Button.Image::live:checked:hovered": {"color": LIVE_SEL},
            "Button.Image::live:disabled": {"color": 0x608A8777},
            "Button.Image::live:selected": {"color": ITEM_DARK},
            "Button::live": {"background_color": 0x0, "margin": 0, "margin_width": 1},
            "Button::live:checked": {"background_color": 0x0},
            "Button::live:hovered": {"background_color": 0x0},
            "Button::live:pressed": {"background_color": 0x0},
            "TreeView.Item.Live": {"color": LIVE_GREEN},
            "TreeView.Item.Live:hovered": {"color": LIVE_SEL},
            "TreeView.Item.Live:selected": {"color": LIVE_SEL},
            # Top of the column section (that has sorting options)
            "TreeView.Header::live_header": {"image_url": Icons().get("lightning"), "color": 0xFF888888},
        }

        self.__stage_model: StageModel = None
        self.__items_sort_policy = LiveColumnSortPolicy.DEFAULT
        self.__live_header_layout = None

    def destroy(self):
        for model in self.__stage_item_models.values():
            model.destroy()
        self.__stage_item_models = {}
        self.__stage_model = None
        if self.__live_header_layout:
            self.__live_header_layout.set_mouse_pressed_fn(None)
            self.__live_header_layout = None

    @property
    def initial_width(self):
        """The width of the column"""
        return ui.Pixel(24)

    @property
    def minimum_width(self):
        return ui.Pixel(24)

    def __on_policy_changed(self):
        stage_model = self.__stage_model
        if not stage_model:
            return

        if self.__items_sort_policy == LiveColumnSortPolicy.LIVE_TO_OFFLINE:
            stage_model.set_items_sort_key_func(
                lambda item: item.in_session if item.payrefs else len(item.payrefs) == 0
            )
        elif self.__items_sort_policy == LiveColumnSortPolicy.OFFLINE_TO_LIVE:
            stage_model.set_items_sort_key_func(
                lambda item: item.in_session if item.payrefs else len(item.payrefs) == 0,
                reverse=True
            )
        else:
            stage_model.set_items_sort_key_func(None)

    def __on_header_clicked(self, x, y, b, m):
        if b != 0 or not self.__stage_model:
            return

        if self.__items_sort_policy == LiveColumnSortPolicy.LIVE_TO_OFFLINE:
            self.__items_sort_policy = LiveColumnSortPolicy.OFFLINE_TO_LIVE
        elif self.__items_sort_policy == LiveColumnSortPolicy.OFFLINE_TO_LIVE:
            self.__items_sort_policy = LiveColumnSortPolicy.DEFAULT
        else:
            self.__items_sort_policy = LiveColumnSortPolicy.LIVE_TO_OFFLINE

        self.__on_policy_changed()

    def build_header(self, **kwargs):
        """Build the header"""

        self.__stage_model = kwargs.get("stage_model", None)
        with ui.ZStack(style=self.WIDGET_STYLE):
            tooltip = "LMB:Sort Live"
            ui.Rectangle(name="hovering", style_type_name_override="TreeView.Header", tooltip=tooltip)
            self.__live_header_layout = ui.HStack()
            with self.__live_header_layout:
                ui.Spacer()
                with ui.VStack(width=0):
                    ui.Spacer()
                    ui.Image(width=22, height=14, name="live_header", style_type_name_override="TreeView.Header")
                    ui.Spacer()
                ui.Spacer()

            self.__live_header_layout.set_tooltip(tooltip)
            self.__live_header_layout.set_mouse_pressed_fn(self.__on_header_clicked)

    async def build_widget(self, item: StageColumnItem, **kwargs):
        """Build the live widget"""

        stage_item = kwargs.get("stage_item", None)
        reloadable = (stage_item.payloads or stage_item.references)
        if not stage_item or not stage_item.prim or not reloadable:
            return

        with ui.ZStack(height=20, style=self.WIDGET_STYLE):
            # Min size
            ui.Spacer(width=22)

            live_model = self.__stage_item_models.get(stage_item, None)
            if not live_model:
                live_model = LiveModel(stage_item)
                self.__stage_item_models[stage_item] = live_model

            tooltip = "Live session status"
            ui.ToolButton(live_model, enabled=True, name="live", tooltip=tooltip)

    def on_stage_items_destroyed(self, items: List[StageItem]):
        if not items:
            return

        for item in items:
            live_model = self.__stage_item_models.pop(item, None)
            if live_model:
                live_model.destroy()

    @property
    def order(self):
        return -103

    @property
    def sortable(self):
        return True
