import carb
import omni
import omni.kit.commands
from pxr import Sdf, UsdGeom, Usd, UsdSkel
from omni.kit.usd_undo import UsdLayerUndo
from typing import List


class ApplySkelBindingAPICommand(omni.kit.commands.Command):
    def __init__(
            self,
            layer: Sdf.Layer = None,
            paths: List[Sdf.Path] = []
    ):
        self._usd_undo = None
        self._layer = layer
        self._paths = paths

    def do(self):
        stage = omni.usd.get_context().get_stage()
        if self._layer is None:
            self._layer = stage.GetEditTarget().GetLayer()
        self._usd_undo = UsdLayerUndo(self._layer)
        for path in self._paths:
            prim = stage.GetPrimAtPath(path)
            if not prim.HasAPI(UsdSkel.BindingAPI):
                self._usd_undo.reserve(path)
                UsdSkel.BindingAPI.Apply(prim)

    def undo(self):
        if self._usd_undo is not None:
            self._usd_undo.undo()


omni.kit.commands.register(ApplySkelBindingAPICommand)
