# 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 omni.ui as ui
import omni.usd
from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidget, UsdPropertyUiEntry
from omni.kit.property.usd.prim_selection_payload import PrimSelectionPayload
from pxr import Usd, UsdSkel, UsdGeom
from omni.kit.property.usd.custom_layout_helper import CustomLayoutFrame, CustomLayoutGroup, CustomLayoutProperty
from .command import *

MESH_ATTRS = ["skel:skeleton", "skel:skinningMethod", "skel:blendShapeTargets"]


class MeshSkelBindingAttributesWidget(UsdPropertiesWidget):
    def __init__(self, title: str):
        super().__init__(title, collapsed=False)
        self._title = title

    def destroy(self):
        pass

    def on_new_payload(self, payload):
        if not super().on_new_payload(payload):
            return False
        if not self._payload or len(self._payload) == 0:
            return False
        for prim_path in payload:
            prim = self._get_prim(prim_path)
            if not prim:
                return False
            if prim.HasAPI(UsdSkel.BindingAPI) and prim.IsA(UsdGeom.Mesh):
                return True
        return False

    def _filter_props_to_build(self, props):
        return [prop for prop in props if prop.GetName() in MESH_ATTRS]

    def _customize_props_layout(self, props):
        frame = CustomLayoutFrame(hide_extra=True)

        def spacer_build_fn(*args):
            ui.Spacer(height=10)

        with frame:
            CustomLayoutProperty("skel:skeleton", "Skeleton")
            CustomLayoutProperty(None, None, build_fn=spacer_build_fn)
            CustomLayoutProperty("skel:skinningMethod", "Skinning Method")
            CustomLayoutProperty(None, None, build_fn=spacer_build_fn)
            CustomLayoutProperty("skel:blendShapeTargets", "Blendshape Targets")

        return frame.apply(props)

    def get_additional_kwargs(self, ui_attr):
        random_prim = Usd.Prim()
        prim_same_type = True
        if self._payload:
            prim_paths = self._payload.get_paths()
            for prim_path in prim_paths:
                prim = self._get_prim(prim_path)
                if not prim:
                    continue
                if not random_prim:
                    random_prim = prim
                elif random_prim.GetTypeName() != prim.GetTypeName():
                    prim_same_type = False
                    break

        if ui_attr.prop_name == "skel:skeleton" and prim_same_type and random_prim and random_prim.IsA(UsdGeom.Mesh):
            return None, {"target_picker_filter_type_list": [UsdSkel.Skeleton], "targets_limit": 1}
        elif ui_attr.prop_name == "skel:blendShapeTargets" and prim_same_type and random_prim and random_prim.IsA(UsdGeom.Mesh):
            return None, {"target_picker_filter_type_list": [UsdSkel.BlendShape], "targets_limit": 0}
        return None, {"targets_limit": 0}
