package ru.dbotthepony.mc.otm.container

import it.unimi.dsi.fastutil.ints.IntAVLTreeSet
import it.unimi.dsi.fastutil.ints.IntArrayList
import it.unimi.dsi.fastutil.ints.IntArraySet
import it.unimi.dsi.fastutil.ints.IntSet
import it.unimi.dsi.fastutil.objects.Object2ObjectFunction
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenCustomHashMap
import net.minecraft.world.Container
import net.minecraft.world.item.ItemStack
import net.minecraft.world.item.enchantment.EnchantmentHelper.hasVanishingCurse
import net.minecraftforge.common.capabilities.Capability
import net.minecraftforge.fluids.capability.IFluidHandler
import ru.dbotthepony.mc.otm.core.addAll
import ru.dbotthepony.mc.otm.core.collect.iterator
import ru.dbotthepony.mc.otm.core.collect.nonEmpty
import ru.dbotthepony.mc.otm.core.isNotEmpty
import kotlin.math.roundToInt

operator fun Container.set(index: Int, value: ItemStack) = setItem(index, value)
operator fun Container.get(index: Int): ItemStack = getItem(index)

operator fun IFluidHandler.get(index: Int) = getFluidInTank(index)

fun Container.addItem(stack: ItemStack, range: IntRange, simulate: Boolean = false): ItemStack {
	if (this is MatteryContainer) {
		return this.addItem(stack, range, simulate)
	}

	if (range.last >= containerSize || range.first < 0)
		throw IllegalArgumentException("Invalid range: $range")

	if (stack.isEmpty)
		return stack

	val copy = stack.copy()

	// двигаем в одинаковые слоты
	for (slot in range) {
		if (ItemStack.isSameItemSameTags(this[slot], copy)) {
			val slotStack = this[slot]
			val slotLimit = maxStackSize.coerceAtMost(slotStack.maxStackSize)

			if (slotStack.count < slotLimit) {
				val newCount = (slotStack.count + copy.count).coerceAtMost(slotLimit)
				val diff = newCount - slotStack.count

				if (!simulate) {
					slotStack.count = newCount
					setChanged()
				}

				copy.shrink(diff)

				if (copy.isEmpty) {
					return copy
				}
			}
		}
	}

	// двигаем в пустые слоты
	for (slot in range) {
		if (this[slot].isEmpty) {
			val diff = copy.count.coerceAtMost(maxStackSize.coerceAtMost(copy.maxStackSize))

			if (!simulate) {
				val copyToPut = copy.copy()
				copyToPut.count = diff
				this[slot] = copyToPut
				setChanged()
			}

			copy.shrink(diff)

			if (copy.isEmpty) {
				return copy
			}
		}
	}

	return copy
}

fun Container.addItem(stack: ItemStack, simulate: Boolean): ItemStack = addItem(stack, 0 until containerSize, simulate)

fun Container.vanishCursedItems() {
	val iterator = iterator()

	for (item in iterator) {
		if (item.isNotEmpty && hasVanishingCurse(item)) {
			iterator.remove()
		}
	}
}

inline fun <T> Container.forEach(cap: Capability<T>, consumer: (Pair<ItemStack, T>) -> Unit) {
	for (value in iterator(cap)) {
		consumer(value)
	}
}

inline fun <T> Container.forEach(cap: Capability<T>, consumer: (ItemStack, T) -> Unit) {
	for ((a, b) in iterator(cap)) {
		consumer(a, b)
	}
}

inline fun <T> Container.forEachItem(cap: Capability<T>, consumer: (ItemStack) -> Unit) {
	for (pair in iterator(cap)) {
		consumer(pair.first)
	}
}

inline fun <T> Container.forEachCapability(cap: Capability<T>, consumer: (T) -> Unit) {
	for (pair in iterator(cap)) {
		consumer(pair.second)
	}
}

inline fun Container.forEach(lambda: (ItemStack) -> Unit) {
	for (value in iterator()) {
		lambda(value)
	}
}

inline fun Container.forEachNonEmpty(lambda: (ItemStack) -> Unit) {
	for (value in iterator().nonEmpty()) {
		lambda(value)
	}
}

fun Container.balance(slots: IntSet, checkForEmpty: Boolean = true) {
	if (slots.isEmpty() || checkForEmpty && !slots.any { getItem(it).isNotEmpty }) return

	val empty = IntArrayList()
	val itemTypes = Object2ObjectOpenCustomHashMap<ItemStack, IntAVLTreeSet>(ItemStackHashStrategy)

	for (i in slots.intIterator()) {
		val item = getItem(i)

		if (item.isEmpty) {
			empty.add(i)
		} else {
			itemTypes.computeIfAbsent(item, Object2ObjectFunction { IntAVLTreeSet() }).add(i)
		}
	}

	if (itemTypes.isEmpty())
		return

	// только один вид предмета, просто балансируем его во все слоты
	if (itemTypes.size == 1) {
		val (item, list) = itemTypes.entries.first()
		var count = 0
		for (i in list.intIterator()) count += getItem(i).count
		var ghostChanges = false

		// всего предметов меньше, чем слотов
		if (count < slots.size) {
			for (slot in list.intIterator()) {
				getItem(slot).count = 1
				ghostChanges = true
			}

			count -= list.size

			while (count > 0 && empty.isNotEmpty()) {
				setItem(empty.removeInt(0), item.copyWithCount(1))
				count--
			}

			if (count > 0) {
				getItem(list.firstInt()).count += count
				ghostChanges = true
			}
		} else {
			// всего предметов больше, чем слотов
			val perSlot = count / slots.size
			var leftover = count - perSlot * slots.size

			for (i in slots.intIterator()) {
				val target = perSlot + leftover.coerceAtMost(1)

				if (getItem(i).count != target) {
					setItem(i, item.copyWithCount(target))
				}

				if (leftover > 0) leftover--
			}
		}

		if (ghostChanges) setChanged()
		return
	}

	// а вот тут уже проблемы
	// для упрощения задачи, выполним рекурсивное разбитие задачи на более простые,
	// где балансировка будет происходить только между пустыми слотами и предметами одного типа

	// если у нас нет пустых слотов, просто балансируем между заполненными слотами
	if (empty.isEmpty) {
		for (set in itemTypes.values) {
			balance(set)
		}

		return
	} else if (empty.size == 1) {
		// только один пустой слот, отдадим самому "жирному" предмету
		val type = itemTypes.entries.stream().max { a, b -> b.value.intStream().sum().compareTo(a.value.intStream().sum()) }.orElseThrow()
		type.value.add(empty.getInt(0))
		balance(type.value)

		for ((a, set) in itemTypes) {
			if (a !== type.key) {
				balance(set)
			}
		}

		return
	}

	// определяем общее количество предметов
	val totalCount = itemTypes.values.stream().mapToInt { it.stream().mapToInt { getItem(it).count }.sum() }.sum().toDouble()
	val totalEmpty = empty.size

	// определяем доли предметов по их количеству к общему количеству,
	// что позволит нам выделить претендентов на пустые слоты
	for (list in itemTypes.values) {
		if (empty.isEmpty) break // ошибка округления

		val perc = list.stream().mapToInt { getItem(it).count }.sum() / totalCount

		for (i in 0 until (perc * totalEmpty).roundToInt()) {
			list.add(empty.removeInt(0))
			if (empty.isEmpty) break // ошибка округления
		}
	}

	if (empty.isNotEmpty()) {
		// ошибка округления
		itemTypes.values.stream().max { a, b -> b.intStream().sum().compareTo(a.intStream().sum()) }.orElseThrow().add(empty.removeInt(0))
	}

	for (list in itemTypes.values) {
		balance(list)
	}
}

fun Container.balance(slots: Iterator<Int>) {
	balance(IntArraySet().also { it.addAll(slots) })
}

fun Container.balance(slots: Iterable<Int>) {
	balance(IntArraySet().also { it.addAll(slots) })
}

fun Container.balance(slots: IntRange) {
	balance(IntArraySet().also { it.addAll(slots) })
}

fun Container.balance(startSlot: Int = 0, endSlot: Int = containerSize - 1) {
	require(startSlot <= endSlot) { "Invalid slot range: $startSlot .. $endSlot" }
	var any = false

	for (i in startSlot .. endSlot) {
		if (getItem(i).isNotEmpty) {
			any = true
			break
		}
	}

	if (!any) return
	balance(IntArraySet(endSlot - startSlot + 1).also { for (i in startSlot .. endSlot) it.add(i) }, false)
}
