# noqa: PLC0302
"""Utilities that handles interactions with the set of OmniGraph graphs and their contents.

The main class is og.Controller, which is an aggregate that implements each of the smaller interfaces that handle
various pieces of the graph manipulation. They are mainly kept separate to avoid a big monolithic implementation
class, while providing a single point of contact for graph manipulation due to the multiple interfaces it inherits.

All of the classes used here follow a pattern that allows the main functions to be called using either an object or a
class, accessing the same underlying functionality but with potentially different arguments.

.. code-block:: python

    class SharedMethodNames:
        def __init__(self, *args, **kwargs):
            # Replace the classmethod in instantiated objects with the internal object-based implementation
            self.dual_method = self.__dual_method_obj

        @staticmethod
        def __dual_method(obj, *args, **kwargs):
            # Implementation of the actual method

        @classmethod
        def dual_method(cls, *args, **kwargs):
            # Any manipulation of the arguments for class-based use is done first, then the implementation is called
            cls.__dual_method(cls, *args, **kwargs)

        def __dual_method_obj(self, *args, **kwargs):
            # Any manipulation of the arguments for object-based use is done first, then the implementation is called
            self.__dual_method(self, *args, **kwargs)

For example, here is an implementation of a class containing a "square()" function that will square the value it owns.
The object will initialized the value as part of its __init__() function, whereas the class method will take the value
as one of its arguments.

.. code-block:: python

    class FlexibleSquare:
        def __init__(self, value: float):
            self.square = self.__square_obj
            self.__value = value

        @staticmethod
        def __square(obj, value: float) -> float:
            return value * value

        @classmethod
        def square(cls, value: float) -> float:
            return cls.__square(cls, value)

        def __square_obj(self) -> float:
            return self.__square(self, self.__value)

    print(f"Square of 3 is {FlexibleSquare.square(3)}")
    five = FlexibleSquare(5)
    print(f"Square of 5 is {five.square()}")

It might also make sense to allow override in the method for the object's method:

    def __square_obj(self, value: float = None) -> float:
        return self.__square(self, self.__value if value is None else value)

    print(f"Square of 3 is {FlexibleSquare.square(3)}")
    five = FlexibleSquare(5)
    print(f"Square of 5 is {five.square()}")
    print(f"Square of 4 is {five.square(4)}")
"""
from __future__ import annotations

import asyncio
from contextlib import suppress
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import omni.graph.core as og
import omni.graph.tools.ogn as ogn
from omni.kit import undo
from pxr import Sdf, Usd

from .data_view import DataView
from .errors import OmniGraphError
from .graph_controller import GraphController
from .node_controller import NodeController
from .object_lookup import ObjectLookup
from .typing import (
    AttributeSpec_t,
    AttributeSpecs_t,
    AttributeType_t,
    GraphSpec_t,
    GraphSpecs_t,
    NewNode_t,
    NodeSpec_t,
    NodeType_t,
    Path_t,
    PrimAttrs_t,
    VariableName_t,
    VariableType_t,
)
from .utils import AttributeValues_t, ValueToSet_t, _flatten_arguments, _Unspecified

# Dictionary type mapping the relative path of a node that was created by the edit() method to the node/prim it created
# TODO: path_to_object_map + graph_path could be broken into a separate class for easier handling
PathToObjectMap_t = Dict[str, Union[og.Node, Usd.Prim]]


# ==============================================================================================================
def _get_as_list(value: Union[Any, List[Any]]) -> List[Any]:
    """Utility to ensure that a passed-in item is returned as a list.
    Args:
        value: Value or list of values to convert
    Returns:
        The value itself if it was already a list, otherwise a single element list containing the value
    """
    return value if isinstance(value, list) else [value]


# ==============================================================================================================
def _find_actual_path(path_id: Path_t, root_path: str, path_to_object_map: Dict[str, Any]) -> str:
    """Finds the actual creation path, applying node mapping and graph path prepending when needed"""
    node_path = str(path_id)
    if node_path.startswith("/"):
        return node_path
    try:
        node_object = path_to_object_map[node_path]
        return node_object.get_prim_path() if isinstance(node_object, og.Node) else node_object
    except KeyError:
        return f"{root_path}/{node_path}"


# ==============================================================================================================
class Controller(GraphController, NodeController, DataView, ObjectLookup):
    # begin-controller-docs
    r"""Class to provide a simple interface to a variety OmniGraph manipulation functions.

    Provides functions for creating nodes, making and breaking connections, and setting values.
    Graph manipulation functions are undoable, value changes are not.

    Functions are set up to be as flexible as possible, accepting a wide variety of argument variations.

    Here is a summary of the interface methods you can access through this class, grouped by interface class. The ones
    marked with an "*" indicate that they have both a classmethod version and an object method version. In those cases
    the methods use object member values which have corresponding arguments in the classmethod version. (e.g. if the
    method uses the "update_usd" member value then the method will also have a "update_usd:bool" argument)

    Controller
        \* edit             Perform a collection of edits on a graph (the union of all interfaces)

        async evaluate   Runs evaluation on one or more graphs as a waitable (typically called from async tests)

        evaluate_sync    Runs evaluation on one or more graphs immediately

    ObjectLookup
        attribute        Looks up an og.Attribute from a description

        attribute_path   Looks up an attribute string path from a description

        attribute_type   Looks up an og.Type from a description

        graph            Looks up an og.Graph from a description

        node             Looks up an og.Node from a description

        node_path        Looks up a node string path from a description

        prim             Looks up an Usd.Prim from a description

        prim_path        Looks up a Usd.Prim string path from a description

        split_graph_from_node_path  Separate a graph and a relative node path from a full node path

    GraphController
        \* connect          Makes connections between attribute pairs

        \* create_graph     Creates a new og.Graph

        \* create_node      Creates a new og.Node

        \* create_prim      Creates a new Usd.Prim

        \* create_variable  Creates a new og.IVariable

        \* delete_node      Deletes a list of og.Nodes

        \* disconnect       Breaks connections between attribute pairs

        \* disconnect_all   Breaks all connections to and from specific attributes

        \* expose_prim      Expose a USD prim to OmniGraph through an importing node

    NodeController
        \* create_attribute Create a new dynamic attribute on a node

        \* remove_attribute Remove an existing dynamic attribute from a node

        safe_node_name   Get a node name in a way that's safe for USD

    DataView
        \* get              Get the value of an attribute's data

        \* get_array_size   Get the number of elements in an array attribute's data

        \* set              Set the value of an attribute's data

    Class Attributes:
        class Keys: Helper for managing the keywords needed for the edit() function

    Attributes:
        __path_to_object_map: Mapping from the node or prim path specified to the created node or prim

    Raises:
        OmniGraphError: If the requested operation could not be performed

    """
    # end-controller-docs

    Keys = ogn.GraphSetupKeys

    TYPE_CHECKING = True
    """If True then verbose type checking will happen in various locations so that more legible error messages can
    be emitted. Set it to False to run a little bit faster, with errors just reporting raw Python error messages.
    """

    # --------------------------------------------------------------------------------------------------------------
    def __init__(self, *args, **kwargs):
        """Set up state information. You only need to create an instance of the Controller if you are going to
        use the edit() function more than once, when it needs to remember the node mapping used for creation.
        Args are passed on to the parent classes who have inits and interpreted by them as they see fit.

        Args:
            graph_id: Must be by keyword. If specified then operations are performed on this graph_id unless it is
                      overridden in a particular function call.

        Check the help information for :py:meth:`GraphController.__init__`, :py:meth:`NodeController.__init__`,
        :py:meth:`DataView.__init__`, and :py:meth:`ObjectLookup.__init__` for details on what constructor arguments
        are accepted.
        """
        (self.__graph_id, self.__path_to_object_map, self.__update_usd, self.__undoable) = _flatten_arguments(
            optional=[("graph_id", None), ("path_to_object_map", None), ("update_usd", True), ("undoable", True)],
            args=args,
            kwargs=kwargs,
        )
        GraphController.__init__(self, *args, **kwargs)
        NodeController.__init__(self, *args, **kwargs)
        DataView.__init__(self, *args, **kwargs)
        ObjectLookup.__init__(self)

        # Dual function methods that can be called either from an object or directly from the class
        self.edit = self.__edit_obj
        self.evaluate = self.__evaluate_obj
        self.evaluate_sync = self.__evaluate_sync_obj

    # --------------------------------------------------------------------------------------------------------------
    @dataclass
    class _EditArgs:
        """Collection of shared arguments that are passed to a bunch of the edit implementation methods

        Members:
            path_to_object_map: Dictionary of the mappings of the node names to instantiated node paths
            graph_path: Path location of the graph on which the edits take place
            update_usd: Should editing operations immediately update the USD backing?
            undoable: Should editing operations be undoable?
        """

        path_to_object_map: PathToObjectMap_t = None
        graph_path: str = None
        update_usd: bool = None
        undoable: bool = None

    # --------------------------------------------------------------------------------------------------------------
    @classmethod
    async def evaluate(obj, *args, **kwargs):  # noqa: N804,PLC0202,PLE0202
        """Wait for the next Graph evaluation cycle - await this function to ensure it is finished before returning.

        This function can be called either from the class or using an instantiated object. The first argument is
        positional, being either the class or object.

        .. code-block:: python

            await og.Controller.evaluate()  # Evaluates all graphs
            controller = og.Controller.edit("/TestGraph", {})
            await controller.evaluate()  # Evaluates only "/TestGraph"
            await og.Controller.evaluate("/TestGraph")  # Evaluates only "/TestGraph"
            controller.evaluate(graph_id="/OtherGraph")  # Evaluates only "/OtherGraph" (not its own "/TestGraph")

        Args:
            obj: Either cls or self depending on how the function was called
            graph_id (GraphSpecs_t): Graph or list of graphs to evaluate - None means all existing graphs
        """
        await obj.__evaluate(obj, args=args, kwargs=kwargs)

    async def __evaluate_obj(self, *args, **kwargs) -> Any:
        """Implements :py:meth:`.Controller.evaluate` when called as an object method"""
        await self.__evaluate(self, graph_id=self.__graph_id, args=args, kwargs=kwargs)

    @staticmethod
    async def __evaluate(
        obj,
        graph_id: Optional[GraphSpecs_t] = None,
        args: List[Any] = None,
        kwargs: Dict[str, Any] = None,
    ):
        """Implements :py:meth:`.Controller.evaluate`"""
        (graph_id,) = _flatten_arguments(
            optional=[("graph_id", graph_id)],
            args=args,
            kwargs=kwargs,
        )
        graphs = obj.graph(graph_id)
        if graphs is None:
            graphs = og.get_all_graphs()
        elif not isinstance(graphs, list):
            graphs = [graphs]

        for graph in graphs:
            graph.evaluate()

    # --------------------------------------------------------------------------------------------------------------
    @classmethod
    def evaluate_sync(obj, *args, **kwargs):  # noqa: N804,PLC0202,PLE0202
        """Run the next Graph evaluation cycle immediately.

        This function can be called either from the class or using an instantiated object. The first argument is
        positional, being either the class or object.

        .. code-block:: python

            og.Controller.evaluate_sync()  # Evaluates all graphs
            controller = og.Controller.edit("/TestGraph", {})
            controller.evaluate_sync()  # Evaluates only "/TestGraph"
            og.Controller.evaluate_sync("/TestGraph")  # Evaluates only "/TestGraph"
            controller.evaluate_sync(graph_id="/OtherGraph")  # Evaluates only "/OtherGraph" (not its own "/TestGraph")

        Args:
            obj: Either cls or self depending on how evaluate_sync() was called
            graph_id (GraphSpecs_t): Graph or list of graphs to evaluate - None means all existing graphs
        """
        return obj.__evaluate_sync(obj, *args, **kwargs)

    def __evaluate_sync_obj(self, *args, **kwargs) -> Any:
        """Implements :py:meth:`.Controller.evaluate_sync` when called as an object method"""
        return self.__evaluate_sync(self, *args, **kwargs)

    @staticmethod
    def __evaluate_sync(obj, graph_id: Optional[GraphSpecs_t] = None):
        """Implements :py:meth:`.Controller.evaluate_sync`"""
        created_loop = False
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            created_loop = True
            loop = asyncio.new_event_loop()
        loop.run_until_complete(obj.evaluate(graph_id))
        if created_loop:
            loop.close()

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def __mapped_node_path(node_path: str, path_to_object_map: PathToObjectMap_t) -> str:
        """Returns the node path after subjecting it to the object path mapping"""
        try:
            node = path_to_object_map[node_path]
            return node.get_prim_path() if isinstance(node, og.Node) else str(node.GetPrimPath())
        except KeyError:
            return node_path

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def __mapped_attribute_spec(  # noqa: PLW0238
        attr_spec: AttributeSpec_t, path_to_object_map: PathToObjectMap_t
    ) -> AttributeSpec_t:
        """Returns the attribute spec with any necessary mapping of local path to fully created path made"""
        # Check for Node.Attr format
        if isinstance(attr_spec, og.Attribute):
            return attr_spec
        if isinstance(attr_spec, str):
            paths = attr_spec.split(".")
            # Attribute without node has no mapping
            if len(paths) == 1:
                return attr_spec
            # Remap node path and rejoin with attribute
            if len(paths) == 2:
                return ".".join([Controller.__mapped_node_path(paths[0], path_to_object_map), paths[1]])
            raise og.OmniGraphError(f"Attribute spec can only have one '.' separator - saw '{attr_spec}'")
        if isinstance(attr_spec, (tuple, list)):
            if len(attr_spec) != 2:
                raise og.OmniGraphError(f"Attribute spec should be a (name, node) pair - saw '{attr_spec}'")
            # Make it reversible when it can be inferred
            if isinstance(attr_spec[0], og.Node):
                return (attr_spec[1], attr_spec[0])
            if isinstance(attr_spec[1], og.Node):
                return attr_spec
            return (attr_spec[0], Controller.__mapped_node_path(attr_spec[1], path_to_object_map))

        # Nothing else has a mapping
        return attr_spec

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_create_nodes(
        obj,
        nodes_to_create: List[Tuple[NewNode_t, NodeType_t]],
        edit_args: Controller._EditArgs,
    ):
        """Creates a list of nodes, updating the path_to_object_map to include the new mappings

        Args:
            nodes_to_create: List of nodes that are to be created
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Raises:
            og.OmniGraphError: If any of the nodes could not be created
        """
        if obj.TYPE_CHECKING:
            for element in nodes_to_create:
                type_wrong = (not isinstance(element, (tuple, list))) or len(element) != 2
                type_wrong |= type(element[0]) not in [str, og.Node, Sdf.Path, Usd.Prim]
                type_wrong |= type(element[1]) not in [str, og.NodeType, og.Node, Usd.Prim]
                if type_wrong:
                    raise og.OmniGraphError(
                        f"Node creation spec must be (node_path, node_type) pairs - got '{element}'"
                    )

        # Before creation make sure the graph path is prepended or the mapped location is found to the node
        # path if the path is not absolute
        nodes_constructed = [
            obj.create_node(
                _find_actual_path(node_path, edit_args.graph_path, edit_args.path_to_object_map),
                node_type,
                update_usd=edit_args.update_usd,
                undoable=edit_args.undoable,
            )
            for node_path, node_type in nodes_to_create
        ]
        just_names = [name for (name, _node_type) in nodes_to_create]
        edit_args.path_to_object_map.update(dict(zip(just_names, nodes_constructed)))

        return nodes_constructed

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_delete_nodes(
        obj,
        nodes_to_delete: List[NodeSpec_t],
        edit_args: Controller._EditArgs,
    ):
        """Deletes a list of nodes, updating the path_to_object_map to remove any that are no longer in it

        Args:
            nodes_to_delete: List of nodes that are to be deleted
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Raises:
            og.OmniGraphError: If the nodes could not be found or could not be deleted
        """
        if obj.TYPE_CHECKING:
            for element in nodes_to_delete:
                if isinstance(element, og.Node) and not element.is_valid():
                    raise og.OmniGraphError(f"Node deletion spec points to an invalid node - '{element}'")
                if (
                    not isinstance(element, str)
                    and not isinstance(element, Usd.Prim)
                    and not isinstance(element, og.Node)
                ):
                    raise og.OmniGraphError(f"Node deletion spec must be a string, node, or prim - got '{element}'")

        for element in set(nodes_to_delete):
            # Make sure the path to object map no longer has the reference to the deleted node
            if isinstance(element, og.Node):
                for node_path, node in edit_args.path_to_object_map.items():
                    if node == element:
                        del edit_args.path_to_object_map[node_path]
                        break
                obj.delete_node(element, update_usd=edit_args.update_usd, undoable=edit_args.undoable)
            elif isinstance(element, Usd.Prim):
                omnigraph_node = og.get_node_by_path(element.GetPrimPath())
                if omnigraph_node is not None and omnigraph_node.is_valid():
                    for node_path, node in edit_args.path_to_object_map.items():
                        if node == omnigraph_node:
                            del edit_args.path_to_object_map[node_path]
                            break
                obj.delete_node(element, update_usd=edit_args.update_usd, undoable=edit_args.undoable)
            else:
                if element in edit_args.path_to_object_map:
                    node_path = og.ObjectLookup.node_path(edit_args.path_to_object_map[element])
                    del edit_args.path_to_object_map[element]
                    obj.delete_node(node_path, update_usd=edit_args.update_usd, undoable=edit_args.undoable)
                else:
                    obj.delete_node(
                        f"{edit_args.graph_path}/{element}",
                        update_usd=edit_args.update_usd,
                        undoable=edit_args.undoable,
                    )

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_connections(
        obj,
        connection_definitions: List[Tuple[AttributeSpec_t, AttributeSpec_t]],
        breaking_connections: bool,
        edit_args: Controller._EditArgs,
    ):
        """Process the edit section connections to be made or broken

        Args:
            connection_definitions: List of (src, dst) attribute pairs of the connections to be processed
            breaking_connections: If True then disconnect the pairs, otherwise connect the pairs
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Raises:
            og.OmniGraphError: If any of the attributes could not be found or the connections failed
        """
        for src_spec, dst_spec in connection_definitions:
            src_attr = obj.__mapped_attribute_spec(src_spec, edit_args.path_to_object_map)  # noqa: PLW0212
            dst_attr = obj.__mapped_attribute_spec(dst_spec, edit_args.path_to_object_map)  # noqa: PLW0212
            if breaking_connections:
                obj.disconnect(src_attr, dst_attr, update_usd=edit_args.update_usd, edit_args=edit_args.undoable)
            else:
                obj.connect(src_attr, dst_attr, update_usd=edit_args.update_usd, edit_args=edit_args.undoable)

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_disconnect_all(obj, attribute_specs: AttributeSpecs_t, edit_args: Controller._EditArgs):
        """Process the edit section connections to be made or broken

        Args:
            path_to_object_map: The map in use with local path names to created objects
            attribute_specs: List of attribute from which all connections are to be broken
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Raises:
            og.OmniGraphError: If the attribute could not be found or the disconnect failed
        """
        attribute_specs = attribute_specs if isinstance(attribute_specs, list) else [attribute_specs]
        for attribute_spec in attribute_specs:
            attribute = (obj.__mapped_attribute_spec(attribute_spec, edit_args.path_to_object_map),)  # noqa: PLW0212
            obj.disconnect_all(attribute, update_usd=edit_args.update_usd, undoable=edit_args.undoable)

    # --------------------------------------------------------------------------------------------------------------
    PrimCreationData_t = Union[Tuple[Path_t, PrimAttrs_t], Tuple[Path_t, str], Tuple[Path_t, PrimAttrs_t, str]]

    @staticmethod
    def _process_create_prims(
        obj,
        prim_definitions: List[PrimCreationData_t],
        edit_args: Controller._EditArgs,
    ) -> List[Usd.Prim]:
        """Process the edit section specifying prim creation

        Args:
            prim_definitions: List of prim paths with optional attribute values and prim type to create on the prim
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Returns:
            List of prims that were created that correspond to the definitions

        Raises:
            og.OmniGraphError if there was a problem with the prim path or attribute definition
        """

        def __create_one_prim(prim_definition: obj.PrimCreationData_t) -> Usd.Prim:
            """Normalize the inputs to turn them all into types suitable for calling create_prim with"""
            prim_path = None
            prim_attributes = {}
            prim_type = None
            if isinstance(prim_definition, (str, Sdf.Path)):
                prim_path = prim_definition
            elif isinstance(prim_definition, (tuple, list)) and isinstance(prim_definition[0], (str, Sdf.Path)):
                if len(prim_definition) == 2:
                    if isinstance(prim_definition[1], str):
                        (prim_path, prim_type) = prim_definition
                    elif isinstance(prim_definition[1], dict):
                        (prim_path, prim_attributes) = prim_definition
                elif (
                    len(prim_definition) == 3
                    and isinstance(prim_definition[1], dict)
                    and isinstance(prim_definition[2], str)
                ):
                    (prim_path, prim_attributes, prim_type) = prim_definition
            if prim_path is None:
                raise og.OmniGraphError(
                    f"Prim definitions must be name, (name,values), (name,type) or (name,values,type) -"
                    f" '{prim_definition}' is not recognized"
                )

            return obj.create_prim(
                _find_actual_path(str(prim_path), "", {}), prim_attributes, prim_type, undoable=edit_args.undoable
            )

        # Loop through all prim definitions, creating them as we go
        return [__create_one_prim(prim_definition) for prim_definition in prim_definitions]

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_expose_prims(
        obj,
        exposure_definitions: List[GraphController.ExposePrimNode_t],
        edit_args: Controller._EditArgs,
    ) -> List[og.Node]:
        """Process the edit section specifying prim creation

        Args:
            exposure_definitions: List of (exposure_type, prim, node_path) mappings giving the type of exposure,
                                  prim to expose, and new node path at which it will be exposed
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Returns:
            List of nodes that were created to expose the specified prims

        Raises:
            og.OmniGraphError if there was a problem with the prim or node path
        """
        nodes_constructed = [
            obj.expose_prim(
                exposure_type,
                _find_actual_path(prim_id, "", edit_args.path_to_object_map) if isinstance(prim_id, str) else prim_id,
                _find_actual_path(node_path, edit_args.graph_path, edit_args.path_to_object_map),
                update_usd=edit_args.update_usd,
                undoable=edit_args.undoable,
            )
            for (exposure_type, prim_id, node_path) in exposure_definitions
        ]
        just_names = [node_path for (_, _, node_path) in exposure_definitions]
        edit_args.path_to_object_map.update(dict(zip(just_names, nodes_constructed)))

        return nodes_constructed

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_set_values(obj, value_definitions: AttributeValues_t, edit_args: Controller._EditArgs):
        """Process the edit section specifying value setting

        Args:
            value_definitions: List of (AttributeSpec_t, value) tuples indicating the attributes and the values to set
                               on them. The attribute spec will have the path_to_object_map applied to them if they
                               contain a relative node path. Optionally it will be a 3-tuple where the third member
                               is an AttributeType_t that specifies a resolve type for the attribute. This is only
                               valid for extended attribute types.
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Raises:
            og.OmniGraphError: If the attributes could not be found, the values were invalid, or setting failed
        """

        def _process_set_value(attribute_id: AttributeSpec_t, value: ValueToSet_t, type_info: AttributeType_t = None):
            """Sets a single value on an attribute"""
            attribute_spec = obj.__mapped_attribute_spec(attribute_id, edit_args.path_to_object_map)  # noqa: PLW0212
            if type_info is not None:
                value = og.TypedValue(value, type_info)
            obj.set(
                og.ObjectLookup.attribute(attribute_spec),
                value,
                update_usd=edit_args.update_usd,
                undoable=edit_args.undoable,
            )

        if isinstance(value_definitions, list):
            _ = [_process_set_value(*value_definition) for value_definition in value_definitions]
        else:
            _process_set_value(*value_definitions)

    # --------------------------------------------------------------------------------------------------------------
    @staticmethod
    def _process_create_variable(
        obj, variable_definitions: List[Tuple[VariableName_t, VariableType_t]], edit_args: Controller._EditArgs
    ) -> List[og.IVariable]:
        """Process the edit section that creates new variables

        Args:
            variable_definition: List of descriptions for variables to be created
            edit_args: The common arguments used to configure editing operations :py:class:`Controller._EditArgs`

        Returns:
            List of variables created
        """
        return [
            obj.create_variable(edit_args.graph_path, name, var_type, undoable=edit_args.undoable)
            for (name, var_type) in variable_definitions
        ]

    # --------------------------------------------------------------------------------------------------------------
    @classmethod
    def edit(  # noqa: PLC0202,PLE0202
        obj, *args, **kwargs  # noqa: N804
    ) -> Tuple[og.Graph, List[og.Node], List[Usd.Prim], PathToObjectMap_t]:
        """Edit and/or create an OmniGraph from the given description.

        This function provides a single call that will make a set of modifications to an OmniGraph. It can be used to
        create a new graph from scratch or to make changes to an existing graph.

        If the "undoable" mode is not set to False then a single undo will revert everything done via this call.

        The description below contains different sections that perform different operations on the graph. They are
        always done in the order listed to minimize conflicts. If you need to execute them in a different order then use
        multiple calls to edit().

        This function can be called either from the class or using an instantiated object. When callling from an object
        context the arguments passed in will take precedence over the arguments provided to the constructor. Here are
        some of the legal ways to call the edit function:

        .. code-block:: python

            # For example purposes "cmds()" is a function that returns a dictionary of editing commands

            (graph, _, _, _) = og.Controller.edit("/TestGraph", cmds())

            new_controller = og.Controller(graph)
            new_controller.edit(cmds())

            og.Controller.edit(graph, cmds())
            new_controller.edit(graph, cmds())
            new_controller.edit(graph, cmds(), undoable=False)  # Overrides the normal undoable state of new_controller

        Below is the list of the allowed operations, in the order in which they will be performed.

        The parameters are described as lists, though if you have only one parameter for a given operation you can pass
        it without it being in a list.

        .. code-block:: python

            { OPERATION: [List, Of, Arguments] }
            { OPERATION: SingleArgument }

        For brevity the shortcut "keys = og.Controller.Keys" is assumed to exist.

            - keys.DELETE_NODES: NodeSpecs_t

                Deletes a node or list of nodes. If the node specification is a relative path then it must be in the
                list of full paths that the controller created in a previous call to edit().

                { keys.DELETE_NODES: ["NodeInGraph", "/Some/Other/Graph/Node", my_omnigraph_node] }

            - keys.CREATE_NODES: [(Path_t, NodeType_t)]

                Constructs a node of the given type at the given path. If the path is a relative one then it is
                added directly under the graph being edited. A map is remembered between the given path, relative or
                absolute, and the node created at it so that further editing functions can refer to the node by that
                name directly rather than trying to infer the final full name.

                { keys.CREATE_NODES: [("NodeInGraph", "omni.graph.tutorial.SimpleData"),
                                      ("Inner/Node/Path", og.NodeType(node_type_name))] }

            - keys.CREATE_PRIMS: [(Path_t, {ATTR_NAME: (AttributeType_t, ATTR_VALUE)}, Optional(PRIM_TYPE))]

                Constructs a prim at path "PRIM_PATH" containing a set of attributes with specified types and values.
                Only those attribute types supported by USD can be used here, though the type specification can be in
                OGN form - invalid types result in an error. Whereas relative paths on nodes are treated as being
                relative to the graph, for prims a relative path is relative to the stage root. Prims are not allowed
                inside an OmniGraph and attempts to create one there will result in an error.

                Note that the PRIM_TYPE can appear with or without an attribute definition. (Many prim types are part
                of a schema and do not require explicit attributes to be added.)

                { keys.PRIMS: [("/World/Prim", {"speed": ("double", 1.0)}),
                               ("/World/Cube", "Cube"),
                               ("RootPrim", {
                               "mass": (Type(BaseDataType.DOUBLE), 3.0),
                               "force:gravity": ("double", 32.0)
                               })]}

            - keys.CONNECT: [(AttributeSpec_t, AttributeSpec_t)]

                Makes a connection between the given source and destination attributes. The local name of a newly
                created node may be made as part of the node in the spec, or the node portion of the attribute path:

                { keys.CONNECT: [("NodeInGraph.outputs:value", ("inputs:value", "NodeInGraph"))]}

            - keys.DISCONNECT: [(AttributeSpec_t, AttributeSpec_t)]

                Breaks a connection between the given source and destination attributes. The local name of a newly
                created node may be made as part of the node in the spec, or the node portion of the attribute path:

                { keys.DISCONNECT: [("NodeInGraph.outputs:value", ("inputs:value", "NodeInGraph"))]}

            - keys.EXPOSE_PRIMS: [(cls.PrimExposureType, Prim_t, NewNode_t)]

                Exposes a prim to OmniGraph through creation of one of the node types designed to do that.
                The first member of the tuple is the method used to expose the prim. The prim path is
                the second member of the tuple and it must already exist in the USD stage. The third member
                of the tuple is a node name with the same restrictions as the name in the CREATE_NODES edit.

                { keys.EXPOSE_PRIMS: [(cls.PrimExposureType.AS_BUNDLE, "/World/Cube", "BundledCube")] }

            - keys.SET_VALUES: [AttributeSpec_t, Any] or [AttributeSpec_t, Any, AttributeType_t]

                Sets the value of the given list of attributes.

                { keys.SET_VALUES[("/World/Graph/Node.inputs:attr1", 1.0), (("inputs:attr2", node), 2.0)] }

                In the case of extended attribute types you may also need to supply a data type, which is the type the
                attribute will resolve to when the value is set. To supply a type, add it as a third parameter to the
                attribute value:

                { keys.SET_VALUES[("inputs:ext1", node), 2.0, "float"] }

            - keys.CREATE_VARIABLES: [(VariableName_t, VariableType_t)]

                Constructs a variable on the graph with the given name and type. The type can be specified
                as either a ogn type string (e.g. "float4"), or as an og.Type (e.g. og.Type(og.BaseDataType.FLOAT))

                { keys.CREATE_VARIABLES[("velocity", "float3"), ("count", og.Type(og.BaseDataType.INT))] }

        Here's a simple call that first deletes an existing node "/World/PushGraph/OldNode", then creates two nodes of
        type "omni.graph.tutorials.SimpleData", connects their "a_int" attributes, disconnects their "a_float"
        attributes and sets the input "a_int" of the source node to the value 5. It also creates two unused USD Prim
        nodes, one with a float attribute named "attrFloat" with value 2.0, and the other with a boolean attribute
        named "attrBool" with the value true.

        .. code-block:: python

            controller = og.Controller()
            keys = og.Controller.Keys
            (graph, nodes_constructed, prims_constructed, path_to_object_map) = controller.edit("/World/PushGraph", {
                keys.DELETIONS: [
                    "OldNode"
                ],
                keys.CREATE_NODES: [
                    ("src", "omni.graph.tutorials.SimpleData"),
                    ("dst", "omni.graph.tutorials.SimpleData")
                    ],
                keys.CREATE_PRIMS: [
                    ("Prim1", {"attrFloat": ("float", 2.0)),
                    ("Prim2", {"attrBool": ("bool", true)),
                    ],
                keys.CONNECT: [
                    ("src.outputs:a_int", "dst.inputs:a_int")
                    ],
                keys.DISCONNECT: [
                    ("src.outputs:a_float", "dst.inputs:a_float")
                    ],
                keys.SET_VALUES: [
                    ("src.inputs:a_int", 5)
                    ]
                keys.CREATE_VARIABLES: [
                    ("a_float_var", og.Type(og.BaseDataType.FLOAT)),
                    ("a_bool_var", "bool")
                ]
                }
            )

        .. note::

            The controller object remembers where nodes are created so that you can use short forms for the node paths
            for convenience. That's why in the above graph the node paths for creation and the node specifications in
            the connections just says "src" and "dst". As they have no leading "/" they are treated as being relative
            to the graph path and will actually end up in "/World/PushGraph/src" and "/World/PushGraph/dst".

            Node specifications with a leading "/" are assumed to be absolute paths and must be inside the path of an
            existing or created graph. e.g. the NODES reference could have been "/World/PushGraph/src", however using
            "/src" would have been an error.

            This node path mapping is remembered across multiple calls to edit() so you can always use the shortform
            so long as you use the same Controller object.

        Args:
            obj: Either cls or self depending on how edit() was called
            graph_id: Identifier that says which graph is being edited. See :py:meth:`GraphController.create_graph`
                      for the data types accepted for the graph description.
            edit_commands: Dictionary of commands and parameters indicating what modifications are to be made
                           to the specified graph. A strict set of keys is accepted. Each of the values in the
                           dictionary can be either a single value or a list of values of the proscribed type.
            path_to_object_map: Dictionary of relative paths mapped on to their full path after creation so that the
                                edit_commands can use either full paths or short-forms to specify the nodes.
            update_usd: If specified then override whether to update the USD after the operations (default False)
            undoable: If True the operation is added to the undo queue, else it is done immediately and forgotten
                      (default True)

        Returns:
            A 4-tuple consisting of:
                - the og.Graph being used for the operation
                - the list of og.Nodes created by the operation
                - the list of Usd.Prims created by the operation
                - the map of node/prim path name to the created og.Node/Usd.Prim objects. Can usually be ignored; it
                  will be used internally to manage the shortform names of the paths used in multiple commands.

        Raises:
            og.OmniGraphError if any of the graph creation instructions could not be fulfilled
            The graph will be left in the partially constructed state it reached at the time of the error
        """
        return obj.__edit(obj, args=args, kwargs=kwargs)

    def __edit_obj(self, *args, **kwargs):
        """Implements :py:meth:`Controller.edit` as an object method"""
        results = self.__edit(
            self,
            graph_id=self.__graph_id,
            path_to_object_map=self.__path_to_object_map,
            update_usd=self.__update_usd,
            undoable=self.__undoable,
            args=args,
            kwargs=kwargs,
        )
        self.__graph_id = results[0]
        self.__path_to_object_map = results[3]
        return results

    @staticmethod
    def __edit(
        obj,
        graph_id: Union[GraphSpec_t, Dict[str, Any]] = _Unspecified,
        edit_commands: Optional[Dict[str, Any]] = None,
        path_to_object_map: PathToObjectMap_t = None,
        update_usd: bool = True,
        undoable: bool = True,
        args: List[str] = None,
        kwargs: Dict[str, Any] = None,
    ) -> Tuple[og.Graph, List[og.Node], List[Usd.Prim]]:
        """Implements :py:meth:`Controller.edit`"""
        (graph_id, edit_commands, path_to_object_map, update_usd, undoable) = _flatten_arguments(
            mandatory=[("graph_id", graph_id)],
            optional=[
                ("edit_commands", edit_commands),
                ("path_to_object_map", path_to_object_map),
                ("update_usd", update_usd),
                ("undoable", undoable),
            ],
            args=args,
            kwargs=kwargs,
        )
        if edit_commands is None:
            edit_commands = {}
        if path_to_object_map is None:
            path_to_object_map = {}

        # If the undo or usd update states were specified then add them to the arguments everywhere
        shared_kwargs = Controller._EditArgs(
            update_usd=update_usd,
            undoable=undoable,
            path_to_object_map=path_to_object_map,
        )

        # Separate the actual work so that it can be configured before being called
        def __do_the_edits():
            # A dictionary ID always indicates the graph should be created, otherwise check if it already exists
            graph = obj.graph(graph_id) if not isinstance(graph_id, dict) else None
            nodes_constructed = []
            prims_constructed = []

            # If the graph couldn't be found then infer that it should be created with the given specification.
            if graph is None:
                graph = obj.create_graph(graph_id)

            # The graph could be in an illegal state when halfway through editing operations. This will disable it
            # until all operations are completed. It's up to the caller to ensure that the state is legal after all
            # operations are completed.
            graph_was_disabled = graph.is_disabled()
            graph.set_disabled(True)
            graph_path = graph.get_path_to_graph()
            shared_kwargs.graph_path = graph_path

            try:
                # Syntax check first
                for instruction, _data in edit_commands.items():
                    if instruction not in obj.Keys.ALL:
                        raise OmniGraphError(f"Unknown graph edit operation - `{instruction}` not in {obj.Keys.ALL}")

                # Delete first because we may be creating nodes with the same names
                with suppress(KeyError):
                    obj._process_delete_nodes(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.DELETE_NODES]),
                        edit_args=shared_kwargs,
                    )

                # Variables next, so they exist when nodes are created
                with suppress(KeyError):
                    obj._process_create_variable(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.CREATE_VARIABLES]),
                        edit_args=shared_kwargs,
                    )

                # Create nodes next since connect and set may need them
                nodes_constructed = []
                with suppress(KeyError):
                    nodes_constructed = obj._process_create_nodes(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.CREATE_NODES]),
                        edit_args=shared_kwargs,
                    )

                # Prims next as they may be used in connections
                with suppress(KeyError):
                    prims_constructed = obj._process_create_prims(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.CREATE_PRIMS]),
                        edit_args=shared_kwargs,
                    )

                # Exposure of a raw prim to OmniGraph through one of the prim interface nodes
                with suppress(KeyError):
                    nodes_constructed += obj._process_expose_prims(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.EXPOSE_PRIMS]),
                        edit_args=shared_kwargs,
                    )

                # Connections next as setting values may override their data
                with suppress(KeyError):
                    obj._process_connections(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.CONNECT]),
                        breaking_connections=False,
                        edit_args=shared_kwargs,
                    )

                # Disconnections may have been created by the connections list so they're next, though that's unlikely
                with suppress(KeyError):
                    obj._process_connections(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.DISCONNECT]),
                        breaking_connections=True,
                        edit_args=shared_kwargs,
                    )

                # Disconnections of all attributes is last for connection changes as it will read existing connections
                with suppress(KeyError):
                    obj._process_disconnect_all(  # noqa: PLW0212
                        obj,
                        _get_as_list(edit_commands[obj.Keys.DISCONNECT_ALL]),
                        edit_args=shared_kwargs,
                    )

                # Now that everything is in place it is safe to set the values
                # TODO: Setting values has two parameters (on_gpu, and update_usd) that are not accounted for here.
                #       Extra information should be provided to allow for them. Until then the defaults will be used.
                with suppress(KeyError):
                    obj._process_set_values(  # noqa: PLW0212
                        obj, _get_as_list(edit_commands[obj.Keys.SET_VALUES]), edit_args=shared_kwargs
                    )
            finally:
                # Really important that this always gets reset
                graph.set_disabled(graph_was_disabled)

            return (graph, nodes_constructed, prims_constructed, path_to_object_map)

        # Put every operation into a single undo umbrella. They will handle queueing undoable operations internally.
        if undoable:
            with undo.group():
                return __do_the_edits()

        return __do_the_edits()
