Untitled

import hou
import viewerstate.utils as su
import toolutils

class MultiRigPoseContainerState(object):
    
    def __init__(self, state_name, scene_viewer):
        self.state_name = state_name
        self.scene_viewer = scene_viewer
        self.rigpose_nodes = []  # List of rigpose nodes to manipulate

    # Enter the viewer state and handle node selection or fallback
    def onEnter(self, kwargs):
        self.scene_viewer.setPromptMessage("Multi Rig Pose State Active: Select and Edit Multiple Rigs")
        
        # Try to collect selected rigpose nodes first
        if not self._collectSelectedRigPoseNodes():
            # If no nodes are selected, collect all rigpose nodes inside the container
            node = kwargs['node']  # Get the current container node
            self._collectAllRigPoseNodes(node)
        
        if not self.rigpose_nodes:
            hou.ui.displayMessage("No rigpose nodes found!")

    # Collect rigpose nodes from the current selection
    def _collectSelectedRigPoseNodes(self):
        # Get the selected nodes from the network editor
        selected_nodes = hou.selectedNodes()

        # Filter the selection to only include `rigpose` nodes
        self.rigpose_nodes = [
            node for node in selected_nodes
            if node.type().name() == "rigpose"
        ]

        return len(self.rigpose_nodes) > 0  # Return True if nodes were selected

    # Fallback: Collect all rigpose nodes inside the container HDA
    def _collectAllRigPoseNodes(self, container_node):
        # Dive into the container node and find all rigpose nodes
        self.rigpose_nodes = [
            child_node for child_node in container_node.allSubChildren()
            if child_node.type().name() == "rigpose"
        ]

    # Update the viewport for multiple rigs
    def onDraw(self, kwargs):
        handle = kwargs['draw_handle']
        viewport = kwargs['viewport']

        # Draw each rigpose node's rig
        for rigpose_node in self.rigpose_nodes:
            if rigpose_node is not None:
                geo = rigpose_node.geometry()
                su.drawRigPose(handle, geo, viewport, color=(0.5, 0.5, 1.0))

    # Handle mouse events (translate, rotate, etc.)
    def onMouseEvent(self, kwargs):
        ui_event = kwargs['ui_event']
        device = ui_event.device()
        mouse_position = device.mousePosition()

        # Apply manipulation logic to each rigpose node
        for rigpose_node in self.rigpose_nodes:
            # Add your translation, rotation, or scaling code here
            pass

    # Exit the viewer state
    def onExit(self, kwargs):
        self.scene_viewer.setPromptMessage("Exiting Multi Rig Pose State")

# Register the custom viewer state
def createViewerStateTemplate():
    template = hou.ViewerStateTemplate(
        "multi_rigpose_container", "Multi Rig Pose Container Viewer", hou.sopNodeTypeCategory()
    )
    template.bindFactory(MultiRigPoseContainerState)
    template.bindHandle("xform", "handle_transform")
    return template

# Register the template in Houdini
hou.ui.registerViewerState(createViewerStateTemplate())