# 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
from omni.kit.property.usd.prim_selection_payload import PrimSelectionPayload
from pxr import Sdf, Usd, UsdSkel
from omni.kit.property.usd.custom_layout_helper import CustomLayoutFrame, CustomLayoutProperty
from .command import *


class SkelAnimationAttributesWidget(UsdPropertiesWidget):
    def __init__(self, title: str):
        super().__init__(title, collapsed=False)
        self._title = title
        from omni.kit.property.usd import PrimPathWidget
        self._add_button_menus = []
        self._add_button_menus.append(
            PrimPathWidget.add_button_menu_entry(
                "Animation/Skeletal Animation",
                show_fn=self.on_show,
                onclick_fn=self.on_click
            )
        )

    def destroy(self):
        for menu_entry in self._add_button_menus:
            PrimPathWidget.remove_button_menu_entry(menu_entry)

    def on_show(self, objects: dict):
        if "prim_list" not in objects or "stage" not in objects:
            return False
        stage = objects["stage"]
        if not stage:
            return False
        prim_list = objects["prim_list"]
        if len(prim_list) < 1:
            return False
        for item in prim_list:
            if isinstance(item, Sdf.Path):
                prim = stage.GetPrimAtPath(item)
            elif isinstance(item, Usd.Prim):
                prim = item
            if not prim.HasAPI(UsdSkel.BindingAPI) and prim.IsA(UsdSkel.Skeleton):
                return True
        return False

    def on_click(self, payload: PrimSelectionPayload):
        if payload is None:
            return Sdf.Path.emptyPath
        payload_paths = payload.get_paths()
        omni.kit.commands.execute("ApplySkelBindingAPICommand", paths=payload_paths)

    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(UsdSkel.Skeleton) or prim.IsA(UsdSkel.Root)):
                return True
        return False

    def _filter_props_to_build(self, props):
        return [prop for prop in props if prop.GetName() == "skel:animationSource"]

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

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

        with frame:
            CustomLayoutProperty("skel:animationSource", "Animation Source")
            CustomLayoutProperty(None, None, build_fn=spacer_build_fn)

        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:animationSource" and prim_same_type and random_prim and (random_prim.IsA(UsdSkel.Skeleton) or random_prim.IsA(UsdSkel.Root)):
            return None, {"target_picker_filter_type_list": [UsdSkel.Animation], "targets_limit": 1}
        return None, {"targets_limit": 0}
