import carb
import omni.usd
from pxr import Sdf, UsdShade
from omni.kit.window.property.property_scheme_delegate import PropertySchemeDelegate


class MaterialPrimSchemeDelegate(PropertySchemeDelegate):
    def get_widgets(self, payload):
        widgets_to_build = []
        if self._should_enable_delegate(payload):
            widgets_to_build.append("path")
            widgets_to_build.append("material")
        return widgets_to_build

    def get_unwanted_widgets(self, payload):
        unwanted_widgets_to_build = []
        if self._should_enable_delegate(payload):
            unwanted_widgets_to_build.append("transform")
        return unwanted_widgets_to_build

    def _should_enable_delegate(self, payload):
        stage = payload.get_stage()
        if stage:
            for prim_path in payload:
                prim = stage.GetPrimAtPath(prim_path)
                if prim:
                    if not prim.IsA(UsdShade.Material):
                        return False
            return True
        return False


class ShaderPrimSchemeDelegate(PropertySchemeDelegate):
    def get_widgets(self, payload):
        widgets_to_build = []
        if self._should_enable_delegate(payload):
            widgets_to_build.append("path")
            widgets_to_build.append("shader")

        return widgets_to_build

    def get_unwanted_widgets(self, payload):
        unwanted_widgets_to_build = []
        if self._should_enable_delegate(payload):
            unwanted_widgets_to_build.append("transform")
        return unwanted_widgets_to_build

    def _should_enable_delegate(self, payload):
        stage = payload.get_stage()
        if stage:
            for prim_path in payload:
                prim = stage.GetPrimAtPath(prim_path)
                if prim:
                    if not prim.IsA(UsdShade.Shader):
                        return False
            return True
        return False
