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

from typing import List
import omni.kit.commands
import omni.ui as ui
from .usd_model_base import UsdBase
from pxr import Sdf, Usd


class VariantSetModel(ui.AbstractItemModel, UsdBase):
    def __init__(self, stage: Usd.Stage,  object_paths: List[Sdf.Path], variant_set_name: str, self_refresh: bool):
        UsdBase.__init__(self, stage, object_paths, self_refresh, {})
        ui.AbstractItemModel.__init__(self)

        self._variant_set_name = variant_set_name
        self._variant_names = []

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

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

        return item.model

    def begin_edit(self, item):
        UsdBase.begin_edit(self)

    def end_edit(self, item):
        UsdBase.end_edit(self)

    def set_value(self, value, comp=-1):
        # if one than one item selected then don't early exit as
        # anchor might be same but others may not
        if value == self._value and len(self._variant_names) == 1:
            return

        omni.kit.commands.execute(
            "SelectVariantPrim",
            prim_path=self._object_paths,
            vset_name=self._variant_set_name,
            var_name=value,
        )

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

        index = model.as_int
        if self.set_value(self._variant_names[index].model.get_value_as_string()):
            self._item_changed(None)

    def _update_variant_names(self):
        self._variant_names = []

        vset = self._get_variant_set()
        if vset:
            vnames = vset.GetVariantNames()

            class AllowedTokenItem(ui.AbstractItem):
                def __init__(self, name):
                    super().__init__()
                    self.model = ui.SimpleStringModel(name)

            self._variant_names.append(AllowedTokenItem(""))  # Empty variant
            for name in vnames:
                self._variant_names.append(AllowedTokenItem(name))

    def _update_value(self, force=False):
        if self._update_value_objects(force, self._get_objects()):
            # TODO don't have to do this every time. Just needed when "VariantNames" actually changed
            self._update_variant_names()

            index = -1
            for i in range(0, len(self._variant_names)):
                if self._variant_names[i].model.get_value_as_string() == self._value:
                    index = i

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

    def _on_dirty(self):
        self._item_changed(None)

    def _read_value(self, object: Usd.Object, time_code: Usd.TimeCode):
        vsets = object.GetVariantSets()
        vset = vsets.GetVariantSet(self._variant_set_name)
        return vset.GetVariantSelection()

    def _get_variant_set(self):
        prims = self._get_objects()
        prim = prims[0] if len(prims) > 0 else None
        if prim:
            vsets = prim.GetVariantSets()
            return vsets.GetVariantSet(self._variant_set_name)
        return None


# Temporary commands. Expected to be moved when we have better support for variants
class SelectVariantPrimCommand(omni.kit.commands.Command):
    def __init__(self, prim_path: str, vset_name: str, var_name: str, usd_context_name: str = ""):
        self._usd_context = omni.usd.get_context(usd_context_name)
        self._prim_path = prim_path if isinstance(prim_path, list) else [prim_path]
        self._vset_name = vset_name
        self._var_name = var_name

    def do(self):
        stage = self._usd_context.get_stage()
        self._previous_selection = {}
        for prim_path in self._prim_path:
            prim = stage.GetPrimAtPath(prim_path)
            vset = prim.GetVariantSets().GetVariantSet(self._vset_name)

            self._previous_selection[prim_path] = vset.GetVariantSelection()
            if self._var_name:
                vset.SetVariantSelection(self._var_name)
            else:
                vset.ClearVariantSelection()

    def undo(self):
        stage = self._usd_context.get_stage()
        for prim_path in self._prim_path:
            prim = stage.GetPrimAtPath(prim_path)
            vset = prim.GetVariantSets().GetVariantSet(self._vset_name)

            if prim_path in self._previous_selection:
                vset.SetVariantSelection(self._previous_selection[prim_path])
            else:
                vset.ClearVariantSelection()

        self._previous_selection = {}


omni.kit.commands.register_all_commands_in_module(__name__)
