# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
from .builders import (
   BitfieldWidgetBuilder,
   LocalSpaceVelocitiesWidgetBuilder,
   QuatEulerRotationBuilder,
   RelationshipWidgetBuilder,
   ModelWithWidgetBuilder,
   ReferenceFrameIsCenterOfMassWidgetBuilder,
)
from .builders import GearingWidgetBuilder, HideWidgetBuilder, PrettyPrintTokenComboBuilder
from .widgets import FixedTendonWidget, SpatialTendonWidget
from .widgets import MassAPIWidget, JointWidget, RevoluteJointWidget, PrismaticJointWidget, DistanceJointWidget, SphericalJointWidget
from .widgets import LimitWidget, DriveWidget, JointAxisWidget, ExtendedSceneWidget, FixedJointWidget, JointStateWidget, GearJointWidget, RackAndPinionJoint
from. widgets import ExtendedColliderWidget
from .widgets import ExtendedRigidMaterialWidget
from .widgets import (
    ExtendedDeformableBodyWidgetDeprecated,
    ExtendedDeformableSurfaceWidgetDeprecated,
    ExtendedTetrahedralMeshWidgetDeprecated,
    ExtendedDeformalbeSurfaceMaterialWidgetDeprecated,
    ExtendedAttachmentWidgetDeprecated,
    ExtendedComputeAutoAttachmentWidgetDeprecated,
    ExtendedParticleClothWidgetDeprecated,
    ExtendedAutoParticleClothWidgetDeprecated, 
)
from .widgets import (
    ExtendedDeformableBodyWidget,
    ExtendedSurfaceDeformableMaterialWidget,
    ExtendedAutoDeformableAttachmentWidget,
    ExtendedAttachmentElementFilterWidget,
)
from .widgets import (
    ExtendedVehicleContextWidget,
    ExtendedVehicleWidget,
    ExtendedVehicleWheelWidget,
    ExtendedVehicleSuspensionWidget,
    ExtendedVehicleTireWidget,
    ExtendedVehicleWheelAttachmentWidget,
    ExtendedVehicleDriveStandardWidget,
    ExtendedVehicleMultiWheelDifferentialWidget,
    ExtendedVehicleControllerBaseWidget,
    ExtendedVehicleControllerWidget,
    ExtendedMimicJointWidget,
)
from .widgets import ExtendedParticleSystemWidget
from .widgets import LocalSpaceVelocitiesWidget
from pxr import UsdPhysics, UsdGeom, PhysxSchema
import omni.physx.bindings._physx as pxb

custom_builders = {
    "physxRigidBody:lockedPosAxis": [BitfieldWidgetBuilder, ["X", "Y", "Z"]],
    "physxRigidBody:lockedRotAxis": [BitfieldWidgetBuilder, ["X", "Y", "Z"]],
    "physics:body0": [RelationshipWidgetBuilder, [UsdGeom.Xformable], 1],
    "physics:body1": [RelationshipWidgetBuilder, [UsdGeom.Xformable], 1],
    "physics:localRot0": [QuatEulerRotationBuilder],
    "physics:localRot1": [QuatEulerRotationBuilder],
    "physics:principalAxes": [QuatEulerRotationBuilder],
    "physics:simulationOwner": [RelationshipWidgetBuilder, [UsdPhysics.Scene]],
    "physics:filteredGroups": [RelationshipWidgetBuilder, [UsdPhysics.CollisionGroup]],
    "physxContactReport:reportPairs": [RelationshipWidgetBuilder, [], 0, lambda prim: prim.HasAPI(UsdPhysics.CollisionAPI)],
    "physics:hinge0": [RelationshipWidgetBuilder, [UsdPhysics.RevoluteJoint], 1],
    "physics:hinge1": [RelationshipWidgetBuilder, [UsdPhysics.RevoluteJoint], 1],
    "physics:hinge": [RelationshipWidgetBuilder, [UsdPhysics.RevoluteJoint], 1],
    "physics:prismatic": [RelationshipWidgetBuilder, [UsdPhysics.PrismaticJoint], 1],
    "groundMaterials": [RelationshipWidgetBuilder, [], 0, lambda prim: prim.HasAPI(PhysxSchema.PhysxMaterialAPI)],
    "physxVehicleContext:verticalAxis": [PrettyPrintTokenComboBuilder, ["X+", "X-", "Y+", "Y-", "Z+", "Z-", "undefined"]],
    "physxVehicleContext:longitudinalAxis": [PrettyPrintTokenComboBuilder, ["X+", "X-", "Y+", "Y-", "Z+", "Z-", "undefined"]],
    "physxVehicleTire:frictionTable": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.IsA(PhysxSchema.PhysxVehicleTireFrictionTable)],
    "physxVehicleWheelAttachment:wheel": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleWheelAPI)],
    "physxVehicleWheelAttachment:tire": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleTireAPI)],
    "physxVehicleWheelAttachment:suspension": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleSuspensionAPI)],
    "physxVehicleWheelAttachment:suspensionFrameOrientation": [QuatEulerRotationBuilder],
    "physxVehicleWheelAttachment:wheelFrameOrientation": [QuatEulerRotationBuilder],
    "physxVehicleWheelAttachment:collisionGroup": [RelationshipWidgetBuilder, [UsdPhysics.CollisionGroup], 1],
    "physxVehicleDriveStandard:engine": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleEngineAPI)],
    "physxVehicleDriveStandard:gears": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleGearsAPI)],
    "physxVehicleDriveStandard:autoGearBox": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleAutoGearBoxAPI)],
    "physxVehicleDriveStandard:clutch": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleClutchAPI)],
    "physxVehicle:drive": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxVehicleDriveBasicAPI) or prim.HasAPI(PhysxSchema.PhysxVehicleDriveStandardAPI)],
    PhysxSchema.Tokens.referenceFrameIsCenterOfMass: [ReferenceFrameIsCenterOfMassWidgetBuilder, PhysxSchema.Tokens.referenceFrameIsCenterOfMass],
    "physxPhysics:simulationOwner": [RelationshipWidgetBuilder, [UsdPhysics.Scene, PhysxSchema.PhysxParticleSystem], 1],
    "rigid": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(UsdPhysics.RigidBodyAPI)],
    # DEPRECATED
    "deformable": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxParticleClothAPI)],
    #~DEPRECATED
    "physxScene:broadphaseType": [ModelWithWidgetBuilder],
    "physxScene:collisionSystem": [ModelWithWidgetBuilder],
    "state:angular:physics:position": [ModelWithWidgetBuilder],
    "state:angular:physics:velocity": [ModelWithWidgetBuilder],
    "state:linear:physics:position": [ModelWithWidgetBuilder],
    "state:linear:physics:velocity": [ModelWithWidgetBuilder],
    "state:rotX:physics:position": [ModelWithWidgetBuilder],
    "state:rotX:physics:velocity": [ModelWithWidgetBuilder],
    "state:rotY:physics:position": [ModelWithWidgetBuilder],
    "state:rotY:physics:velocity": [ModelWithWidgetBuilder],
    "state:rotZ:physics:position": [ModelWithWidgetBuilder],
    "state:rotZ:physics:velocity": [ModelWithWidgetBuilder],
    "state:transX:physics:position": [ModelWithWidgetBuilder],
    "state:transX:physics:velocity": [ModelWithWidgetBuilder],
    "state:transY:physics:position": [ModelWithWidgetBuilder],
    "state:transY:physics:velocity": [ModelWithWidgetBuilder],
    "state:transZ:physics:position": [ModelWithWidgetBuilder],
    "state:transZ:physics:velocity": [ModelWithWidgetBuilder],
    # DEPRECATED
    "physxDeformable:kinematicEnabled": [ModelWithWidgetBuilder],
    "physxDeformable:simulationHexahedralResolution": [ModelWithWidgetBuilder],
    "physxDeformable:collisionSimplificationRemeshing": [ModelWithWidgetBuilder],
    "physxDeformable:collisionSimplificationRemeshingResolution": [ModelWithWidgetBuilder],
    "physxDeformable:collisionSimplificationTargetTriangleCount": [ModelWithWidgetBuilder],
    "physxDeformable:collisionSimplificationForceConforming": [ModelWithWidgetBuilder],
    "physxAutoAttachment:enableDeformableVertexAttachments": [ModelWithWidgetBuilder],
    "physxAutoAttachment:deformableVertexOverlapOffset": [ModelWithWidgetBuilder],
    "physxAutoAttachment:enableRigidSurfaceAttachments": [ModelWithWidgetBuilder],
    "physxAutoAttachment:rigidSurfaceSamplingDistance": [ModelWithWidgetBuilder],
    "physxAutoAttachment:enableCollisionFiltering": [ModelWithWidgetBuilder],
    "physxAutoAttachment:collisionFilteringOffset": [ModelWithWidgetBuilder],
    "physxAutoAttachment:maskShapes": [RelationshipWidgetBuilder, [], 0, lambda prim: prim.IsA(UsdGeom.Sphere) or prim.IsA(UsdGeom.Cube) or prim.IsA(UsdGeom.Capsule)],
    #~DEPRECATED
    "physxAutoDeformableAttachment:enableDeformableVertexAttachments": [ModelWithWidgetBuilder],
    "physxAutoDeformableAttachment:deformableVertexOverlapOffset": [ModelWithWidgetBuilder],
    "physxAutoDeformableAttachment:enableRigidSurfaceAttachments": [ModelWithWidgetBuilder],
    "physxAutoDeformableAttachment:rigidSurfaceSamplingDistance": [ModelWithWidgetBuilder],
    "physxAutoDeformableAttachment:enableCollisionFiltering": [ModelWithWidgetBuilder],
    "physxAutoDeformableAttachment:collisionFilteringOffset": [ModelWithWidgetBuilder],
    "physxAutoDeformableAttachment:maskShapes": [RelationshipWidgetBuilder, [], 0, lambda prim: prim.IsA(UsdGeom.Sphere) or prim.IsA(UsdGeom.Cube) or prim.IsA(UsdGeom.Capsule)],
    "parentLink": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.HasAPI(PhysxSchema.PhysxTendonAttachmentAPI) or prim.HasAPI(PhysxSchema.PhysxTendonAttachmentRootAPI)],
    "jointAxis": [HideWidgetBuilder],
    "physics:approximation": [PrettyPrintTokenComboBuilder, ["Triangle Mesh", "Convex Decomposition", "Convex Hull", "Bounding Sphere", "Bounding Cube", "Mesh Simplification"],
                              [('sdf', 'SDF Mesh'),('sphereFill', 'Sphere Approximation')]],
    "actor0": [RelationshipWidgetBuilder, [UsdGeom.Xformable], 1],
    "actor1": [RelationshipWidgetBuilder, [UsdGeom.Xformable], 1],
    "simulationOwner": [RelationshipWidgetBuilder, [UsdPhysics.Scene], 1],
    pxb.METADATA_ATTRIBUTE_NAME_LOCALSPACEVELOCITIES: [LocalSpaceVelocitiesWidgetBuilder, pxb.METADATA_ATTRIBUTE_NAME_LOCALSPACEVELOCITIES],
    "physxMaterial:compliantContactDamping": [ModelWithWidgetBuilder],
    "physxMimicJoint:rotX:referenceJoint": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.IsA(UsdPhysics.Joint)],
    "physxMimicJoint:rotY:referenceJoint": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.IsA(UsdPhysics.Joint)],
    "physxMimicJoint:rotZ:referenceJoint": [RelationshipWidgetBuilder, [], 1, lambda prim: prim.IsA(UsdPhysics.Joint)],
    "physxMimicJoint:rotX:referenceJointAxis": [ModelWithWidgetBuilder],
    "physxMimicJoint:rotY:referenceJointAxis": [ModelWithWidgetBuilder],
    "physxMimicJoint:rotZ:referenceJointAxis": [ModelWithWidgetBuilder],
}

custom_widgets = {
    UsdPhysics.MassAPI: MassAPIWidget,
    UsdPhysics.Joint: JointWidget,
    UsdPhysics.FixedJoint: FixedJointWidget,
    UsdPhysics.RevoluteJoint: RevoluteJointWidget,
    UsdPhysics.PrismaticJoint: PrismaticJointWidget,
    UsdPhysics.DistanceJoint: DistanceJointWidget,
    UsdPhysics.SphericalJoint: SphericalJointWidget,
    PhysxSchema.PhysxPhysicsGearJoint: GearJointWidget,
    PhysxSchema.PhysxPhysicsRackAndPinionJoint: RackAndPinionJoint,
    UsdPhysics.LimitAPI: LimitWidget,
    UsdPhysics.DriveAPI: DriveWidget,
    "PhysxJointAxisAPI": JointAxisWidget,
    PhysxSchema.JointStateAPI: JointStateWidget,
    UsdPhysics.Scene: ExtendedSceneWidget,
    UsdPhysics.MaterialAPI: ExtendedRigidMaterialWidget,
    PhysxSchema.PhysxTendonAxisAPI: FixedTendonWidget,
    PhysxSchema.PhysxTendonAxisRootAPI: FixedTendonWidget,
    PhysxSchema.PhysxTendonAttachmentRootAPI: SpatialTendonWidget,
    PhysxSchema.PhysxTendonAttachmentAPI: SpatialTendonWidget,
    PhysxSchema.PhysxTendonAttachmentLeafAPI: SpatialTendonWidget,
    UsdPhysics.CollisionAPI: ExtendedColliderWidget,
    UsdPhysics.RigidBodyAPI: LocalSpaceVelocitiesWidget,
    # DEPRECATED
    PhysxSchema.PhysxDeformableBodyAPI: ExtendedDeformableBodyWidgetDeprecated,
    PhysxSchema.PhysxDeformableSurfaceAPI: ExtendedDeformableSurfaceWidgetDeprecated,
    PhysxSchema.PhysxDeformableSurfaceMaterialAPI: ExtendedDeformalbeSurfaceMaterialWidgetDeprecated,
    PhysxSchema.PhysxPhysicsAttachment: ExtendedAttachmentWidgetDeprecated,
    PhysxSchema.PhysxAutoAttachmentAPI: ExtendedComputeAutoAttachmentWidgetDeprecated,
    PhysxSchema.TetrahedralMesh: ExtendedTetrahedralMeshWidgetDeprecated,
    #~DEPRECATED
    "OmniPhysicsDeformableBodyAPI": ExtendedDeformableBodyWidget,
    "OmniPhysicsSurfaceDeformableMaterialAPI": ExtendedSurfaceDeformableMaterialWidget,
    "PhysxAutoDeformableAttachmentAPI": ExtendedAutoDeformableAttachmentWidget,
    "OmniPhysicsVtxVtxAttachment": ExtendedAttachmentElementFilterWidget,
    "OmniPhysicsVtxTriAttachment": ExtendedAttachmentElementFilterWidget,
    "OmniPhysicsVtxTetAttachment": ExtendedAttachmentElementFilterWidget,
    "OmniPhysicsVtxXformAttachment": ExtendedAttachmentElementFilterWidget,
    "OmniPhysicsTriTriAttachment": ExtendedAttachmentElementFilterWidget,
    "OmniPhysicsElementCollisionFilter": ExtendedAttachmentElementFilterWidget,
    PhysxSchema.PhysxVehicleContextAPI: ExtendedVehicleContextWidget,
    PhysxSchema.PhysxVehicleAPI: ExtendedVehicleWidget,
    PhysxSchema.PhysxVehicleWheelAPI: ExtendedVehicleWheelWidget,
    PhysxSchema.PhysxVehicleSuspensionAPI: ExtendedVehicleSuspensionWidget,
    PhysxSchema.PhysxVehicleTireAPI: ExtendedVehicleTireWidget,
    PhysxSchema.PhysxVehicleWheelAttachmentAPI: ExtendedVehicleWheelAttachmentWidget,
    PhysxSchema.PhysxVehicleDriveStandardAPI: ExtendedVehicleDriveStandardWidget,
    PhysxSchema.PhysxVehicleMultiWheelDifferentialAPI: ExtendedVehicleMultiWheelDifferentialWidget,
    PhysxSchema.PhysxVehicleControllerAPI: ExtendedVehicleControllerWidget,
    PhysxSchema.PhysxVehicleTankControllerAPI: ExtendedVehicleControllerBaseWidget,
    # DEPRECATED
    PhysxSchema.PhysxParticleClothAPI: ExtendedParticleClothWidgetDeprecated,
    PhysxSchema.PhysxAutoParticleClothAPI: ExtendedAutoParticleClothWidgetDeprecated,
    #~DEPRECATED
    PhysxSchema.PhysxParticleSystem: ExtendedParticleSystemWidget,
    PhysxSchema.PhysxMimicJointAPI: ExtendedMimicJointWidget,
}
