package p.studio.utilities.structures;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;

public final class DependencyGraphAnaliser<N> {
    private final Comparator<? super N> nodeOrder;

    public DependencyGraphAnaliser(final Comparator<? super N> nodeOrder) {
        this.nodeOrder = nodeOrder;
    }

    public Analysis<N> analyse(final Map<N, ? extends Collection<N>> dependenciesByNode) {
        final var canonicalGraph = canonicalize(dependenciesByNode);
        final var indegree = new HashMap<N, Integer>();
        final var reverse = new HashMap<N, LinkedHashSet<N>>();
        for (final var node : canonicalGraph.keySet()) {
            indegree.put(node, 0);
            reverse.put(node, new LinkedHashSet<>());
        }

        for (final var entry : canonicalGraph.entrySet()) {
            for (final var dependency : entry.getValue()) {
                indegree.put(entry.getKey(), indegree.get(entry.getKey()) + 1);
                reverse.get(dependency).add(entry.getKey());
            }
        }

        final var ready = new PriorityQueue<N>(nodeOrder);
        for (final var entry : indegree.entrySet()) {
            if (entry.getValue() == 0) {
                ready.add(entry.getKey());
            }
        }

        final var traversalOrder = new ArrayList<N>(canonicalGraph.size());
        while (!ready.isEmpty()) {
            final var current = ready.remove();
            traversalOrder.add(current);
            for (final var dependent : reverse.getOrDefault(current, new LinkedHashSet<>())) {
                final var next = indegree.get(dependent) - 1;
                indegree.put(dependent, next);
                if (next == 0) {
                    ready.add(dependent);
                }
            }
        }

        final var stronglyConnectedComponents = tarjan(canonicalGraph);
        final var cycleComponents = new ArrayList<ReadOnlyList<N>>();
        for (final var scc : stronglyConnectedComponents) {
            if (scc.size() > 1) {
                cycleComponents.add(scc);
                continue;
            }
            final var node = scc.getFirst();
            if (canonicalGraph.getOrDefault(node, new LinkedHashSet<>()).contains(node)) {
                cycleComponents.add(scc);
            }
        }

        return new Analysis<>(
                ReadOnlyList.wrap(traversalOrder),
                ReadOnlyList.wrap(stronglyConnectedComponents),
                ReadOnlyList.wrap(cycleComponents),
                traversalOrder.size() == canonicalGraph.size());
    }

    private LinkedHashMap<N, LinkedHashSet<N>> canonicalize(final Map<N, ? extends Collection<N>> dependenciesByNode) {
        final var graph = new LinkedHashMap<N, LinkedHashSet<N>>();
        for (final var entry : dependenciesByNode.entrySet()) {
            graph.computeIfAbsent(entry.getKey(), ignored -> new LinkedHashSet<>());
            for (final var dependency : entry.getValue()) {
                graph.computeIfAbsent(dependency, ignored -> new LinkedHashSet<>());
                graph.get(entry.getKey()).add(dependency);
            }
        }
        return graph;
    }

    private List<ReadOnlyList<N>> tarjan(final LinkedHashMap<N, LinkedHashSet<N>> graph) {
        final var indexByNode = new HashMap<N, Integer>();
        final var lowLinkByNode = new HashMap<N, Integer>();
        final var inStack = new HashSet<N>();
        final var stack = new ArrayDeque<N>();
        final var components = new ArrayList<ReadOnlyList<N>>();
        final var indexRef = new int[]{0};

        for (final var node : graph.keySet()) {
            if (!indexByNode.containsKey(node)) {
                strongConnect(node, graph, indexByNode, lowLinkByNode, inStack, stack, components, indexRef);
            }
        }

        return components;
    }

    private void strongConnect(
            final N node,
            final Map<N, LinkedHashSet<N>> graph,
            final Map<N, Integer> indexByNode,
            final Map<N, Integer> lowLinkByNode,
            final Set<N> inStack,
            final Deque<N> stack,
            final List<ReadOnlyList<N>> components,
            final int[] indexRef) {
        indexByNode.put(node, indexRef[0]);
        lowLinkByNode.put(node, indexRef[0]);
        indexRef[0]++;
        stack.push(node);
        inStack.add(node);

        for (final var dependency : graph.getOrDefault(node, new LinkedHashSet<>())) {
            if (!indexByNode.containsKey(dependency)) {
                strongConnect(dependency, graph, indexByNode, lowLinkByNode, inStack, stack, components, indexRef);
                lowLinkByNode.put(node, Math.min(lowLinkByNode.get(node), lowLinkByNode.get(dependency)));
            } else if (inStack.contains(dependency)) {
                lowLinkByNode.put(node, Math.min(lowLinkByNode.get(node), indexByNode.get(dependency)));
            }
        }

        if (lowLinkByNode.get(node).equals(indexByNode.get(node))) {
            final var component = new ArrayList<N>();
            while (true) {
                final var current = stack.pop();
                inStack.remove(current);
                component.add(current);
                if (current.equals(node)) {
                    break;
                }
            }
            component.sort(nodeOrder);
            components.add(ReadOnlyList.wrap(component));
        }
    }

    public record Analysis<N>(
            ReadOnlyList<N> traversalOrder,
            ReadOnlyList<ReadOnlyList<N>> stronglyConnectedComponents,
            ReadOnlyList<ReadOnlyList<N>> cycleComponents,
            boolean acyclic) {
    }
}
