package net.minecraft.util;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/** AFTER — fixed: isCyclic passes visited set — O(E) per call, O(E²) total */
public class DependencySorter<K, V extends DependencySorter.Entry<K>> {
    private final Map<K, V> contents = new HashMap<>();

    public DependencySorter<K, V> addEntry(K id, V value) {
        this.contents.put(id, value);
        return this;
    }

    private void visitDependenciesAndElement(Multimap<K, K> dependencies, Set<K> alreadyVisited, K id, BiConsumer<K, V> output) {
        if (!alreadyVisited.add(id)) return;
        dependencies.get(id).forEach(dep -> visitDependenciesAndElement(dependencies, alreadyVisited, dep, output));
        V current = this.contents.get(id);
        if (current != null) output.accept(id, current);
    }

    // FIX: visited set prevents exponential revisiting of diamond nodes
    private static <K> boolean isCyclic(Multimap<K, K> directDependencies, K from, K to, Set<K> visited) {
        if (!visited.add(to)) return false; // already explored — no cycle via here
        Collection<K> dependencies = directDependencies.get(to);
        if (dependencies.contains(from)) return true;
        return dependencies.stream().anyMatch(dep -> isCyclic(directDependencies, from, dep, visited));
    }

    private static <K> void addDependencyIfNotCyclic(Multimap<K, K> directDependencies, K from, K to) {
        if (!isCyclic(directDependencies, from, to, new HashSet<>())) directDependencies.put(from, to);
    }

    public void orderByDependencies(BiConsumer<K, V> output) {
        HashMultimap<K, K> directDependencies = HashMultimap.create();
        this.contents.forEach((id, value) -> value.visitRequiredDependencies(dep -> addDependencyIfNotCyclic(directDependencies, id, dep)));
        this.contents.forEach((id, value) -> value.visitOptionalDependencies(dep -> addDependencyIfNotCyclic(directDependencies, id, dep)));
        Set<K> alreadyVisited = new HashSet<>();
        this.contents.keySet().forEach(id -> visitDependenciesAndElement(directDependencies, alreadyVisited, id, output));
    }

    public interface Entry<K> {
        void visitRequiredDependencies(Consumer<K> consumer);
        void visitOptionalDependencies(Consumer<K> consumer);
    }
}