UnityImageProcessing_Package/Scripts/Editor/Windows/ImageProcessingGraphViewWindow.cs

using System.Collections.Generic;
using System.Linq;
using ImageProcessingGraph.Editor.Unity_Image_Processing.Scripts.Editor.Windows;
using ImageProcessingGraph.Editor.Windows;
using UnityEditor;
using UnityEditor.Experimental.GraphView;
using UnityEngine;
using UnityEngine.UIElements;

namespace ImageProcessingGraph.Editor
{
    public class ImageProcessingGraphViewWindow : GraphView
    {
        internal ImageProcessingGraphAsset asset;
        private SerializedObject serializedObject;
        private ImageProcessingGraphEditorWindow window;
        
        public ImageProcessingGraphEditorWindow Window => window;

        public List<ImageProcessingGraphNodeVisual> graphNodes;
        public Dictionary<string, ImageProcessingGraphNodeVisual> nodeDictionary;
        public Dictionary<Edge, GraphConnection> connectionDictionary;
        
        internal ImageProcessingGraphSearchProvider searchProvider;
        internal ImageProcessingGraphEdgeConnectorListener edgeConnectorListener;
        
        public ImageProcessingGraphViewWindow(SerializedObject obeject, ImageProcessingGraphEditorWindow window)
        {
            this.serializedObject = obeject;
            this.asset = obeject.targetObject as ImageProcessingGraphAsset;
            
            this.graphNodes = new List<ImageProcessingGraphNodeVisual>();
            nodeDictionary = new Dictionary<string, ImageProcessingGraphNodeVisual>();
            connectionDictionary = new Dictionary<Edge, GraphConnection>();
            
            searchProvider = ScriptableObject.CreateInstance<ImageProcessingGraphSearchProvider>();
            searchProvider.graph = this;

            edgeConnectorListener = new ImageProcessingGraphEdgeConnectorListener(this);
            
            this.nodeCreationRequest = ShowSearchWindow;

            this.window = window;
            
            StyleSheet styleSheet = AssetDatabase.LoadAssetAtPath<StyleSheet>("Assets/Unity Image Processing/GraphView.uss");
            styleSheets.Add(styleSheet);
            
            GridBackground background = new GridBackground();

            background.name = "Grid";
            
            
            Add(background);
            
            background.SendToBack();
            
            this.AddManipulator(new ContentDragger());
            this.AddManipulator(new SelectionDragger());
            this.AddManipulator(new RectangleSelector());
            this.AddManipulator(new ClickSelector());
            this.AddManipulator(new ContentZoomer() );
            
            DrawNodes();
            DrawConnections();

            graphViewChanged += OnGraphViewChanged;
            Undo.undoRedoEvent += UndoEvent;
        }
        
        private ImageProcessingGraphNodeVisual GetNode(string NodeID)
        {
            ImageProcessingGraphNodeVisual node = null;
            nodeDictionary.TryGetValue(NodeID, out node);
            return node;
        }

        public override List<Port> GetCompatiblePorts(Port startPort, NodeAdapter nodeAdapter)
        {
            List<Port> compatiblePorts = new List<Port>();


            foreach (var node in graphNodes)
            {
                // Prevent connections to self
                if (node == startPort.node)
                    continue;
                
                foreach (var port in node.inputContainer.Children().Concat(node.outputContainer.Children()).OfType<Port>())
                {
                    // Prevent connecting input to input or output to output
                    if (port.direction == startPort.direction)
                        continue;

                    if (port.portType != startPort.portType)
                    {
                        if (DoesConversionNodeExist())
                        {

                        }
                        else
                            continue;
                    }
                    
                    // Prevent connection if it creates a cycle
                    if (startPort.direction == Direction.Output && CreatesCycle(startPort, port))
                        continue;

                    if (startPort.direction == Direction.Input && CreatesCycle(port, startPort))
                        continue;

                    compatiblePorts.Add(port);
                }   
            }

            return compatiblePorts;
        }

        private void UndoEvent(in UndoRedoInfo undo)
        {
            DrawNodes();
            
        }

        private GraphViewChange OnGraphViewChanged(GraphViewChange graphviewchange)
        {
            
            if (graphviewchange.movedElements != null)
            {
                Undo.RecordObject(serializedObject.targetObject, "Moved Graph Elements");
                foreach (var VARIABLE in graphviewchange.movedElements.OfType<ImageProcessingGraphNodeVisual>())
                {
                    VARIABLE.SavePosition();
                }
            }
            
            if (graphviewchange.elementsToRemove != null)
            {
                List<ImageProcessingGraphNodeVisual> nodesToRemove = graphviewchange.elementsToRemove.OfType<ImageProcessingGraphNodeVisual>().ToList();

                if (nodesToRemove.Count > 0)
                {
                    Undo.RecordObject(serializedObject.targetObject, "Remove Node");
                    
                    foreach (var VARIABLE in nodesToRemove)
                    {
                        RemoveNode(VARIABLE);
                    }   
                }

                foreach (var VARIABLE in graphviewchange.elementsToRemove.OfType<Edge>())
                {
                    RemoveEdge(VARIABLE);
                }
                
            }

            if (graphviewchange.edgesToCreate != null)
            {
                foreach (Edge edge in graphviewchange.edgesToCreate)
                {
                    Undo.RecordObject(serializedObject.targetObject, "Created Connection");
                    CreateEdge(edge);
                }
            }
            
            return graphviewchange;
        }

        #region Edges
    
        void CreateEdge(Edge edge)
        {
            ImageProcessingGraphNodeVisual outputNode = (ImageProcessingGraphNodeVisual)edge.output.node;
            ImageProcessingGraphNodeVisual inputNode = (ImageProcessingGraphNodeVisual)edge.input.node;
            
            
            int outputIndex = outputNode.OutputPorts.IndexOf(edge.output);
            int inputIndex = inputNode.InputPorts.IndexOf(edge.input);

            string inputType = inputNode.GraphNode.GetType().Name;
            string outputType = outputNode.GraphNode.GetType().Name;

            GraphConnection connection = new GraphConnection(inputNode.GraphNode.ID,inputIndex, inputType,outputNode.GraphNode.ID, outputIndex, outputType);
            
            asset.Connections.Add(connection);
        }
        
        private void RemoveEdge(Edge variable)
        {
            if (connectionDictionary.TryGetValue(variable, out GraphConnection connection))
            {
                asset.Connections.Remove(connection);
                connectionDictionary.Remove(variable);
            }
        }
        
        private void DrawConnections()
        {
            if (asset.Connections != null)
            {
                foreach (GraphConnection conn in asset.Connections)
                {
                    ImageProcessingGraphNodeVisual inputNode = GetNode(conn.inputPort.nodeID);
                    ImageProcessingGraphNodeVisual outputNode = GetNode(conn.outputPort.nodeID);

                    if (inputNode != null && outputNode != null)
                    {
                        Port inPort = inputNode.InputPorts[conn.inputPort.portID];
                        Port outPort = outputNode.OutputPorts[conn.outputPort.portID];

                        Edge edge = inPort.ConnectTo(outPort);
                        AddElement(edge);
                        connectionDictionary.Add(edge, conn);
                    }
                    
                }
            }
        }

        
        private bool DoesConversionNodeExist()
        {
            return false;
        }

        #region Cycle Logic

        private bool CreatesCycle(Port fromPort, Port toPort)
        {
            var visited = new HashSet<Node>();
            return HasPathTo(toPort.node, fromPort.node, visited);
        }

        private bool HasPathTo(Node current, Node target, HashSet<Node> visited)
        {
            if (current == target)
                return true;

            if (visited.Contains(current))
                return false;

            visited.Add(current);

            foreach (var outputPort in current.outputContainer.Children().OfType<Port>())
            {
                foreach (var edge in outputPort.connections)
                {
                    var nextNode = edge.input.node;
                    if (HasPathTo(nextNode, target, visited))
                        return true;
                }
            }

            return false;
        }

        #endregion
        
        #endregion
        
        #region Nodes
        private void DrawNodes()
        {
            foreach (ImageProcessingGraphNodeVisual node in graphNodes)
            {
                RemoveElement(node);
            }
            graphNodes.Clear();
            
            foreach (KeyValuePair<string, ImageProcessingGraphNodeVisual> node in nodeDictionary)
            {
                RemoveElement(node.Value);
            }
            nodeDictionary.Clear();
            
            foreach (var variable in asset.Nodes)
            {
                AddNodeToGraph(variable);
            }
        }
        
        private void AddNodeToGraph(BaseImageNode node)
        {
            node.typeName = node.GetType().AssemblyQualifiedName;

            ImageProcessingGraphNodeVisual editorNode = new ImageProcessingGraphNodeVisual(node, this);
            editorNode.SetPosition(node.Position);

            graphNodes.Add(editorNode);
            nodeDictionary.Add(node.ID, editorNode);
            
            
            AddElement(editorNode);
        }
        
        public void Add(BaseImageNode node)
        {
            Undo.RecordObject(serializedObject.targetObject, "Added Node");
            asset.Nodes.Add(node);
            
            serializedObject.Update();

            AddNodeToGraph(node);

        }
        
        private void RemoveNode(ImageProcessingGraphNodeVisual variable)
        {
            List<GraphConnection> connectionsToRemove = new List<GraphConnection>();
            
            foreach (var connection in asset.Connections)
            {
                if (connection.inputPort.nodeID == variable.GraphNode.ID || connection.outputPort.nodeID == variable.GraphNode.ID)
                    connectionsToRemove.Add(connection);
            }

            foreach (var connection in connectionsToRemove)
            {
                asset.Connections.Remove(connection);
            }
            
            asset.Nodes.Remove(variable.GraphNode);
            nodeDictionary.Remove(variable.GraphNode.ID);
            graphNodes.Remove(variable);
            serializedObject.Update(); 
        }
        
        #endregion

        #region Search Window

        private void ShowSearchWindow(NodeCreationContext obj)
        {
            searchProvider.target = (VisualElement)focusController.focusedElement;
            SearchWindow.Open(new SearchWindowContext(obj.screenMousePosition), searchProvider);
        }

        #endregion
        
        
    }
}