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.Optional;
import java.util.Set;

public class HashAssociation<T, U> implements Association<T, U> {

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

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

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

	@Override
	public boolean link(T src, U dest) {
		Objects.requireNonNull(src);
		Objects.requireNonNull(dest);
		// If the destination is associated to null
		var previousSrc = inverseAssociations.get(dest);

		// Return if the association already exists
		if (src.equals(previousSrc)) {
			return false;
		}

		// Remove the previous association if present
		if (previousSrc != null) {
			inverseAssociations.remove(dest);
			var destinations = associations.get(previousSrc);
			destinations.remove(dest);
			if (destinations.isEmpty()) {
				associations.remove(previousSrc);
			}
		}

		if (!associations.computeIfAbsent(src, s -> new ObjectOpenHashSet<>(0)).add(dest)) {
			throw new IllegalStateException("Association was partially linked");
		}
		return true;
	}

	@Override
	public boolean unlink(T src, U dest) {
		Objects.requireNonNull(src);
		Objects.requireNonNull(dest);

		var linkedSrc = inverseAssociations.get(dest);

		// Check if the link is different
		if (!Objects.equals(src, linkedSrc)) {
			return false;
		}

		inverseAssociations.remove(dest);
		var destinations = associations.get(src);
		if (!destinations.remove(dest)) {
			throw new IllegalStateException("Association was partially linked");
		}
		if (destinations.isEmpty()) {
			associations.remove(src);
		}

		return true;
	}

	@Override
	public Set<U> unlink(T src) {
		Objects.requireNonNull(src);

		var destinations = associations.remove(src);

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

		for (U destination : destinations) {
			if (!Objects.equals(src, inverseAssociations.remove(destination))) {
				throw new IllegalStateException("Association was partially linked");
			}
		}

		return destinations;
	}

	@Override
	public Optional<T> unlinkFromSource(U dest) {
		Objects.requireNonNull(dest);

		var previousSrc = inverseAssociations.remove(dest);

		if (previousSrc == null) {
			return Optional.empty();
		}

		var destinations = associations.get(previousSrc);
		destinations.remove(dest);
		if (destinations.isEmpty()) {
			associations.remove(previousSrc);
		}

		return Optional.of(previousSrc);
	}

	@Override
	public boolean hasAnyLink(T src) {
		Objects.requireNonNull(src);
		var destinations = associations.get(src);
		return destinations != null && !destinations.isEmpty();
	}

	@Override
	public boolean hasLinkSource(U dest) {
		Objects.requireNonNull(dest);
		return inverseAssociations.containsKey(dest);
	}

	@Override
	public boolean hasLink(T src, U dest) {
		Objects.requireNonNull(src);
		Objects.requireNonNull(dest);
		return Objects.equals(src, inverseAssociations.get(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 Optional<T> getLinkSource(U dest) {
		Objects.requireNonNull(dest);

		return Optional.ofNullable(inverseAssociations.get(dest));
	}

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

	@Override
	public int size() {
		return inverseAssociations.size();
	}

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

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

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