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

import asyncio

import carb.settings
from omni.kit.async_engine import run_coroutine
import omni.kit.window.property as p
import omni.usd as ou
import omni.ui as ui
from pxr import Usd, Sdf, Tf

from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidget

from omni.kit.window.property.templates import LABEL_WIDTH, LABEL_HEIGHT


class ComboListItem(ui.AbstractItem):
    def __init__(self, item):
        super().__init__()
        self.model = ui.SimpleStringModel(item)


class ComboListModel(ui.AbstractItemModel):
    def __init__(self, item_list, default_index):
        super().__init__()
        self._default_index = default_index
        self._current_index = ui.SimpleIntModel(default_index)
        self._current_index.add_value_changed_fn(lambda a: self._item_changed(None))
        self._item_list = item_list
        self._items = []
        if item_list:
            for item in item_list:
                self._items.append(ComboListItem(item))

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

    def get_item_list(self):
        return self._item_list

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

    def get_current_index(self):
        return self._current_index.get_value_as_int()

    def set_current_index(self, index):
        self._current_index.set_value(index)

    def get_current_string(self):
        return self._items[self._current_index.get_value_as_int()].model.get_value_as_string()

    def is_default(self):
        return self.get_current_index() == self._default_index

class IndexPropertiesBaseWidget(UsdPropertiesWidget):

    USD_RENDERSETTINGS = 'nvindex.renderSettings'

    def __init__(self, title, element_types):
        super().__init__(title, collapsed=False, multi_edit=False)
        if not isinstance(element_types, list):
            element_types = [element_types]
        self.element_types = element_types
        w = p.get_window()
        w.register_widget("prim", self.element_types[0], self)
        settings = carb.settings.get_settings()
        self.show_advanced_settings = settings.get("/nvindex/showAdvancedSettings")

    def clean(self):
        self.reset()
        super().clean()

    def reset(self):
        super().reset()

    def on_shutdown(self):
        w = p.get_window()
        w.unregister_widget("prim", self.element_types[0])

    def on_new_payload(self, payload):
        if not super().on_new_payload(payload):
            return False

        if len(self._payload) == 0:
            return False

        for path in self._payload:
            prim = self._get_prim(path)
            if prim is None:
                return False
            if prim.GetTypeName() in self.element_types:
                return True

        return False

    def get_prim(self):
        return ou.get_context().get_stage().GetPrimAtPath(self._payload[-1])

    # def get_script(self):
    #     stage = ou.get_context().get_stage()
    #     return stage.GetPrimAtPath(self._payload[-1]).GetAttribute("inputs:script").Get() or ""

    def set_render_setting(self, key, value, default):
        prim = self.get_prim()

        #render_settings = prim.GetCustomDataByKey(self.USD_RENDERSETTINGS)
        path = self.USD_RENDERSETTINGS + ":" + key

        if value == default:
            prim.ClearCustomDataByKey(path )
            # if render_settings and key in render_settings:
            #     del render_settings[key]
        else:
            prim.SetCustomDataByKey(path, value)
            # if render_settings == None:
            #     render_settings = {}
            # render_settings[key] = value

        # if render_settings:
        #     prim.SetCustomDataByKey(self.USD_RENDERSETTINGS, render_settings)
        # else:
        #     prim.ClearCustomDataByKey(self.USD_RENDERSETTINGS)

    def set_volume_type(self, combo_model, item):
        prim = self.get_prim()
        volume_type = combo_model.get_current_string()
        if volume_type ==  "default":
            prim.GetAttribute("nvindex:type").Clear()
        else:
            prim.CreateAttribute("nvindex:type", Sdf.ValueTypeNames.Token, True, Sdf.VariabilityUniform).Set(volume_type)

    def _on_domain_changed(self, *args):
        pass
        # self.target_attrib_name_model._set_dirty()

    def create_combo_list_model(self, current_value, items_list):
        index = 0
        if current_value in items_list:
            index = items_list.index(current_value)
        return ComboListModel(items_list, index)

    def build_label(self, caption, full_width=False):
        if full_width:
            ui.Label(caption, name="label", word_wrap=True, height=LABEL_HEIGHT)
        else:
            ui.Label(caption, name="label", word_wrap=True, width=LABEL_WIDTH, height=LABEL_HEIGHT)

    def store_value(self, key, value, default = '', valueType = Sdf.ValueTypeNames.String, custom = False):
        prim = self.get_prim()
        if value == default:
            prim.GetAttribute(key).Clear()
        else:
            prim.CreateAttribute(key, valueType, custom, Sdf.VariabilityUniform).Set(value)

    def get_render_setting_value(self, key, default, render_settings):
        if render_settings and key in render_settings:
            return render_settings[key]
        else:
            return default

    def build_render_setting_combo_box(self, caption, key, items_list, render_settings):
        with ui.HStack():
            self.build_label(caption)
            current_value = self.get_render_setting_value(key, "default", render_settings)

            index = 0
            if current_value in items_list:
                index = items_list.index(current_value)

            widget = ui.ComboBox(self.create_combo_list_model(current_value, items_list), name=key)
            widget.model.add_item_changed_fn(lambda m, i: self.set_render_setting(key, m.get_current_string(), "default"))

    def build_render_setting_float(self, caption, key, default, render_settings):
        with ui.HStack():
            self.build_label(caption)
            current_value = self.get_render_setting_value(key, default, render_settings)

            widget = ui.FloatField(name=key)
            widget.model.set_value(current_value)
            widget.model.add_value_changed_fn(lambda m: self.set_render_setting(key, m.get_value_as_float(), default))


    def build_render_setting_int(self, caption, key, default, render_settings):
        with ui.HStack():
            self.build_label(caption)
            current_value = self.get_render_setting_value(key, default, render_settings)

            widget = ui.IntField(name=key)
            widget.model.set_value(current_value)
            widget.model.add_value_changed_fn(lambda m: self.set_render_setting(key, m.get_value_as_int(), default))

    def build_render_setting_bool(self, caption, key, default, render_settings):
        with ui.HStack():
            self.build_label(caption)
            current_value = self.get_render_setting_value(key, default, render_settings)

            widget = ui.CheckBox(name=key)
            widget.model.set_value(current_value)
            widget.model.add_value_changed_fn(lambda m: self.set_render_setting(key, m.get_value_as_bool(), default))

    # Copied from UsdPropertiesWidget with minor changes (force rebuild on
    # non-property change, to handle modified custom data)
    def _on_usd_changed(self, notice, stage):
        if stage != self._payload.get_stage():
            return

        if not self._collapsable_frame:
            return

        if len(self._payload) == 0:
            return

        if self._pending_rebuild_task is not None:
            return

        dirty_paths = set()

        for path in notice.GetResyncedPaths():
            if path in self._payload:
                self.request_rebuild()
                return
            elif path.GetPrimPath() in self._payload:
                prop = stage.GetPropertyAtPath(path)
                if (not prop.IsValid()) != (self._models.get(path) is None):
                    self.request_rebuild()
                    return
                else:
                    dirty_paths.add(path)

        for path in notice.GetChangedInfoOnlyPaths():
            dirty_paths.add(path)
            if path in self._payload and not path.IsPropertyPath():
                # Handle changed custom data
                self.request_rebuild()
                return

        self._pending_dirty_paths.update(dirty_paths)

        if self._pending_dirty_task_or_future is None:
            self._pending_dirty_task_or_future = run_coroutine(self._delayed_dirty_handler())

    def build_items(self):
        self.reset()

        if len(self._payload) == 0:
            return

        prim_path = self._payload[-1]
        last_prim = self._get_prim(prim_path)
        if not last_prim:
            return

        stage = last_prim.GetStage()
        if not stage:
            return

        self._listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, self._on_usd_changed, stage)

        self.build_properties(stage, last_prim, prim_path)
