package org.warp.commonutils.type;

import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
import it.unimi.dsi.fastutil.objects.ObjectSet;
import java.util.Objects;
import java.util.Set;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public class HashMultiAssociation<T, U> implements MultiAssociation<T, U> {

	private final Object2ObjectOpenHashMap<T, ObjectOpenHashSet<U>> associations;
	private final Object2ObjectOpenHashMap<U, ObjectOpenHashSet<T>> inverseAssociations;

	public HashMultiAssociation() {
		this.associations = new Object2ObjectOpenHashMap<>();
		this.inverseAssociations = new Object2ObjectOpenHashMap<>();
	}

	private HashMultiAssociation(Object2ObjectOpenHashMap<T, ObjectOpenHashSet<U>> associations,
			Object2ObjectOpenHashMap<U, ObjectOpenHashSet<T>> inverseAssociations) {
		this.associations = associations;
		this.inverseAssociations = inverseAssociations;
	}

	/**
	 * Remove the source from this destination
	 */
	@Nullable
	private boolean removeSourceFromDest(@NotNull U dest, @NotNull T src) {
		var sources = inverseAssociations.get(dest);
		if (sources != null) {
			sources.remove(src);

			// Keep the map clean
			if (sources.isEmpty()) {
				inverseAssociations.remove(dest);
			}

			return true;
		}
		return false;
	}

	/**
	 * Remove the destination from this source
	 */
	@Nullable
	private boolean removeDestFromSource(@NotNull T src, @NotNull U dest) {
		var dests = associations.get(src);
		if (dests != null) {
			dests.remove(dest);

			// Keep the map clean
			if (dests.isEmpty()) {
				associations.remove(src);
			}

			return true;
		}
		return false;
	}

	@Override
	public boolean link(T src, U dest) {
		return associations.computeIfAbsent(src, s -> new ObjectOpenHashSet<>(1)).add(dest);
	}

	@Override
	public boolean unlink(T src, U dest) {
		var removed = removeDestFromSource(src, dest);

		if (removed) {
			removeSourceFromDest(dest, src);
			return true;
		}

		return false;
	}

	@Override
	public Set<U> unlink(T src) {
		var dests = associations.remove(src);

		if (dests == null) {
			return Set.of();
		}

		for (U dest : dests) {
			removeSourceFromDest(dest, src);
		}
		return dests;
	}

	@Override
	public Set<T> unlinkFromSource(U dest) {
		var sources = inverseAssociations.remove(dest);

		if (sources == null) {
			return Set.of();
		}

		for (T source : sources) {
			removeDestFromSource(source, dest);
		}
		return sources;
	}

	@Override
	public boolean hasAnyLink(T src) {
		var dests = associations.get(src);
		if (dests == null) {
			return false;
		}
		return !dests.isEmpty();
	}

	@Override
	public boolean hasAnyLinkSource(U dest) {
		var sources = inverseAssociations.get(dest);
		if (sources == null) {
			return false;
		}
		return !sources.isEmpty();
	}

	@Override
	public boolean hasLink(T src, U dest) {
		var dests = associations.get(src);
		if (dests == null) {
			return false;
		}
		return dests.contains(dest);
	}

	@Override
	public Set<U> getLinks(T src) {
		Objects.requireNonNull(src);
		var dests = associations.get(src);
		if (dests == null) return Set.of();
		return dests.clone();
	}

	@Override
	public Set<T> getLinkSources(U dest) {
		Objects.requireNonNull(dest);
		var sources = inverseAssociations.get(dest);
		if (sources == null) return Set.of();
		return sources.clone();
	}

	@Override
	public void clear() {
		associations.clear();
		inverseAssociations.clear();
	}

	@Override
	public int size() {
		return Math.max(associations.size(), inverseAssociations.size());
	}

	@Override
	public ObjectSet<T> getSources() {
		return associations.clone().keySet();
	}

	@Override
	public ObjectSet<U> getDestinations() {
		return inverseAssociations.clone().keySet();
	}

	@SuppressWarnings("MethodDoesntCallSuperMethod")
	@Override
	public HashMultiAssociation<T, U> clone() {
		var associationsClone = associations.clone();
		associationsClone.replaceAll((item, destinations) -> destinations.clone());

		var inverseAssociationsClone = inverseAssociations.clone();
		inverseAssociationsClone.replaceAll((item, sources) -> sources.clone());

		return new HashMultiAssociation<>(associationsClone, inverseAssociationsClone);
	}
}