graph

@dataclass
class AbstractSerializable(abc.ABC):
    """An abstract serializable container."""

    name: str = field(default_factory=str)
    description: str = field(default_factory=str)

    def export_to_file(self, filepath: str):
        with open(filepath, "w", encoding="utf-8") as file_handler:
            file_handler.write(self.to_json())

    def to_json(self, sorted_keys: Optional[bool] = False) -> str:
        as_primitive = serialization.serialize_to_primitive(self)
        return json.dumps(as_primitive, indent=4, sort_keys=sorted_keys)

    @classmethod
    def import_from_file(cls, filepath: str) -> Self:
        data = pathlib.Path(filepath).read_text(encoding="utf-8")
        return cls.from_json(data)

    @classmethod
    def from_json(cls, json_str: str) -> Self:
        data = json.loads(json_str)
        return cls.from_dict(data)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> Any:
        return serialization.deserialize_from_primitive(data)

@dataclass
class DispatchGraph(Graph):
    """A Graph that only stores dispatch nodes for submission."""

    all_nodes: List[AbstractDispatchTask] = field(default_factory=list)
    default_task_cls: str = field(default_factory=str)  # Store class name

    def get_associated_manager(self) -> str:
        managers = {
            node.get_associated_manager()
            for node in self.all_nodes
        }
        if len(managers) == 1:
            return managers.pop()

        # Either no value or multiple values for custom managers
        # within the dispatch graph. This is unexpected.
        raise RuntimeError(
            "Ambiguous dispatch graph, cannot"
            f"identify render manager from {managers}."
        )

    def create_node(
        self,
        dispatch_node_type: str,
        label: Optional[str] = None
    ) -> AbstractDispatchTask:
        """Create a container from a specific metadata container type."""
        meta_types = {
            klass.__name__: klass
            for klass in get_dispatch_node_types()
        }

        # Initialize new metadata container.
        if dispatch_node_type not in meta_types:
            raise ValueError(
                f"Invalid dispatch node type {dispatch_node_type}. "
                f"Available types are: {meta_types}."
            )

        meta_type = meta_types[dispatch_node_type]
        new_container_name = self._get_new_name(dispatch_node_type)
        new_container = meta_type(
            name=new_container_name,
            label=label,
            class_name=dispatch_node_type,
        )
        self.all_nodes.append(new_container)
        return new_container

    def delete_node(
        self,
        dispatch_node: AbstractDispatchTask,
    ):
        """Delete a dispatch node from the graph."""
        if not isinstance(dispatch_node, AbstractDispatchTask):
            raise TypeError(f"Invalid node provided {dispatch_node}")

        if dispatch_node not in self.all_nodes:
            raise ValueError(
                f"Invalid item to remove. {dispatch_node} "
                "is not part of the graph."
            )

        self.all_nodes.remove(dispatch_node)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "DispatchGraph":
        data = serialization.deserialize_from_primitive(data)

        # Backward compatible with previous serialization method.
        # TODO remove legacy serialization format support in a future version.
        if isinstance(data, dict):
            logger.warning(
                "DispatchGraph.from_dict: using legacy serialization format"
            )
            dispatch_type_map = {
                klass.__name__: klass
                for klass in get_dispatch_node_types()
            }
            all_nodes = []
            for meta_node_data in data.get("all_nodes", []):
                class_name = meta_node_data.get("class_name") or next(
                    (name for name in dispatch_type_map
                     if name in meta_node_data["name"]),
                    None
                )
                klass = (
                    dispatch_type_map.get(class_name)
                    if class_name else None
                )
                if klass is None:
                    raise ValueError(
                        f"Unsupported node data: {meta_node_data}."
                    )
                all_nodes.append(klass.from_dict(meta_node_data))

            return cls(
                name=data.get("name"),
                description=data.get("description"),
                all_nodes=all_nodes,
                default_task_cls=data.get("default_task_cls", ""),
                ui=data.get("ui"),
            )

        return data

@dataclass
class Graph(AbstractSerializable):
    """A container of nodes."""

    all_nodes: List[Node] = field(default_factory=list)

    ui: Optional[ui_utils.UiGraphMetadata] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __str__(self) -> str:
        return f"<Graph {self.name} {len(self.all_nodes)} node(s)>"

    def get_nodes(
        self,
        node_types: Union[None, List[str]] = None,
    ) -> List[Node]:
        """Return all nodes with potential filter on type."""
        if node_types:
            return [
                node for node in self.all_nodes if node.node_type in node_types
            ]

        return copy.copy(self.all_nodes)

    def get_node_by_name(self, name: str) -> Node:
        for node in self.all_nodes:
            if node.name == name:
                return node
        raise ValueError(f"No node named {name} in graph.")

    def create_node(
            self,
            node_type: str,
            label: Optional[str] = None
        ) -> Node:
        """Create a node from a specific node type."""
        plugins = PluginRegistry().plugin_list()

        # Initialize new node from plugin desc.
        if node_type not in plugins:
            raise ValueError(
                f"Invalid node type {node_type}. "
                f"Available types are: {tuple(plugins)}"
            )

        new_node_name = self._get_new_name(node_type)
        new_node = Node(
            name=new_node_name,
            plugin_ref=PluginRef(
                name=node_type,
                version=plugins[node_type].version,
            ),
            label=label,
        )
        self.all_nodes.append(new_node)
        return new_node

    def delete_node(self, node: Node):
        """Delete a node from the graph."""
        if not isinstance(node, Node):
            raise TypeError(f"Invalid node provided {node}")

        # Check that node is present.
        elif node not in self.all_nodes:
            raise ValueError(
                f"Invalid node to remove. {node} is not part of the graph."
            )

        # Execution node: remove and adjust graph connections.
        to_remove = self.all_nodes.pop(self.all_nodes.index(node))

        # Remove potential broken connections.
        for other_node in self.all_nodes:
            for input_name in other_node.inputs:
                # Skip if to_remove is not connected this other node.
                conn_values = other_node[input_name]
                if not isinstance(conn_values, list):
                    conn_values = [conn_values]

                if (
                    conn_values == []
                    or not isinstance(conn_values[0], NodeConnection)
                ):
                    continue

                for conn in conn_values:
                    connected_node_name = conn.origin_node_name
                    if connected_node_name == to_remove.name:
                        other_node.disconnect_input(input_name)

        # Delete node.
        del to_remove

    def _get_new_name(self, item_type: str) -> str:
        """Compute a new meaningful yet unique item name."""
        similar_nodes = self.get_nodes(node_types=[item_type])

        if similar_nodes:
            regex = r"\D*(?P<index>\d*)"
            indexes = [
                re.match(regex, node.name).groupdict()["index"]
                for node in similar_nodes
            ]
            node_idx = max(map(int, indexes)) + 1

        else:
            node_idx = 1

        # Technically node name just have to be unique, stuff such
        # as "NodeType3" makes it easier to troubleshoot.
        return f"{item_type}{node_idx}"

    def duplicate_nodes(self, nodes: List[Node]) -> List[Node]:
        """ Duplicate provided nodes in the graph.
        """
        new_nodes = []
        node_names = []
        graph_node_names = {node.name for node in self.all_nodes}

        # duplicate nodes
        for node in nodes:
            node_names.append(node.name)
            new_node = copy.deepcopy(node)
            new_node.name = self._get_new_name(node.node_type)
            self.all_nodes.append(new_node)
            new_nodes.append(new_node)

        node_idx_by_name = {name: idx for idx, name in enumerate(node_names)}

        # duplicate connections
        for new_node in new_nodes:
            for input_name, conns in new_node.input_connections.items():
                is_multi_input = True
                if not isinstance(conns, list):  # handle non-multi inputs
                    is_multi_input = False
                    conns = [conns]

                idx = 0
                for input_conn in conns:
                    if (
                        isinstance(input_conn, NodeConnection)
                        and input_conn.origin_node_name not in graph_node_names
                    ):
                        if is_multi_input:
                            new_node.input_connections[input_name].pop(idx)
                            continue
                        else:
                            new_node.input_connections[input_name] = None
                            break

                    new_conn = copy.copy(input_conn)
                    if (
                        isinstance(input_conn, NodeConnection)
                        and input_conn.origin_node_name in node_idx_by_name
                    ):
                        node_idx = node_idx_by_name[
                            input_conn.origin_node_name
                        ]
                        new_conn.origin_node_name = new_nodes[node_idx].name

                    if not is_multi_input:
                        new_node.input_connections[input_name] = new_conn
                        break
                    new_node.input_connections[input_name][idx] = new_conn
                    idx += 1

        return new_nodes

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Graph":
        data = super().from_dict(data)

        # Backward compatible with previous serialization method.
        # TODO remove legacy serialization format support in a future version.
        if isinstance(data, dict):
            logger.warning(
                "Graph.from_dict: using legacy serialization format"
            )
            all_nodes = []
            for node_data in data.get("all_nodes", []):
                try:
                    all_nodes.append(Node.from_dict(node_data))

                except TypeError as error:
                    raise ValueError(
                        f"Unsupported node data: {node_data}."
                    ) from error

            return cls(
                name=data.get("name"),
                description=data.get("description"),
                all_nodes=all_nodes,
                ui=data.get("ui"),
            )

        return data