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

import net.minecraft.world.Container
import net.minecraftforge.common.capabilities.Capability
import net.minecraftforge.common.capabilities.ICapabilityProvider
import ru.dbotthepony.mc.otm.container.iterator
import ru.dbotthepony.mc.otm.core.ifPresentK

open class CapabilityIterator<P : ICapabilityProvider, T>(protected open val parent: Iterator<P>, private val cap: Capability<T>) : Iterator<Pair<P, T>> {
	private var provider: P? = null
	private var capability: T? = null
	private var searched = false

	private fun search() {
		searched = true

		if (provider != null) {
			return
		}

		for (provider in parent) {
			provider.getCapability(cap).ifPresentK {
				this.provider = provider
				capability = it
				return
			}
		}
	}

	override fun hasNext(): Boolean {
		if (!searched) {
			search()
		}

		return provider != null
	}

	override fun next(): Pair<P, T> {
		if (!searched) {
			search()
		}

		val provider = provider ?: throw IllegalStateException("No next element")
		val capability = capability ?: throw IllegalStateException("No next element")
		this.provider = null
		this.capability = null
		this.searched = false
		return provider to capability
	}
}

class MutableCapabilityIterator<P : ICapabilityProvider, T>(override val parent: MutableIterator<P>, cap: Capability<T>) : CapabilityIterator<P, T>(parent, cap), MutableIterator<Pair<P, T>> {
	override fun remove() {
		parent.remove()
	}
}

fun <T> Container.iterator(cap: Capability<T>) = CapabilityIterator(iterator().nonEmpty(), cap)
fun <P : ICapabilityProvider, T> Iterator<P>.filtered(cap: Capability<T>) = CapabilityIterator(this, cap)
fun <P : ICapabilityProvider, T> MutableIterator<P>.filtered(cap: Capability<T>) = MutableCapabilityIterator(this, cap)
