# noqa: PLC0302
"""
Commands that modify the topology of the OmniGraph
"""
import asyncio
from contextlib import suppress
from typing import Any, List, Optional, Tuple, Union

import carb
import omni.graph.core as og
import omni.kit
import omni.kit.commands
import omni.usd

from .object_lookup import ObjectLookup
from .typing import Attribute_t, AttributeType_t, ExtendedAttribute_t, Node_t


# ==============================================================================================================
def get_attr_from_string(attr_str, create_node=False, graph=None) -> Optional[og.Attribute]:
    """
    Args:
        attr_str: The full path of the attribute
        create_node: If True, and if the node does not exist, create an OG Prim node
        graph: The relevent graph, or the current graph if None
    """
    # TODO: When og.Settings.PRIM_NODES is removed this function can be replaced by og.Controller.attribute()
    if og.Settings()(og.Settings.PRIM_NODES):
        tokens = attr_str.split(".")
        if len(tokens) != 2:
            # Output bundles are prims and so have a "/" separator, not a "."
            tokens = attr_str.rsplit("/", 1)
            if len(tokens) != 2:
                return None
        node_path = tokens[0]
        attr_name = tokens[1]
        if not graph:
            node = og.get_node_by_path(node_path)
            if node:
                graph = node.get_graph()
                while graph.get_parent_graph():
                    graph = graph.get_parent_graph()
            else:
                graph = og.get_current_graph()
        if create_node:
            node = graph.create_node(node_path, "omni.graph.core.Prim", True)
        else:
            node = graph.get_node(node_path)
        if not node.is_valid():
            return None

        return node.get_attribute(attr_name)

    return og.Controller.attribute(attr_str)


# ==============================================================================================================
class ConnectPrimCommand(omni.kit.commands.Command):
    """
    Connect Prim **Command**.  Connects a bundle attribute to a prim.  This can be both for
    bundle purposes or for "pure relationship" type connections where we just want a relationship
    that points to a prim, without the connotations associated with bundles.

    Args:
        attr: The relationship attribute. This can be specified in the form
              of an attribute from a node, or a string that denotes the path to the attribute in USD.

        prim_path: The path to the prim that is to be the target of the relationship.
        is_bundle_connection: Whether this connection represents a bundle connection or just a
                              regular relationship
    """

    def __init__(self, attr: Union[str, og.Attribute], prim_path: str, is_bundle_connection: bool):
        self._attr_in = attr
        self._attr = None
        self._prim_path = prim_path
        self._is_bundle_connection = is_bundle_connection
        self._modify_usd = True
        self._did_doit = False

    def do(self):
        if isinstance(self._attr_in, str):
            if len(self._attr_in.split(".")) != 2:
                carb.log_error("Malformed attribute string in ConnectPrimCommand " + self._attr_in)
                return
            self._attr = get_attr_from_string(self._attr_in, create_node=False)
            if self._attr is None or not self._attr.is_valid():
                carb.log_error("Unable to retrieve attribute in ConnectPrimCommand " + self._attr_in)
                return
        else:
            self._attr = self._attr_in

        node = self._attr.get_node()
        if not node.is_backed_by_usd():
            self._modify_usd = False

        if self._attr.connectPrim(self._prim_path, self._modify_usd, self._is_bundle_connection):
            self._did_doit = True

    def undo(self):
        if self._did_doit:
            self._attr.disconnectPrim(self._prim_path, self._modify_usd, self._is_bundle_connection)


# ==============================================================================================================
class DisconnectPrimCommand(omni.kit.commands.Command):
    """
    Disconnect Prim **Command**.  Disconnects a bundle attribute from a prim.  This can be both for
    bundle purposes or for "pure relationship" type connections where we just want a relationship
    that points to a prim, without the connotations associated with bundles.

    Args:
        attr: The relationship attribute.  This can be specified in the form
              of an attribute from a node, or a string that denotes the path to the attribute in USD.

        prim_path: The path to the prim that is to be the target of the relationship.
        is_bundle_connection: Whether this connection represents a bundle connection or just a
                              regular relationship
    """

    def __init__(self, attr: Union[str, og.Attribute], prim_path: str, is_bundle_connection: bool):
        self._attr_in = attr
        self._attr = None
        self._prim_path = prim_path
        self._is_bundle_connection = is_bundle_connection
        self._modify_usd = True
        self._did_doit = False

    def do(self):
        if isinstance(self._attr_in, str):
            if len(self._attr_in.split(".")) != 2:
                carb.log_error("Malformed attribute string in DisconnectPrimCommand " + self._attr_in)
                return
            self._attr = get_attr_from_string(self._attr_in, create_node=False)
            if self._attr is None or not self._attr.is_valid():
                return
        else:
            self._attr = self._attr_in

        node = self._attr.get_node()
        if not node.is_backed_by_usd():
            self._modify_usd = False

        if self._attr.disconnectPrim(self._prim_path, self._modify_usd, self._is_bundle_connection):
            self._did_doit = True

    def undo(self):
        if self._did_doit:
            self._attr.connectPrim(self._prim_path, self._modify_usd, self._is_bundle_connection)


# ==============================================================================================================
class ConnectAttrsCommand(omni.kit.commands.Command):
    """
    Connect Attrs **Command**.  Causes two attributes to be connected together in the omnigraph

    Args:
        src_attr: The source (upstream) attribute.  This can be specified in the form
                   of an attribute from a node, or a string that denotes the path to the attribute
                   in USD.  If specified as a string path, if the node happens to be a prim and
                   the prim doesn't yet have prim node created for it, this command will create
                   one automatically.
        dest_attr: The destination (downstream) attribute.  This can be specified in the form
                   of an attribute from a node, or a string that denotes the path to the attribute
                   in USD.  If specified as a string path, if the node happens to be a prim and
                   the prim doesn't yet have prim node created for it, this command will create
                   one automatically.
        modify_usd: Whether to modify the underlying usd stage with this connection
        connection_type: Whether this is regular connection or something ore fancy, like data only
                         and execution only connections
    """

    def __init__(
        self,
        src_attr: Union[str, og.Attribute],
        dest_attr: Union[str, og.Attribute],
        modify_usd: bool,
        connection_type: og.ConnectionType = og.ConnectionType.CONNECTION_TYPE_REGULAR,
    ):
        self._src_attr = src_attr
        self._dst_attr = dest_attr
        self._modify_usd = modify_usd
        self._connection_type = connection_type
        self._did_doit = False
        self._value = None

    # ---------------------------------------------------------------------------------------------
    @staticmethod
    def _get_from_attr_path(attr_in: Union[og.Attribute, str]) -> Optional[og.Attribute]:
        """Gets the attribute from the given path and will create the node if it doesn't exist"""
        # TODO: When og.Settings.PRIM_NODES is removed this function can be replaced by og.Controller.attribute()
        if isinstance(attr_in, str):
            attr = get_attr_from_string(attr_in, create_node=False)
            if not attr:
                # We were not able to find a valid node / attribute with the given path.  Either
                # this is an invalid path, or it's the path to a prim for which we have not created
                # a node yet. We'll assume the latter here and try to create a node. If the path is
                # bogus, the prim node creation will fail.
                attr = get_attr_from_string(attr_in, create_node=True)
                if not attr:
                    carb.log_error(f"Malformed attribute string in ConnectAttrsCommand: {attr_in}")
                    return None
            return attr

        return attr_in

    # ---------------------------------------------------------------------------------------------
    @staticmethod
    def do_immediate(
        src_attr: Attribute_t,
        dest_attr: Attribute_t,
        modify_usd: bool,
        connection_type: og.ConnectionType = og.ConnectionType.CONNECTION_TYPE_REGULAR,
        return_old_value: bool = False,
    ):
        _src = src_attr
        _dst = dest_attr
        failure_return = (_src, _dst, connection_type, None, False) if return_old_value else False
        src_attr = ConnectAttrsCommand._get_from_attr_path(_src)
        dest_attr = ConnectAttrsCommand._get_from_attr_path(_dst)
        if not src_attr:
            carb.log_warn(f"Could not connect to unknown src attribute {_src}")
            return failure_return
        if not dest_attr:
            carb.log_warn(f"Could not connect to unknown destination attribute {_dst}")
            return failure_return

        src_node = src_attr.get_node()
        dest_node = dest_attr.get_node()
        nodes_in_usd = src_node.is_backed_by_usd() and dest_node.is_backed_by_usd()

        connected = src_attr.is_connected(dest_attr)
        value = None
        did_doit = False
        if not connected:
            # If there is USD backing for the destination attribute, we prefer to stash the value using the USD
            # mechanism because it is more robust currently.  og.Controller.get() can fail if we haven't loaded the
            # value into Fabric yet when this connection happens.
            if nodes_in_usd:
                with suppress(og.OmniGraphError):
                    usd_attr = og.ObjectLookup.usd_attribute(dest_attr)
                    value = usd_attr.Get()
            else:
                value = og.Controller(attribute=dest_attr).get()
            connection_info = og.ConnectionInfo(dest_attr, connection_type)
            if src_attr.connectEx(connection_info, modify_usd):
                did_doit = True

        if return_old_value:
            return (src_attr.get_path(), dest_attr.get_path(), connection_type, value, did_doit)

        return True

    # ---------------------------------------------------------------------------------------------
    def do(self):
        (
            self._src_attr,
            self._dst_attr,
            self._connection_type,
            self._value,
            self._did_doit,
        ) = ConnectAttrsCommand.do_immediate(
            self._src_attr, self._dst_attr, self._modify_usd, self._connection_type, return_old_value=True
        )

    # ---------------------------------------------------------------------------------------------
    def undo(self):
        if self._did_doit:
            src_attr = ConnectAttrsCommand._get_from_attr_path(self._src_attr)
            dst_attr = ConnectAttrsCommand._get_from_attr_path(self._dst_attr)
            if not src_attr or not dst_attr:
                return
            src_attr.disconnect(dst_attr, self._modify_usd)
            if self._value is not None:
                src_node = src_attr.get_node()
                dest_node = dst_attr.get_node()
                nodes_in_usd = src_node.is_backed_by_usd() and dest_node.is_backed_by_usd()
                if nodes_in_usd:
                    with suppress(og.OmniGraphError):
                        usd_attr = og.ObjectLookup.usd_attribute(dst_attr)
                        usd_attr.Set(self._value)
                else:
                    og.Controller(og.Controller.attribute(self._dst_attr)).set(self._value)
            elif dst_attr.get_type_name() != "bundle":
                carb.log_warn("No value to restore after undo of ConnectAttrsCommand")
            self._did_doit = False


# ==============================================================================================================
class DisconnectAllAttrsCommand(omni.kit.commands.Command):
    """
    Disconnect All Attrs **Command**.  Breaks every connection to and from an OmniGraph attribute

    Args:
        attr: The attribute to be disconnected
        modify_usd: Whether to modify the underlying usd stage with this connection
    """

    def __init__(self, attr: og.Attribute, modify_usd: bool):
        """Remember the command parameters"""
        self._attr = attr
        self._modify_usd = modify_usd
        self._disconnections = []
        self._graph = attr.get_node().get_graph()
        self._undo_cursor = 0  # Current location in the _disconnections of the undo/redo - usually 0 or len()-1

        while self._graph.get_parent_graph():
            self._graph = self._graph.get_parent_graph()

    def __remember_disconnection(self, upstream_attribute: og.Attribute, downstream_attribute: og.Attribute):
        """Add enough information to the disconnection to be able to recreate it in undo. There's no guarantee the
        same og.Attribute object will exist after a series of operations but the paths will be the same."""
        upstream_path = upstream_attribute.get_node().get_prim_path()
        upstream_name = upstream_attribute.get_name()
        downstream_path = downstream_attribute.get_node().get_prim_path()
        downstream_name = downstream_attribute.get_name()
        self._disconnections.append((upstream_path, upstream_name, downstream_path, downstream_name))

    def __get_attributes_from_connection_info(
        self, connection_info: Tuple[str, str, str, str]
    ) -> Tuple[og.Attribute, og.Attribute]:
        """Extract the upstream and downstream attribute from the stored connection information created by
        the __remember_disconnection function"""
        upstream_path, upstream_name, downstream_path, downstream_name = connection_info
        upstream_attribute = self._graph.get_node(upstream_path).get_attribute(upstream_name)
        downstream_attribute = self._graph.get_node(downstream_path).get_attribute(downstream_name)
        return (upstream_attribute, downstream_attribute)

    @staticmethod
    def do_immediate(attr: og.Attribute, modify_usd: bool) -> bool:
        """Do the disconnections without remembering the previous connections
        Args:
            attr: Attribute to be disconnected
            modify_usd: Is the USD to be immediately updated after the disconnections?
        Returns:
            True if all disconnections succeeded
        """
        status = True
        upstream_connections = attr.get_upstream_connections()
        for upstream_attribute in upstream_connections:
            status = upstream_attribute.disconnect(attr, modify_usd) and status
        downstream_connections = attr.get_downstream_connections()
        for downstream_attribute in downstream_connections:
            status = attr.disconnect(downstream_attribute, modify_usd) and status
        return status

    def do(self):
        """Perform the disconnections, remembering what was disconnected"""
        if self._attr is not None:
            upstream_connections = self._attr.get_upstream_connections()
            for upstream_attribute in upstream_connections:
                if upstream_attribute.disconnect(self._attr, self._modify_usd):
                    self.__remember_disconnection(upstream_attribute, self._attr)
            downstream_connections = self._attr.get_downstream_connections()
            for downstream_attribute in downstream_connections:
                if self._attr.disconnect(downstream_attribute, self._modify_usd):
                    self.__remember_disconnection(self._attr, downstream_attribute)
            self._undo_cursor = len(self._disconnections)
            # Mark the operation as being done once, to avoid accessing potentially stale data
            self._attr = None
        else:
            while self._undo_cursor < len(self._disconnections):
                upstream_attribute, downstream_attribute = self.__get_attributes_from_connection_info(
                    self._disconnections[self._undo_cursor]
                )
                if upstream_attribute.disconnect(downstream_attribute, self._modify_usd):
                    self._undo_cursor += 1
                else:
                    return

    def undo(self):
        """Use the remembered disconnections to reestablish the connections"""
        while self._undo_cursor > 0:
            self._undo_cursor -= 1
            upstream_attribute, downstream_attribute = self.__get_attributes_from_connection_info(
                self._disconnections[self._undo_cursor]
            )
            if not upstream_attribute.connect(downstream_attribute, self._modify_usd):
                self._undo_cursor += 1
                return


# ==============================================================================================================
class DisconnectAttrsCommand(omni.kit.commands.Command):
    """
    Disconnect Attrs **Command**.  Causes two attrs to be disconnected in OmniGraph

    Args:
        src_attr: The source (upstream) attribute
        dest_attr: The destination (downstream) attribute
        modify_usd: Whether to modify the underlying usd stage with this connection
    """

    def __init__(self, src_attr: Attribute_t, dest_attr: Attribute_t, modify_usd: bool):
        self._src_attr = src_attr.get_path()
        self._dst_attr = dest_attr.get_path()
        self._modify_usd = modify_usd
        self._graph = src_attr.get_node().get_graph()
        self._did_doit = False

        while self._graph.get_parent_graph():
            self._graph = self._graph.get_parent_graph()

    @staticmethod
    def do_immediate(
        src_attr: Attribute_t,
        dest_attr: Attribute_t,
        modify_usd: bool,
    ):
        connected = src_attr.is_connected(dest_attr)
        if connected:
            # since the dest attr is being driven, there is no need to save the state
            # of the driven attribute.  When undo happens, it'll simply be driven again
            return src_attr.disconnect(dest_attr, modify_usd)

        carb.log_warn(f"Could not find connection between {src_attr} and {dest_attr} to remove")
        return False

    def do(self):
        src_attr = get_attr_from_string(self._src_attr, create_node=False, graph=self._graph)
        dst_attr = get_attr_from_string(self._dst_attr, create_node=False, graph=self._graph)
        if src_attr is None:
            carb.log_warn(f"Could not find source attribute {self._src_attr} to disconnect from")
            return
        if dst_attr is None:
            carb.log_warn(f"Could not find destination attribute {self._dst_attr} to disconnect from")
            return
        self._did_doit = DisconnectAttrsCommand.do_immediate(src_attr, dst_attr, self._modify_usd)

    def undo(self):
        if self._did_doit:
            src_attr = get_attr_from_string(self._src_attr, create_node=False, graph=self._graph)
            dst_attr = get_attr_from_string(self._dst_attr, create_node=False, graph=self._graph)
            src_attr.connect(dst_attr, self._modify_usd)
            self._did_doit = False


# ==============================================================================================================
class _AttributeFactory:
    """Helper class shared by the CreateAttr and RemoveAttr commands since both require the ability to create
    and remove attributes, either in do() or in undo(). It is not assumed that it is safe to hang on to
    PyBind objects after the __init__ call so lookup strings are saved instead.
    """

    def __raise_error(self, msg: str):
        """Raise an OmniGraphError populated with common identifying information"""
        self._failure_pending = True
        raise og.OmniGraphError(f"Failed to {msg} using attribute {self._attr_name} on node {self._node_path}")

    def __init__(self, node: Node_t, *args):
        """Create attribute information either from an actual attribute or from individual create parameters"""
        # If one or the other operation failed then undo will attempt the opposite, which should just be skipped
        self._failure_pending = False
        if len(args) == 1:
            self._created = True
            attribute = ObjectLookup.attribute(args[0])
            self._node_path = attribute.get_node().get_prim_path()
            self._attr_name = attribute.get_name()
            self._attr_type = attribute.get_attribute_data().get_type()
            self._attr_port = attribute.get_port_type()
            if attribute.get_extended_type() == og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION:
                self._attr_extended_type = (attribute.get_extended_type, ",".join(attribute.get_union_types()))
            else:
                self._attr_extended_type = attribute.get_extended_type()
            # TODO: Really we should be getting the default from og.Attribute.getDefault(), but it doesn't exist
            self._attr_default = None
        elif len(args) == 5:
            self._created = False
            self._node_path = ObjectLookup.node(node).get_prim_path()
            self._attr_name = args[0]
            self._attr_type = args[1]
            self._attr_port = args[2]
            self._attr_default = args[3]
            self._attr_extended_type = args[4]
        else:
            raise og.OmniGraphError(
                "Attribute factory requires either an og.Attribute or tuple of name, type, port, default, extended_type"
            )

    def create(self):
        """Create a new attribute with saved parameters, checking to make sure it hasn't already been created"""
        if self._failure_pending:
            self._failure_pending = False
            return
        if self._created:
            self.__raise_error("create dynamic attribute twice")
        omg_node = ObjectLookup.node(self._node_path)
        args = [self._attr_name, self._attr_type, self._attr_port, self._attr_default]
        if isinstance(self._attr_extended_type, Tuple):
            args.append(self._attr_extended_type[0])
            if isinstance(self._attr_extended_type[1], list):
                args.append(",".join(self._attr_extended_type[1]))

        else:
            args.append(self._attr_extended_type)
        self._created = omg_node.create_attribute(*args)
        if not self._created:
            self._failure_pending = True
            self.__raise_error("add the dynamic attribute to the node")

    def remove(self) -> bool:
        """Remove an existing dynamic attribute, raising an exception if it was never created"""
        if self._failure_pending:
            self._failure_pending = False
            return
        omg_node = ObjectLookup.node(self._node_path)
        if not self._created or not omg_node.get_attribute_exists(self._attr_name):
            self.__raise_error("remove unknown dynamic attribute")
        self._created = not omg_node.remove_attribute(self._attr_name)
        if self._created:
            self.__raise_error("remove the dynamic attribute from the node")


# ==============================================================================================================
class CreateAttrCommand(omni.kit.commands.Command):
    """
    Create Attribute **Command**.  Adds a new dynamic attribute to a node.

    Args:
        node: Node on which to create the attribute (path or og.Node)
        attr_name: Name of the new attribute, either with or without the port namespace
        attr_type: Type of the new attribute, as an OGN type string or og.Type
        attr_port: Port type of the new attribute, default is og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT
        attr_default: The initial value to set on the attribute, default is None, meaning the type's default is used
        attr_extended_type: The extended type of the attribute, default is
                            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR. If the extended type is
                            og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION then this parameter will be a
                            2-tuple with the second element being a list or comma-separated string of union types

    If there is any problem with the do/undo of the attribute creation a warning is issued and the command fails.
    """

    def __init__(
        self,
        node: Node_t,
        attr_name: str,
        attr_type: AttributeType_t,
        attr_port: Optional[og.AttributePortType] = og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
        attr_default: Optional[Any] = None,
        attr_extended_type: Optional[ExtendedAttribute_t] = og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR,
    ):
        """The information in here is only meant to survive until the do() function is called. After that, the
        do() function will store information necessary to undo and redo the operations."""
        super().__init__()
        self._factory = _AttributeFactory(
            node, attr_name, ObjectLookup.attribute_type(attr_type), attr_port, attr_default, attr_extended_type
        )

    @staticmethod
    def do_immediate(
        node: Node_t,
        attr_name: str,
        attr_type: AttributeType_t,
        attr_port: og.AttributePortType = og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT,
        attr_default: Any = None,
        attr_extended_type: ExtendedAttribute_t = og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_REGULAR,
    ):
        """Removes the attribute, raising og.OmniGraphError if it fails"""
        attr_type = ObjectLookup.attribute_type(attr_type)
        _AttributeFactory(node, attr_name, attr_type, attr_port, attr_default, attr_extended_type).create()

    def do(self):
        """Create the dynamic attribute with the specified parameters, raising OmniGraphError if the create failed."""
        try:
            self._factory.create()
        except og.OmniGraphError as error:
            carb.log_warn(str(error))
            return False
        return True

    def undo(self):
        """Remove the created attribute, raising OmniGraphError if the removal failed"""
        try:
            self._factory.remove()
        except og.OmniGraphError as error:
            carb.log_warn(str(error))
            return False
        return True


# ==============================================================================================================
class RemoveAttrCommand(omni.kit.commands.Command):
    """
    Remove Attribute **Command**.  Removes an existing dynamic attribute from a node.

    Args:
        attribute: Name of the attribute to be removed

    If there is any problem with the do/undo of the attribute removal a warning is issue and the command fails
    """

    def __init__(
        self,
        attribute: og.Attribute,
    ):
        """The information in here is only meant to survive until the do() function is called. After that, the
        do() function will store information necessary to undo and redo the operations."""
        super().__init__()
        self._factory = _AttributeFactory(None, attribute)

    @staticmethod
    def do_immediate(attribute: og.Attribute):
        """Removes the dynamic attribute, raising og.OmniGraphError if the removal failed."""
        try:
            _AttributeFactory(None, attribute).remove()
        except og.OmniGraphError as error:
            carb.log_warn(str(error))
            return False
        return True

    def do(self):
        """Remove the dynamic attribute, raising OmniGraphError if the removal failed."""
        try:
            self._factory.remove()
        except og.OmniGraphError as error:
            carb.log_warn(str(error))
            return False
        return True

    def undo(self):
        """Undo the removal of the dynamic attribute, raising OmniGraphError if the creation failed"""
        try:
            self._factory.create()
        except og.OmniGraphError as error:
            carb.log_warn(str(error))
            return False
        return True


# ==============================================================================================================
class CreateNodeCommand(omni.kit.commands.Command):
    """
    Create Node **Command**.  Creates a new compute node of a particular node type in OmniGraph

    Args:
        graph: The graph in which the new compute node should be created
        node_path: The location in the USD stage to add the new compute node
        node_type: The name of the type of compute node to create
        create_usd: Whether to also create an USD prim on the stage for this node

    Raises:
        og.OmniGraphError if a node already exists at the specified node path
    """

    def __init__(self, graph: og.Graph, node_path: str, node_type: str, create_usd: bool):
        self._graph = graph
        self._node_path = node_path
        self._node_type = node_type
        self._modify_usd = create_usd
        self._did_doit = False
        self._node = None

    # Do the operation do_immediately, bypassing the command's interaction with the undo queue
    @staticmethod
    def do_immediate(graph: og.Graph, node_path: str, node_type: str, modify_usd: bool):
        # Don't create a node if one already exists at the specified path.
        node = graph.get_node(node_path)
        if node is not None and node.is_valid():
            raise og.OmniGraphError(f"Node already exists at {node_path}. Cannot create another.")
        return graph.create_node(node_path, node_type, modify_usd)

    def do(self):
        try:
            self._node = CreateNodeCommand.do_immediate(self._graph, self._node_path, self._node_type, self._modify_usd)
        except og.OmniGraphError:
            # This should raise an exception, but for backward compatibility when executed here it will just log a
            # warning and then succeed
            carb.log_warn(f"Node already exists at {self._node_path}. Cannot create another.")
            node = self._node
            self._node = None
            return node

        self._did_doit = True
        return self._node

    def undo(self):
        if self._did_doit and (self._node is not None):
            self._graph.destroy_node(self._node_path, self._modify_usd)
            self._node = None
            self._did_doit = False


# ==============================================================================================================
class CreateGraphAsNodeCommand(omni.kit.commands.Command):
    """
    Create Graph As Node **Command**.  Creates a new graph wrapped by a node.

    Args:
        graph: The graph in which the new wrapper node should be created
        node_name: The name of the node
        graph_path: The path to the graph
        evaluator_name: The name of the evaluator to use for the graph
        is_global_graph: Whether this is a global graph (global level graphs have their own FC)
        backed_by_usd: Whether the constructs are to be backed by USD
        fc_backing_type: What kind of Flatcache backs this graph
        pipeline_stage: What pipeline stage does this graph occupy: simulation, prerender, or postrender
        evaluation_mode: What evaluation mode to use with this graph: Automatic, Standalone or Instanced
    """

    def __init__(
        self,
        graph,
        node_name,
        graph_path,
        evaluator_name,
        is_global_graph,
        backed_by_usd,
        fc_backing_type,
        pipeline_stage,
        evaluation_mode=og.GraphEvaluationMode.GRAPH_EVALUATION_MODE_AUTOMATIC,
    ):
        self._graph = graph
        self._node_name = node_name
        self._graph_path = graph_path
        self._evaluator_name = evaluator_name
        self._is_global_graph = is_global_graph
        self._backed_by_usd = backed_by_usd
        self._fc_backing_type = fc_backing_type
        self._pipeline_stage = pipeline_stage
        self._evaluation_mode = evaluation_mode
        self._did_doit = False
        self._created_node = None

    # Do the operation do_immediately, bypassing the command's interaction with the undo queue
    @staticmethod
    def do_immediate(
        graph,
        node_name,
        graph_path,
        evaluator_name,
        is_global_graph,
        backed_by_usd,
        fc_backing_type,
        pipeline_stage,
        evaluation_mode,
    ) -> og.Node:
        return graph.create_graph_as_node(
            node_name,
            graph_path,
            evaluator_name,
            is_global_graph,
            backed_by_usd,
            fc_backing_type,
            pipeline_stage,
            evaluation_mode,
        )

    # Returns the created node wrapping the graph
    def do(self) -> og.Node:

        node = CreateGraphAsNodeCommand.do_immediate(
            self._graph,
            self._node_name,
            self._graph_path,
            self._evaluator_name,
            self._is_global_graph,
            self._backed_by_usd,
            self._fc_backing_type,
            self._pipeline_stage,
            self._evaluation_mode,
        )
        if node is not None and node.is_valid():
            self._did_doit = True
            self._created_node = node
        return node

    def undo(self):
        if self._did_doit and (self._created_node is not None):
            self._graph.destroy_node(self._created_node.get_prim_path(), self._backed_by_usd)
            self._created_node = None
            self._did_doit = False


# ==============================================================================================================
class CreateSubgraphCommand(omni.kit.commands.Command):
    """
    Create Subgraph **Command**.  Creates a new subgraph node in OmniGraph

    Args:
        graph: The graph in which the new compute node should be created
        subgraph_path: The location in the USD stage to add the new compute node
    """

    def __init__(self, graph, subgraph_path, evaluator=None, create_usd=True):
        self._graph = graph
        self._subgraph_path = subgraph_path
        self._evaluator = evaluator
        self._create_usd = create_usd
        self._did_doit = False
        self._subgraph = None

    def do(self):
        # Don't create a subgraph if one already exists at the specified path.
        subgraph = self._graph.get_node(self._subgraph_path)
        if subgraph is None or not subgraph.is_valid():
            self._subgraph = self._graph.create_subgraph(self._subgraph_path, self._evaluator, self._create_usd)
            self._did_doit = True
            return self._subgraph
        return subgraph

    def undo(self):
        if self._did_doit and (self._subgraph is not None):
            self._graph.destroy_node(self._subgraph_path, True)
            if self._create_usd:
                delete_cmd = omni.usd.commands.DeletePrimsCommand([self._subgraph_path])
                delete_cmd.do()
            self._subgraph = None
            self._did_doit = False


# ==============================================================================================================
class DeleteNodeCommand(omni.kit.commands.Command):
    """
    Delete Node **Command**.  Delete the specified node in the graph

    Args:
        graph: The graph in which the new compute node should be created
        node_path: The location in the USD stage to add the new compute node
        modify_usd: Whether to also delete the USD prim on the stage for this node
    """

    def __init__(self, graph: og.Graph, node_path: str, modify_usd: bool):
        self._graph = graph
        self._node_path = node_path
        self._node_type = None
        self._modify_usd = modify_usd
        self._did_doit = False
        self._upstream_connections = {}
        self._downstream_connections = {}

    # Do the operation do_immediately, bypassing the command's interaction with the undo queue
    @staticmethod
    def do_immediate(graph: og.Graph, node_path: str, modify_usd: bool):
        upstream_connections = {}
        downstream_connections = {}
        did_doit = False
        node = graph.get_node(node_path)
        if node is not None and node.is_valid():
            attributes = node.get_attributes()
            for attribute in attributes:
                _upstream_connections = attribute.get_upstream_connections()
                upstream_connections[attribute.get_name()] = [
                    (src_attr.get_node().get_prim_path(), src_attr.get_name()) for src_attr in _upstream_connections
                ]
                for src_attr in _upstream_connections:
                    src_attr.disconnect(attribute, modify_usd)

                _downstream_connections = attribute.get_downstream_connections()
                downstream_connections[attribute.get_name()] = [
                    (dst_attr.get_node().get_prim_path(), dst_attr.get_name()) for dst_attr in _downstream_connections
                ]
                for dst_attr in _downstream_connections:
                    attribute.disconnect(dst_attr, modify_usd)

            node_type = node.get_type_name()
            if not node_type:
                node_type = node.get_type_name()
            graph.destroy_node(node_path, modify_usd)
            did_doit = True

            return (node_type, upstream_connections, downstream_connections, did_doit)
        return (None, {}, {}, False)

    def do(self):
        (
            self._node_type,
            self._upstream_connections,
            self._downstream_connections,
            self._did_doit,
        ) = DeleteNodeCommand.do_immediate(self._graph, self._node_path, self._modify_usd)

    def undo(self):
        if self._did_doit and (self._node_type is not None):
            node = self._graph.create_node(self._node_path, self._node_type, self._modify_usd)
            for attribute_name, src_connections in self._upstream_connections.items():
                attribute = node.get_attribute(attribute_name)
                for src_node_path, src_attr_name in src_connections:
                    src_node = self._graph.get_node(src_node_path)
                    src_attr = src_node.get_attribute(src_attr_name)
                    src_attr.connect(attribute, self._modify_usd)

            for attribute_name, dest_connections in self._downstream_connections.items():
                attribute = node.get_attribute(attribute_name)
                for dest_node_path, dest_attr_name in dest_connections:
                    dest_node = self._graph.get_node(dest_node_path)
                    dest_attr = dest_node.get_attribute(dest_attr_name)
                    attribute.connect(dest_attr, self._modify_usd)

            self._node_type = None
            self._did_doit = False


# ==============================================================================================================
class CreateVariableCommand(omni.kit.commands.Command):
    """
    Create Variable **Command**.  Creates a new variable of a particular type in OmniGraph

    Args:
        graph: The graph in which the new variable should be created
        variable_name: The name of the new variable
        variable_type: The OmniGraph type of the new variable
        graph_context: The OmniGraph context to use when setting the initial variable value
        variable_value: The initial variable value
    """

    def __init__(
        self,
        graph: og.Graph,
        variable_name: str,
        variable_type: og.Type,
        graph_context: og.GraphContext = None,
        variable_value=None,
    ):
        self._graph = graph
        self._variable_name = variable_name
        self._variable_type = variable_type
        self._graph_context = graph_context
        self._variable_value = variable_value
        self._created_variable = None

    @staticmethod
    def do_immediate(
        graph: og.Graph,
        variable_name: str,
        variable_type: og.Type,
        graph_context: og.GraphContext = None,
        variable_value=None,
    ):
        if not graph or not graph.is_valid() or not variable_name or not variable_type:
            return None

        created_variable = graph.create_variable(variable_name, variable_type)
        if created_variable and variable_value is not None:
            og.Controller.set_variable_default_value(created_variable, variable_value)
        return created_variable

    def do(self):
        self._created_variable = CreateVariableCommand.do_immediate(
            self._graph,
            self._variable_name,
            self._variable_type,
            self._graph_context,
            self._variable_value,
        )
        return self._created_variable

    def undo(self):
        if self._created_variable is None:
            return

        self._graph.remove_variable(self._created_variable)
        self._created_variable = None


# ==============================================================================================================
class RemoveVariableCommand(omni.kit.commands.Command):
    """
    Remove Variable **Command**.  Remove the specified variable in the graph

    Args:
        graph: The graph to remove the variable from
        variable: The OmniGraph IVariable to be removed
        graph_context: The OmniGraph context to use when restoring the variable value on undo
    """

    def __init__(
        self,
        graph: og.Graph,
        variable: og.IVariable,
        graph_context: og.GraphContext = None,
    ):
        self._graph = graph
        self._variable = variable
        self._variable_name = None
        self._variable_type = None
        self._graph_context = graph_context
        self._variable_value = None
        self._removed = False

    def do(self):
        if not self._graph or not self._graph.is_valid() or not self._variable:
            return

        self._variable_name = self._variable.name
        self._variable_type = self._variable.type
        self._variable_value = og.Controller.get_variable_default_value(self._variable)
        self._removed = self._graph.remove_variable(self._variable)

    def undo(self):
        if not self._removed:
            return

        self._variable = self._graph.create_variable(self._variable_name, self._variable_type)
        if self._variable_value is not None:
            og.Controller.set_variable_default_value(self._variable, self._variable_value)


class _OGRestoreConnectionsOnUndo(omni.kit.commands.Command):
    """
    Restore connections between OG nodes on undo. (Does nothing on do or redo.)

    This command is for internal use only. It may be changed or removed
    without notice.

    Args:
        connections (List[(str, str)])
            The connections to be restored. Each element of the list is a tuple
            containing the path strings for the source and destination attributes.
    """

    # When the source prim for a connection is deleted OG removes all traces of the
    # connection. If the deletion is undone OG has no way of knowing that the restored
    # prim had a connection which should also be restored. This command is used to
    # restore those connections on undo.
    def __init__(self, connections: List[Tuple[str, str]]):
        self._path_strings: List[Tuple[str, str]] = connections.copy()
        self.__reconnect_task = None

    def destroy(self):
        if self.__reconnect_task:
            if not self.__reconnect_task.done():
                self.__reconnect_task.cancel()
            self.__reconnect_task = None

    def do(self):
        pass

    def undo(self):
        # We cannot do the reconnection yet because OG won't have created the prim's
        # Node yet.
        if self.__reconnect_task is None or self.__reconnect_task.done():
            self.__reconnect_task = asyncio.ensure_future(self.__do_reconnections())

    async def __do_reconnections(self):
        # Give OG a chance to create the prim's Node.
        await omni.kit.app.get_app().next_update_async()

        for src_attr_str, dest_attr_str in self._path_strings:
            src_attr: og.Attribute = og.Controller.attribute(src_attr_str)
            dest_attr: og.Attribute = og.Controller.attribute(dest_attr_str)
            if src_attr and dest_attr and src_attr not in dest_attr.get_upstream_connections():
                src_attr.connect(dest_attr, modify_usd=True)
        self.__reconnect_task = None
