# 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.
#
import carb
import omni.kit.commands
import omni.timeline
import omni.ui as ui
import omni.usd
from .usd_model_base import UsdBase
from typing import List
from pxr import Usd, Sdf


class MetadataObjectModel(ui.AbstractItemModel, UsdBase):
    """The value model that is reimplemented in Python to watch a USD paths.
       Paths can be either Attribute or Prim paths"""

    def __init__(
        self,
        stage: Usd.Stage,
        object_paths: List[Sdf.Path],
        self_refresh: bool,
        metadata: dict,
        key: str,
        default,
        options: list,
    ):
        UsdBase.__init__(self, stage, object_paths, self_refresh, metadata)
        ui.AbstractItemModel.__init__(self)

        self._key = key
        self._default_value = default
        self._combobox_options = options
        self._update_option()
        self._current_index = ui.SimpleIntModel()
        self._current_index.add_value_changed_fn(self._current_index_changed)

        self._has_index = False
        self._update_value()
        self._has_index = True

    def clean(self):
        UsdBase.clean(self)

    def get_item_children(self, item):
        self._update_value()
        return self._options

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

    def begin_edit(self):
        carb.log_warn("begin_edit not supported in MetadataObjectModel")

    def end_edit(self):
        carb.log_warn("end_edit not supported in MetadataObjectModel")

    def _current_index_changed(self, model):
        if not self._has_index:
            return

        index = model.as_int
        if self.set_value(self._options[index].value):
            self._item_changed(None)

    def _update_option(self):
        class OptionItem(ui.AbstractItem):
            def __init__(self, display_name: str, value: int):
                super().__init__()
                self.model = ui.SimpleStringModel(display_name)
                self.value = value

        self._options = []
        for index, option in enumerate(self._combobox_options):
            self._options.append(OptionItem(option, int(index)))

    def _update_value(self, force=False):
        if self._update_value_objects(force, self._get_objects()):
            index = -1
            for i in range(0, len(self._options)):
                if self._options[i].model.get_value_as_string() == self._value:
                    index = i

            if index != -1 and self._current_index.as_int != index:
                self._has_index = False
                self._current_index.set_value(index)
                self._item_changed(None)
                self._has_index = True

    def _on_dirty(self):
        self._update_value()

    def set_value(self, value, comp=-1):
        if comp != -1:
            carb.log_warn("Arrays not supported in MetadataObjectModel")
            self._update_value(True)  # reset value
            return False

        if not self._ambiguous and not any(self._comp_ambiguous) and value == self._value:
            return False

        self._value = self._combobox_options[int(value)]
        objects = self._get_objects()
        if len(objects) == 0:
            return False

        self._write_value(objects=objects, key=self._key, value=self._value)

        # We just set all the properties to the same value, it's no longer ambiguous
        self._ambiguous = False
        self._comp_ambiguous.clear()
        self.update_control_state()

        return True

    def _read_value(self, object: Usd.Object, time_code: Usd.TimeCode):
        value = object.GetMetadata(self._key)
        if not value:
            value = self._default_value
        return value

    def _write_value(self, objects: list, key: str, value):
        path_list = []
        for object in objects:
            path_list.append(object.GetPath())

        omni.kit.commands.execute("ChangeMetadata", object_paths=path_list, key=key, value=value)
