namespace Nodebox.Extensions;

public static class GraphExtensions {
    extension(Graph graph) {
        public Graph SubGraphWith(Node node) {
            Assert.True(graph.Contains<Node>(node));

            var connectedNodes = new Queue<Node>([node]);
            var subgraph = new Graph();
            var closedNodes = new HashSet<Node>();
            while (connectedNodes.TryDequeue(out var connectedNode)) {
                if (closedNodes.Contains(connectedNode)) {
                    continue;
                }

                subgraph.Add(connectedNode);
                closedNodes.Add(connectedNode);
                foreach (var wireSet in connectedNode.InputWires) {
                    foreach (var wire in wireSet) {
                        connectedNodes.Enqueue(wire.Source);
                    }
                }

                foreach (var wireSet in connectedNode.OutputWires) {
                    foreach (var wire in wireSet) {
                        connectedNodes.Enqueue(wire.Destination);
                    }
                }
            }
            return subgraph;
        }

        public List<Graph> ConnectedSubGraphs() {
            var nodes = new Queue<Node>(graph.GetAll<Node>());
            var graphs = new List<Graph>();
            var closedNodes = new HashSet<Node>();

            while (nodes.TryDequeue(out var node)) {
                if (closedNodes.Contains(node)) {
                    continue;
                }

                var subgraph = graph.SubGraphWith(node);
                graphs.Add(subgraph);
                foreach (var subgraphNode in subgraph.GetAll<Node>()) {
                    closedNodes.Add(subgraphNode);
                }
            }

            return graphs;
        }
    }
}