node_scene.py

# -*- coding: utf-8 -*-
"""
A module containing the representation of the NodeEditor's Scene
"""
from collections import OrderedDict
from nodeeditor.utils_no_qt import dumpException, pp
from nodeeditor.node_serializable import Serializable
from nodeeditor.node_graphics_scene import QDMGraphicsScene
from nodeeditor.node_node import Node
from nodeeditor.node_edge import Edge
from nodeeditor.node_scene_history import SceneHistory
from nodeeditor.node_scene_clipboard import SceneClipboard
import keyvalues3 as kv3
from src.widgets.common import show_message
from src.vsmart.remapping import vsmart_remapping


DEBUG_REMOVE_WARNINGS = False


class InvalidFile(Exception): pass


class Scene(Serializable):
    """Class representing NodeEditor's `Scene`"""
    historyClass = SceneHistory
    clipboardClass = SceneClipboard

    def __init__(self):
        """
        :Instance Attributes:

            - **nodes** - list of `Nodes` in this `Scene`
            - **edges** - list of `Edges` in this `Scene`
            - **history** - Instance of :class:`~nodeeditor.node_scene_history.SceneHistory`
            - **clipboard** - Instance of :class:`~nodeeditor.node_scene_clipboard.SceneClipboard`
            - **scene_width** - width of this `Scene` in pixels
            - **scene_height** - height of this `Scene` in pixels
        """
        super().__init__()
        self.nodes = []
        self.edges = []

        # Variable storage for serialization
        self.m_Variables = []

        # current filename assigned to this scene
        self.filename = None

        self.scene_width = 64000
        self.scene_height = 64000

        # custom flag used to suppress triggering onItemSelected which does a bunch of stuff
        self._silent_selection_events = False
        
        # flag to disable output node validation during deserialization
        self._deserializing = False

        self._has_been_modified = False
        self._last_selected_items = None

        # initialize all listeners
        self._has_been_modified_listeners = []
        self._item_selected_listeners = []
        self._items_deselected_listeners = []

        # here we can store callback for retrieving the class for Nodes
        self.node_class_selector = None
        
        # element ID generator for nodes
        self.element_id_generator = None
        
        # reference to variables manager for serialization
        self.variables_manager = None

        self.initUI()
        self.history = self.historyClass(self)
        self.clipboard = self.clipboardClass(self)

        self.grScene.itemSelected.connect(self.onItemSelected)
        self.grScene.itemsDeselected.connect(self.onItemsDeselected)

    @property
    def has_been_modified(self):
        """
        Has this `Scene` been modified?

        :getter: ``True`` if the `Scene` has been modified
        :setter: set new state. Triggers `Has Been Modified` event
        :type: ``bool``
        """
        return self._has_been_modified

    @has_been_modified.setter
    def has_been_modified(self, value):
        if not self._has_been_modified and value:
            # set it now, because we will be reading it soon
            self._has_been_modified = value

            # call all registered listeners
            for callback in self._has_been_modified_listeners: callback()

        self._has_been_modified = value

    def initUI(self):
        """Set up Graphics Scene Instance"""
        self.grScene = QDMGraphicsScene(self)
        self.grScene.setGrScene(self.scene_width, self.scene_height)

    def getNodeByID(self, node_id: int):
        """
        Find node in the scene according to provided `node_id`

        :param node_id: ID of the node we are looking for
        :type node_id: ``int``
        :return: Found ``Node`` or ``None``
        """
        for node in self.nodes:
            if node.id == node_id:
                return node
        return None


    def setSilentSelectionEvents(self, value: bool=True):
        """Calling this can suppress onItemSelected events to be triggered. This is useful when working with clipboard"""
        self._silent_selection_events = value

    def onItemSelected(self, silent: bool=False):
        """
        Handle Item selection and trigger event `Item Selected`

        :param silent: If ``True`` scene's onItemSelected won't be called and history stamp not stored
        :type silent: ``bool``
        """
        if self._silent_selection_events: return

        current_selected_items = self.getSelectedItems()
        if current_selected_items != self._last_selected_items:
            self._last_selected_items = current_selected_items
            if not silent:
                # we could create some kind of UI which could be serialized,
                # therefore first run all callbacks...
                for callback in self._item_selected_listeners: callback()
                # and store history as a last step always
                self.history.storeHistory("Selection Changed")

    def onItemsDeselected(self, silent: bool=False):
        """
        Handle Items deselection and trigger event `Items Deselected`

        :param silent: If ``True`` scene's onItemsDeselected won't be called and history stamp not stored
        :type silent: ``bool``
        """
        # somehow this event is being triggered when we start dragging file outside of our application
        # or we just loose focus on our app? -- which does not mean we've deselected item in the scene!
        # double check if the selection has actually changed, since
        current_selected_items = self.getSelectedItems()
        if current_selected_items == self._last_selected_items:
            # print("Qt itemsDeselected Invalid Event! Ignoring")
            return

        self.resetLastSelectedStates()
        if current_selected_items == []:
            self._last_selected_items = []
            if not silent:
                self.history.storeHistory("Deselected Everything")
                for callback in self._items_deselected_listeners: callback()


    def isModified(self) -> bool:
        """Is this `Scene` dirty aka `has been modified` ?

        :return: ``True`` if `Scene` has been modified
        :rtype: ``bool``
        """
        return self.has_been_modified

    def getSelectedItems(self) -> list:
        """
        Returns currently selected Graphics Items

        :return: list of ``QGraphicsItems``
        :rtype: list[QGraphicsItem]
        """
        return self.grScene.selectedItems()

    def doDeselectItems(self, silent: bool=False) -> None:
        """
        Deselects everything in scene

        :param silent: If ``True`` scene's onItemsDeselected won't be called
        :type silent: ``bool``
        """
        for item in self.getSelectedItems():
            item.setSelected(False)
        if not silent:
            self.onItemsDeselected()

    # our helper listener functions
    def addHasBeenModifiedListener(self, callback: 'function'):
        """
        Register callback for `Has Been Modified` event

        :param callback: callback function
        """
        self._has_been_modified_listeners.append(callback)

    def addItemSelectedListener(self, callback: 'function'):
        """
        Register callback for `Item Selected` event

        :param callback: callback function
        """
        self._item_selected_listeners.append(callback)

    def addItemsDeselectedListener(self, callback: 'function'):
        """
        Register callback for `Items Deselected` event

        :param callback: callback function
        """
        self._items_deselected_listeners.append(callback)

    def addDragEnterListener(self, callback: 'function'):
        """
        Register callback for `Drag Enter` event

        :param callback: callback function
        """
        self.getView().addDragEnterListener(callback)

    def addDropListener(self, callback: 'function'):
        """
        Register callback for `Drop` event

        :param callback: callback function
        """
        self.getView().addDropListener(callback)

    # custom flag to detect node or edge has been selected....
    def resetLastSelectedStates(self):
        """Resets internal `selected flags` in all `Nodes` and `Edges` in the `Scene`"""
        for node in self.nodes:
            node.grNode._last_selected_state = False
        for edge in self.edges:
            edge.grEdge._last_selected_state = False

    def getView(self) -> 'QGraphicsView':
        """Shortcut for returning `Scene` ``QGraphicsView``

        :return: ``QGraphicsView`` attached to the `Scene`
        :rtype: ``QGraphicsView``
        """
        return self.grScene.views()[0]

    def getItemAt(self, pos: 'QPointF'):
        """Shortcut for retrieving item at provided `Scene` position

        :param pos: scene position
        :type pos: ``QPointF``
        :return: Qt Graphics Item at scene position
        :rtype: ``QGraphicsItem``
        """
        return self.getView().itemAt(pos)

    def addNode(self, node: Node):
        """Add :class:`~nodeeditor.node_node.Node` to this `Scene`

        :param node: :class:`~nodeeditor.node_node.Node` to be added to this `Scene`
        :type node: :class:`~nodeeditor.node_node.Node`
        """
        # Check if this is an output node and if one already exists (only during normal operation, not deserialization)
        if not self._deserializing and hasattr(node, 'op_code'):
            from src.conf import OP_NODE_OUTPUT
            if node.op_code == OP_NODE_OUTPUT:
                existing_output_nodes = [n for n in self.nodes if hasattr(n, 'op_code') and n.op_code == OP_NODE_OUTPUT]
                if existing_output_nodes:
                    show_message(self, "Warning", "An output node already exists in this graph. Please remove it before adding a new one.")
                    raise ValueError("Only one output node is allowed per graph")
        self.nodes.append(node)

    def addEdge(self, edge: Edge):
        """Add :class:`~nodeeditor.node_edge.Edge` to this `Scene`

        :param edge: :class:`~nodeeditor.node_edge.Edge` to be added to this `Scene`
        :return: :class:`~nodeeditor.node_edge.Edge`
        """
        self.edges.append(edge)

    def removeNode(self, node: Node):
        """Remove :class:`~nodeeditor.node_node.Node` from this `Scene`

        :param node: :class:`~nodeeditor.node_node.Node` to be removed from this `Scene`
        :type node: :class:`~nodeeditor.node_node.Node`
        """
        if node in self.nodes: self.nodes.remove(node)
        else:
            if DEBUG_REMOVE_WARNINGS: print("!W:", "Scene::removeNode", "wanna remove nodeeditor", node,
                                            "from self.nodes but it's not in the list!")

    def removeEdge(self, edge: Edge):
        """Remove :class:`~nodeeditor.node_edge.Edge` from this `Scene`

        :param edge: :class:`~nodeeditor.node_edge.Edge` to be remove from this `Scene`
        :return: :class:`~nodeeditor.node_edge.Edge`
        """
        if edge in self.edges: self.edges.remove(edge)
        else:
            if DEBUG_REMOVE_WARNINGS: print("!W:", "Scene::removeEdge", "wanna remove edge", edge,
                                            "from self.edges but it's not in the list!")


    def clear(self):
        """Remove all `Nodes` from this `Scene`. This causes also to remove all `Edges`"""
        while len(self.nodes) > 0:
            self.nodes[0].remove()

        self.has_been_modified = False


    def saveToFile(self, filename: str):
        """
        Save this `Scene` to the file on disk.

        :param filename: where to save this scene
        :type filename: ``str``
        """
        with open(filename, "w") as file:
            from src.common import __version__
            from src.conf import OP_NODE_OUTPUT
            
            # Find output nodes
            output_nodes = [node for node in self.nodes if hasattr(node, 'op_code') and node.op_code == OP_NODE_OUTPUT]
            
            # Ensure only one output node exists
            if len(output_nodes) > 1:
                raise ValueError(f"Multiple output nodes found ({len(output_nodes)}). Only one output node is allowed per file.")
            
            m_Nodes = self.serialize()
            m_Children = []
            m_Editor = {'name': 'VsmartEditor', 'version': __version__}
            m_Choices = []
            
            # If there's an output node, evaluate it and use its value for m_Children
            if output_nodes:
                output_node = output_nodes[0]
                try:
                    # Evaluate the output node to get the final result
                    output_value = output_node.eval()
                    if output_value is not None:
                        # If output_value is a list, use it directly as m_Children
                        if isinstance(output_value, list):
                            m_Children = output_value
                        else:
                            # If it's a single value, wrap it in a list
                            m_Children = [output_value]
                except Exception as e:
                    print(f"Error evaluating output node: {e}")
                    # Continue with empty m_Children if evaluation fails
            
            # Get variables from VariablesManager if available
            if self.variables_manager:
                try:
                    raw_variables = self.variables_manager.serialize_variables()
                    # Convert tuples to lists for KV3 compatibility
                    self.m_Variables = self._convert_tuples_to_lists(raw_variables)
                except Exception as e:
                    print(f"Error serializing variables: {e}")
                    self.m_Variables = []
            
            # Use self.m_Variables for variable serialization
            data = {'generic_data_type': "CSmartPropRoot", 'm_Editor': m_Editor, 'm_Variables': self.m_Variables, 'm_Choices':m_Choices, 'm_Children':m_Children, 'm_Nodes': m_Nodes}
            kv3.write(data,file)
            # print("saving to", filename, "was successfull.")

            self.has_been_modified = False
            self.filename = filename

    def loadFromFile(self, filename: str):
        """
        Load `Scene` from a file on disk

        :param filename: from what file to load the `Scene`
        :type filename: ``str``
        :raises: :class:`~nodeeditor.node_scene.InvalidFile` if there was an error decoding KV3 file
        """

        try:
            data = (kv3.read(filename)).value
            data = vsmart_remapping(data)
            self.filename = filename
            # Restore variables from m_Variables
            self.m_Variables = data.get('m_Variables', [])
            
            # Restore variables to VariablesManager if available
            if self.variables_manager and self.m_Variables:
                try:
                    self.variables_manager.deserialize_variables(self.m_Variables)
                except Exception as e:
                    print(f"Error deserializing variables: {e}")
            
            m_Nodes = data.get('m_Nodes', None)
            if m_Nodes is None:
                # No m_Nodes found - this might be a raw vsmart file, try to convert it
                try:
                    from src.vsmart.fileload import load_vsmart_from_data, show_load_results_dialog
                    created_nodes, errors, warnings = load_vsmart_from_data(data, self, self.element_id_generator)

                    # Show results dialog if there are any errors or warnings
                    if errors or warnings:
                        # Try to get the main window as parent for the dialog
                        parent = None
                        show_load_results_dialog(errors, warnings, parent)

                    # If conversion was successful, mark as not modified
                    if created_nodes and not errors:
                        self.has_been_modified = False

                except Exception as e:
                    # If vsmart conversion fails, raise an InvalidFile error
                    raise InvalidFile(f"File does not contain valid node data (m_Nodes) and vsmart conversion failed: {str(e)}")
            else:
                self.deserialize(m_Nodes)
            self.has_been_modified = False
        except Exception as e:
            dumpException(e)

    def getEdgeClass(self):
        """Return the class representing Edge. Override me if needed"""
        return Edge

    def setNodeClassSelector(self, class_selecting_function: 'functon') -> 'Node class type':
        """
        Set the function which decides what `Node` class to instantiate when deserializing `Scene`.
        If not set, we will always instantiate :class:`~nodeeditor.node_node.Node` for each `Node` in the `Scene`

        :param class_selecting_function: function which returns `Node` class type (not instance) from `Node` serialized ``dict`` data
        :type class_selecting_function: ``function``
        :return: Class Type of `Node` to be instantiated during deserialization
        :rtype: `Node` class type
        """
        self.node_class_selector = class_selecting_function

    def getNodeClassFromData(self, data: dict) -> 'Node class instance':
        """
        Takes `Node` serialized data and determines which `Node Class` to instantiate according the description
        in the serialized Node

        :param data: serialized `Node` object data
        :type data: ``dict``
        :return: Instance of `Node` class to be used in this Scene
        :rtype: `Node` class instance
        """
        return Node if self.node_class_selector is None else self.node_class_selector(data)


    def serialize(self) -> OrderedDict:
        nodes, edges = [], []
        for node in self.nodes:
            newnode = node.serialize()
            if not any (newnode['id'] == a['id'] for a in nodes):
                nodes.append(newnode)
        for edge in self.edges:
            newedge = edge.serialize()
            if not any (newedge['id'] == a['id'] for a in edges):
                edges.append(newedge)
        return OrderedDict([
            ('id', self.id),
            ('scene_width', self.scene_width),
            ('scene_height', self.scene_height),
            ('nodes', nodes),
            ('edges', edges),
        ])

    def deserialize(self, data: dict, hashmap: dict={}, restore_id: bool=True, *args, **kwargs) -> bool:
        hashmap = {}
        
        # Set deserializing flag to disable output node validation
        self._deserializing = True
        
        try:
            if restore_id: self.id = data['id']

            # -- deserialize NODES

            ## Instead of recreating all the nodes, reuse existing ones...
            # get list of all current nodes:
            all_nodes = self.nodes.copy()

            # go through deserialized nodes:
            for node_data in data['nodes']:
                # can we find this node in the scene?
                found = False
                for node in all_nodes:
                    if node.id == node_data['id']:
                        found = node
                        break

                if not found:
                    try:
                        new_node = self.getNodeClassFromData(node_data)(self, element_id_generator=self.element_id_generator)
                        new_node.deserialize(node_data, hashmap, restore_id, *args, **kwargs)
                        new_node.onDeserialized(node_data)
                        # print("New node for", node_data['title'])
                    except:
                        dumpException()
                else:
                    try:
                        found.deserialize(node_data, hashmap, restore_id, *args, **kwargs)
                        found.onDeserialized(node_data)
                        all_nodes.remove(found)
                        # print("Reused", node_data['title'])
                    except: dumpException()

            # remove nodes which are left in the scene and were NOT in the serialized data!
            # that means they were not in the graph before...
            while all_nodes != []:
                node = all_nodes.pop()
                node.remove()


            # -- deserialize EDGES


            ## Instead of recreating all the edges, reuse existing ones...
            # get list of all current edges:
            all_edges = self.edges.copy()

            # go through deserialized edges:
            for edge_data in data['edges']:
                # can we find this node in the scene?
                found = False
                for edge in all_edges:
                    if edge.id == edge_data['id']:
                        found = edge
                        break

                if not found:
                    new_edge = self.getEdgeClass()(self).deserialize(edge_data, hashmap, restore_id, *args, **kwargs)
                    # print("New edge for", edge_data)
                else:
                    found.deserialize(edge_data, hashmap, restore_id, *args, **kwargs)
                    all_edges.remove(found)

            # remove nodes which are left in the scene and were NOT in the serialized data!
            # that means they were not in the graph before...
            while all_edges != []:
                edge = all_edges.pop()
                edge.remove()

        finally:
            # Always reset the deserializing flag
            self._deserializing = False

        return True

    def _convert_tuples_to_lists(self, obj):
        """
        Recursively convert tuples to lists for KV3 compatibility
        """
        if isinstance(obj, tuple):
            return [self._convert_tuples_to_lists(item) for item in obj]
        elif isinstance(obj, list):
            return [self._convert_tuples_to_lists(item) for item in obj]
        elif isinstance(obj, dict):
            return {key: self._convert_tuples_to_lists(value) for key, value in obj.items()}
        else:
            return obj