# 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__ = ["ReloadColumnDelegate"]

from omni.kit.widget.stage.abstract_stage_column_delegate import AbstractStageColumnDelegate, StageColumnItem
from omni.kit.widget.stage import StageItem
from typing import List
from ..models.reload_model import ReloadModel
from ..icons import Icons
from functools import partial
from omni.ui import color as cl
from enum import Enum

import weakref
import omni.ui as ui
import omni.kit.app
import omni.kit.usd.layers as layers
import carb


class ReloadColumnSortPolicy(Enum):
    DEFAULT = 0
    OUTDATE_TO_LATEST = 1
    LATEST_TO_OUTDATE = 2


class ReloadColumnDelegate(AbstractStageColumnDelegate):
    """The column delegate that represents the Auto Reload column"""

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

        self.__reload_layout = None
        self.__stage_item_models = {}

        ITEM_DARK = 0x7723211F
        ICON_DARK = 0x338A8777

        RELOAD = cl("#eb9d00")
        RELOAD_H = cl("#ffaa00")
        AUTO_RELOAD = cl("#34C7FF")
        AUTO_RELOAD_H = cl("#82dcff")

        self.WIDGET_STYLE = {
            "Button.Image::reload": {"image_url": Icons().get("reload"), "color": ICON_DARK},
            "Button.Image::reload:checked": {"color": RELOAD_H},
            "Button.Image::reload:disabled": {"color": 0x608A8777},
            "Button.Image::reload:selected": {"color": ITEM_DARK},
            "Button.Image::reload:checked:hovered": {"color": 0xFFFFFFFF},
            "Button.Image::auto_reload": {"image_url": Icons().get("reload"), "color": AUTO_RELOAD},
            "Button.Image::auto_reload:checked": {"color": AUTO_RELOAD},
            "Button.Image::auto_reload:disabled": {"color": AUTO_RELOAD},
            "Button.Image::auto_reload:checked:hovered": {"color": AUTO_RELOAD_H},
            "Button::reload": {"background_color": 0x0, "margin": 0, "margin_width": 1},
            "Button::reload:checked": {"background_color": 0x0},
            "Button::reload:hovered": {"background_color": 0x0, "color": 0xFFFFFFFF},
            "Button::reload:pressed": {"background_color": 0x0},
            "Button::auto_reload": {"background_color": 0x0, "margin": 0, "margin_width": 1},
            "Button::auto_reload:checked": {"background_color": 0x0},
            "Button::auto_reload:hovered": {"background_color": 0x0, "color": 0xFFFFFFFF},
            "Button::auto_reload:pressed": {"background_color": 0x0},
            "TreeView.Item.Outdated": {"color": RELOAD},
            "TreeView.Item.Outdated:hovered": {"color": RELOAD_H},
            "TreeView.Item.Outdated:selected": {"color": RELOAD_H},
            # Top of the column section (that has sorting options)
            "TreeView.Header::reload_header": {"image_url": Icons().get("reload_dark")},
        }

        self.__stage_model = None
        self.__items_sort_policy = ReloadColumnSortPolicy.DEFAULT

    def destroy(self):
        if self.__reload_layout:
            self.__reload_layout.set_mouse_pressed_fn(None)
            self.__reload_layout = None

        for model in self.__stage_item_models.values():
            model.destroy()
        self.__stage_item_models = {}
        self.__stage_model = 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 == ReloadColumnSortPolicy.OUTDATE_TO_LATEST:
            stage_model.set_items_sort_key_func(
                lambda item: item.is_outdated if item.payrefs else len(item.payrefs) == 0
            )
        elif self.__items_sort_policy == ReloadColumnSortPolicy.LATEST_TO_OUTDATE:
            stage_model.set_items_sort_key_func(
                lambda item: item.is_outdated if item.payrefs else len(item.payrefs) == 0,
                reverse=True
            )
        else:
            stage_model.set_items_sort_key_func(None)

    def __on_reload_clicked(self, weakref_stage_model, x, y, b, m):
        stage_model = weakref_stage_model()
        if not stage_model or not stage_model.stage:
            return

        if b == 1:
            usd_context = omni.usd.get_context_from_stage(stage_model.stage)
            if not usd_context:
                return

            layers_state = layers.get_layers_state(usd_context)

            from omni.kit.widget.live_session_management.utils import reload_outdated_layers

            ref_ids = layers_state.get_outdated_non_sublayer_identifiers()
            reload_outdated_layers(ref_ids, usd_context)
        elif b == 0:
            if self.__items_sort_policy == ReloadColumnSortPolicy.OUTDATE_TO_LATEST:
                self.__items_sort_policy = ReloadColumnSortPolicy.LATEST_TO_OUTDATE
            elif self.__items_sort_policy == ReloadColumnSortPolicy.LATEST_TO_OUTDATE:
                self.__items_sort_policy = ReloadColumnSortPolicy.DEFAULT
            else:
                self.__items_sort_policy = ReloadColumnSortPolicy.OUTDATE_TO_LATEST

            self.__on_policy_changed()

    def build_header(self, **kwargs):
        """Build the header"""
        stage_model = kwargs.get("stage_model", None)
        self.__stage_model = stage_model
        if stage_model:
            with ui.ZStack(style=self.WIDGET_STYLE):
                tooltip = "RMB: Reload All References/Payloads"
                ui.Rectangle(name="hovering", style_type_name_override="TreeView.Header", tooltip=tooltip)
                self.__reload_layout = ui.HStack()
                with self.__reload_layout:
                    ui.Spacer()
                    with ui.VStack(width=0):
                        ui.Spacer()
                        ui.Image(width=22, height=14, name="reload_header", style_type_name_override="TreeView.Header")
                        ui.Spacer()
                    ui.Spacer()

                weakref_stage_model = weakref.ref(stage_model)
                self.__reload_layout.set_mouse_pressed_fn(
                    partial(self.__on_reload_clicked, weakref_stage_model)
                )
        else:
            with ui.HStack(style=self.WIDGET_STYLE):
                ui.Spacer()
                with ui.VStack(width=0):
                    ui.Spacer()
                    ui.Image(width=22, height=14, name="reload_header", style_type_name_override="TreeView.Header")
                    ui.Spacer()
                ui.Spacer()

    async def build_widget(self, item: StageColumnItem, **kwargs):
        """Build the reload 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

        g_auto_reload = carb.settings.get_settings().get_as_bool(layers.SETTINGS_AUTO_RELOAD_NON_SUBLAYERS)

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

            reload_model = self.__stage_item_models.get(stage_item, None)
            if not reload_model:
                reload_model = ReloadModel(stage_item)
                self.__stage_item_models[stage_item] = reload_model

            button = ui.ToolButton(reload_model)

            stub = "payloads" if stage_item.payloads else "references"
            if not reload_model.get_value_as_bool():
                tooltip = "All {} up to date".format(stub)
            else:
                tooltip = "Reload {} on prim".format(stub)

            if stage_item.auto_reload:
                tooltip = "Auto Reload All Enabled" if g_auto_reload else "Auto Reload Enabled"
                button.enabled = False

            if stage_item.is_outdated and g_auto_reload and stage_item.in_session:
                tooltip = "Auto Reload Blocked by Live Session"

            button.tooltip = tooltip
            button.name = "auto_reload" if stage_item.auto_reload else "reload"

    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 -104

    @property
    def sortable(self):
        return True
