import carb
import omni.usd
from pxr import Sdf, UsdMedia, OmniAudioSchema, UsdGeom, UsdLux, UsdSkel
from omni.kit.window.property.property_scheme_delegate import PropertySchemeDelegate


class GeomPrimSchemeDelegate(PropertySchemeDelegate):
    def get_widgets(self, payload):
        widgets_to_build = []
        anchor_prim = None

        stage = payload.get_stage()
        if stage:
            for prim_path in payload:
                prim = stage.GetPrimAtPath(prim_path)
                if prim:
                    if not prim.IsA(UsdGeom.Imageable):
                        return widgets_to_build
                    anchor_prim = prim

        if anchor_prim is None:
            return widgets_to_build

        if (
            anchor_prim
            and anchor_prim.IsA(UsdMedia.SpatialAudio)
            or anchor_prim.IsA(OmniAudioSchema.OmniSound)
            or anchor_prim.IsA(OmniAudioSchema.OmniListener)
        ):
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("media")
            widgets_to_build.append("audio_sound")
            widgets_to_build.append("audio_listener")
            widgets_to_build.append("audio_settings")
            widgets_to_build.append("kind")

        elif anchor_prim and anchor_prim.IsA(UsdSkel.Root): # pragma: no cover
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("character_api")
            widgets_to_build.append("anim_graph_api")
            widgets_to_build.append("omni_graph_api")
            widgets_to_build.append("kind")

        elif anchor_prim.IsA(UsdSkel.Skeleton): # pragma: no cover
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("control_rig")
            widgets_to_build.append("skel_animation")

        elif anchor_prim and anchor_prim.IsA(UsdGeom.Camera):
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("camera")
            widgets_to_build.append("geometry_imageable")
            widgets_to_build.append("kind")

        # https://github.com/PixarAnimationStudios/USD/commit/7540fdf3b2aa6b6faa0fce8e7b4c72b756286f51
        elif anchor_prim and ((hasattr(UsdLux, 'LightAPI') and anchor_prim.HasAPI(UsdLux.LightAPI)) or ((hasattr(UsdLux, 'Light') and anchor_prim.IsA(UsdLux.Light)))):
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("light")
            widgets_to_build.append("geometry_imageable")
            widgets_to_build.append("kind")

        elif anchor_prim and anchor_prim.IsA(UsdGeom.Scope):
            widgets_to_build.append("path")
            widgets_to_build.append("geometry_imageable")
            widgets_to_build.append("character_motion_library")
            widgets_to_build.append("kind")

        elif anchor_prim and anchor_prim.IsA(UsdGeom.Mesh):
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("material_binding")
            widgets_to_build.append("geometry")
            widgets_to_build.append("geometry_imageable")
            widgets_to_build.append("character_interactive_object_api")
            widgets_to_build.append("mesh_skel_binding")
            widgets_to_build.append("kind")

        elif anchor_prim and anchor_prim.IsA(UsdGeom.Xform):
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("material_binding")
            widgets_to_build.append("geometry_imageable")
            widgets_to_build.append("kind")

        elif anchor_prim and anchor_prim.IsA(UsdGeom.Xformable):
            widgets_to_build.append("path")
            widgets_to_build.append("transform")
            widgets_to_build.append("material_binding")
            widgets_to_build.append("geometry")
            widgets_to_build.append("geometry_imageable")
            widgets_to_build.append("character_interactive_object_api")
            widgets_to_build.append("kind")

        return widgets_to_build

    def get_unwanted_widgets(self, payload):
        unwanted_widgets_to_build = []
        anchor_prim = None

        stage = payload.get_stage()
        if stage:
            for prim_path in payload:
                prim = stage.GetPrimAtPath(prim_path)
                if prim:
                    if not prim.IsA(UsdGeom.Imageable):
                        return unwanted_widgets_to_build
                    anchor_prim = prim

        if anchor_prim is None:
            return unwanted_widgets_to_build

        if anchor_prim and anchor_prim.IsA(UsdGeom.Scope):
            unwanted_widgets_to_build.append("transform")
            unwanted_widgets_to_build.append("geometry")
            unwanted_widgets_to_build.append("kind")

        # https://github.com/PixarAnimationStudios/USD/commit/7540fdf3b2aa6b6faa0fce8e7b4c72b756286f51
        elif anchor_prim and ((hasattr(UsdLux, 'LightAPI') and anchor_prim.HasAPI(UsdLux.LightAPI)) or ((hasattr(UsdLux, 'Light') and anchor_prim.IsA(UsdLux.Light)))):
            unwanted_widgets_to_build.append("geometry")

        elif (
            anchor_prim
            and anchor_prim.IsA(UsdMedia.SpatialAudio)
            or anchor_prim.IsA(OmniAudioSchema.Sound)
            or anchor_prim.IsA(OmniAudioSchema.Listener)
        ):
            unwanted_widgets_to_build.append("geometry_imageable")

        elif anchor_prim and anchor_prim.IsA(UsdSkel.Skeleton) or anchor_prim.IsA(UsdSkel.Root):
            unwanted_widgets_to_build.append("geometry")
            unwanted_widgets_to_build.append("geometry_imageable")

        return unwanted_widgets_to_build
