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

import java.util.Optional
import java.util.function.BinaryOperator
import java.util.function.Predicate
import java.util.stream.Collector

// Purpose of Stream API over Iterators is that it is simple enough for JIT to inline most of it,
// unlike actual Streams.

// We lose only one (actual) element of Stream API here: ability to parallelize work,
// except it doesn't (properly) work in Minecraft (Forge) either, due to classloader bug:
// https://github.com/MinecraftForge/EventBus/issues/44

// Aside parallel work, unimplemented Stream API elements can be easily implemented when required

class FilteredIterator<T>(private val parent: Iterator<T>, private val predicate: Predicate<in T>) : Iterator<T> {
	private var foundValue: Any? = Companion

	override fun hasNext(): Boolean {
		if (foundValue === Companion) {
			while (parent.hasNext()) {
				val next = parent.next()

				if (predicate.test(next)) {
					foundValue = next
					return true
				}
			}

			return false
		}

		return true
	}

	@Suppress("unchecked_cast")
	override fun next(): T {
		var foundValue = foundValue

		if (foundValue === Companion) {
			if (!hasNext()) {
				throw NoSuchElementException()
			}

			foundValue = this.foundValue

			if (foundValue === Companion) {
				throw ConcurrentModificationException()
			}
		}

		this.foundValue = Companion
		return foundValue as T
	}

	private companion object
}

class MappingIterator<T, R>(private val parent: Iterator<T>, private val transform: (T) -> R) : Iterator<R> {
	override fun hasNext(): Boolean {
		return parent.hasNext()
	}

	override fun next(): R {
		return transform.invoke(parent.next())
	}
}

class FlatMappingIterator<T, R>(private val parent: Iterator<T>, private val transform: (T) -> Iterator<R>) : Iterator<R> {
	private var current: Iterator<R>? = null

	override fun hasNext(): Boolean {
		while (current?.hasNext() != true && parent.hasNext()) {
			current = transform.invoke(parent.next())
		}

		return current?.hasNext() == true
	}

	override fun next(): R {
		if (!hasNext())
			throw NoSuchElementException()

		return current!!.next()
	}
}

class LimitingIterator<T>(val parent: Iterator<T>, val limit: Long) : Iterator<T> {
	init {
		require(limit > 0) { "Invalid limit $limit" }
	}

	private var found = 0L

	override fun hasNext(): Boolean {
		return found < limit && parent.hasNext()
	}

	override fun next(): T {
		if (found >= limit)
			throw NoSuchElementException()

		return parent.next()
	}
}

class SkippingIterator<T>(val parent: Iterator<T>, skip: Long) : Iterator<T> {
	init {
		require(skip >= 0) { "Invalid skip amount $skip" }
	}

	private var found = skip

	override fun hasNext(): Boolean {
		while (parent.hasNext() && found > 0L) {
			found--
			parent.next()
		}

		return parent.hasNext()
	}

	override fun next(): T {
		if (!hasNext())
			throw NoSuchElementException()

		return parent.next()
	}
}

fun <T> Iterator<T>.filter(condition: Predicate<in T>): Iterator<T> = FilteredIterator(this, condition)
fun <T, R> Iterator<T>.map(mapper: (T) -> R): Iterator<R> = MappingIterator(this, mapper)
fun <T, R> Iterator<T>.map(mapper: java.util.function.Function<T, R>): Iterator<R> = MappingIterator(this, mapper::apply)
fun <T, R> Iterator<T>.flatMap(mapper: (T) -> Iterator<R>): Iterator<R> = FlatMappingIterator(this, mapper)
fun <T, R> Iterator<T>.flatMap(mapper: java.util.function.Function<T, Iterator<R>>): Iterator<R> = FlatMappingIterator(this, mapper::apply)

fun interface O2DFunction<T> {
	fun apply(value: T): Double
}

fun interface O2IFunction<T> {
	fun apply(value: T): Int
}

fun <T> Iterator<T>.mapToDouble(mapper: O2DFunction<T>): it.unimi.dsi.fastutil.doubles.DoubleIterator {
	return object : it.unimi.dsi.fastutil.doubles.DoubleIterator {
		override fun hasNext(): Boolean {
			return this@mapToDouble.hasNext()
		}

		override fun remove() {
			throw UnsupportedOperationException()
		}

		override fun nextDouble(): Double {
			return mapper.apply(this@mapToDouble.next())
		}
	}
}

fun <T> Iterator<T>.mapToInt(mapper: O2IFunction<T>): it.unimi.dsi.fastutil.ints.IntIterator {
	return object : it.unimi.dsi.fastutil.ints.IntIterator {
		override fun hasNext(): Boolean {
			return this@mapToInt.hasNext()
		}

		override fun remove() {
			throw UnsupportedOperationException()
		}

		override fun nextInt(): Int {
			return mapper.apply(this@mapToInt.next())
		}
	}
}

fun it.unimi.dsi.fastutil.doubles.DoubleIterator.sum(): Double {
	var value = 0.0
	while (hasNext()) value += nextDouble()
	return value
}

fun it.unimi.dsi.fastutil.ints.IntIterator.sum(): Int {
	var value = 0
	while (hasNext()) value += nextInt()
	return value
}

fun interface DD2DFunction {
	fun apply(a: Double, b: Double): Double
}

fun interface II2IFunction {
	fun apply(a: Int, b: Int): Int
}

fun it.unimi.dsi.fastutil.doubles.DoubleIterator.reduce(identity: Double, reducer: DD2DFunction): Double {
	var result = identity
	while (hasNext()) result = reducer.apply(result, nextDouble())
	return result
}

fun it.unimi.dsi.fastutil.ints.IntIterator.reduce(identity: Int, reducer: II2IFunction): Int {
	var result = identity
	while (hasNext()) result = reducer.apply(result, nextInt())
	return result
}

fun <T> Iterator<T>.reduce(identity: T, reducer: (T, T) -> T): T {
	var result = identity

	for (value in this) {
		result = reducer.invoke(result, value)
	}

	return result
}

fun <T> Iterator<T>.reduce(identity: T, reducer: BinaryOperator<T>): T = reduce(identity, reducer::apply)
fun <T> Iterator<T?>.filterNotNull(): Iterator<T> = filter { it != null } as Iterator<T>

fun <T> Iterator<T>.anyMatch(predicate: Predicate<in T>): Boolean {
	for (value in this) {
		if (predicate.test(value)) {
			return true
		}
	}

	return false
}

fun <T> Iterator<T>.allMatch(predicate: Predicate<in T>): Boolean {
	for (value in this) {
		if (!predicate.test(value)) {
			return false
		}
	}

	return true
}

fun <T> Iterator<T>.noneMatch(predicate: Predicate<in T>): Boolean {
	for (value in this) {
		if (predicate.test(value)) {
			return false
		}
	}

	return true
}

fun <T, A, R> Iterator<T>.collect(collector: Collector<T, A, R>): R {
	val accumulator = collector.accumulator()
	val instance = collector.supplier().get()

	for (value in this) {
		accumulator.accept(instance, value)
	}

	return collector.finisher().apply(instance)
}

fun <T : Any> Iterator<T>.findFirst(): Optional<T> {
	if (hasNext()) {
		return Optional.of(next())
	}

	return Optional.empty<T>()
}

fun <T : Any> Iterator<T>.findAny() = findFirst()
fun <T> Iterator<T>.limit(limit: Long) = LimitingIterator(this, limit)
fun <T> Iterator<T>.skip(skip: Long) = if (skip == 0L) this else SkippingIterator(this, skip)

inline fun <T> Iterator<T>.forEach(action: (T) -> Unit) {
	for (value in this) {
		action.invoke(value)
	}
}

fun <T : Any> Iterator<T>.min(comparator: Comparator<T>): Optional<T> {
	if (!hasNext()) {
		return Optional.empty<T>()
	}

	var min = next()

	for (value in this) {
		if (comparator.compare(min, value) > 0) {
			min = value
		}
	}

	return Optional.of(min)
}

fun <T : Any> Iterator<T>.max(comparator: Comparator<T>): Optional<T> {
	if (!hasNext()) {
		return Optional.empty<T>()
	}

	var max = next()

	for (value in this) {
		if (comparator.compare(max, value) < 0) {
			max = value
		}
	}

	return Optional.of(max)
}

fun <T> Iterator<T>.peek(peeker: (T) -> Unit): Iterator<T> {
	return object : Iterator<T> {
		override fun hasNext(): Boolean {
			return this@peek.hasNext()
		}

		override fun next(): T {
			return this@peek.next().also(peeker)
		}
	}
}
