package ru.dbotthepony.mc.otm.core.collect

import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenCustomHashSet
import it.unimi.dsi.fastutil.objects.ObjectOpenCustomHashSet
import it.unimi.dsi.fastutil.objects.ObjectSet
import ru.dbotthepony.mc.otm.core.util.HashedWeakReference
import java.lang.ref.ReferenceQueue

class WeakHashSet<E : Any>(initialCapacity: Int = 16, loadFactor: Float = 0.75f, linked: Boolean = false) : MutableSet<E> {
	private val queue = ReferenceQueue<E>()
	private val backing: ObjectSet<Any> = if (linked) ObjectLinkedOpenCustomHashSet(initialCapacity, loadFactor, ReferenceHashStrategy) else ObjectOpenCustomHashSet(initialCapacity, loadFactor, ReferenceHashStrategy)

	private fun purge() {
		var next = queue.poll()

		while (next != null) {
			backing.remove(next)
			next = queue.poll()
		}
	}

	override fun add(element: E): Boolean {
		purge()
		if (element in backing) return false
		return backing.add(HashedWeakReference(element, queue))
	}

	override fun addAll(elements: Collection<E>): Boolean {
		var any = false
		elements.forEach { any = add(it) || any }
		return any
	}

	override fun clear() {
		backing.clear()
		while (queue.poll() != null) {}
	}

	override fun iterator(): MutableIterator<E> {
		purge()
		return backing.iterator().map { (it as HashedWeakReference<E>).get() }.filterNotNull()
	}

	override fun remove(element: E): Boolean {
		purge()
		return backing.remove(element)
	}

	override fun removeAll(elements: Collection<E>): Boolean {
		purge()
		return backing.removeAll(elements)
	}

	override fun retainAll(elements: Collection<E>): Boolean {
		purge()
		return backing.retainAll(elements)
	}

	override val size: Int get() {
		purge()
		return backing.size
	}

	override fun contains(element: E): Boolean {
		purge()
		return backing.contains(element)
	}

	override fun containsAll(elements: Collection<E>): Boolean {
		purge()
		return backing.containsAll(elements)
	}

	override fun isEmpty(): Boolean {
		purge()
		return backing.isEmpty()
	}
}
