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.IntCollection
import it.unimi.dsi.fastutil.ints.IntIterable
import it.unimi.dsi.fastutil.ints.IntIterator
import it.unimi.dsi.fastutil.ints.IntList
import it.unimi.dsi.fastutil.ints.IntOpenHashSet
import it.unimi.dsi.fastutil.ints.IntSet
import it.unimi.dsi.fastutil.objects.Object2ObjectFunction
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenCustomHashMap
import it.unimi.dsi.fastutil.objects.ObjectArrayList
import net.minecraft.world.Container
import net.minecraft.world.inventory.CraftingContainer
import net.minecraft.world.item.ItemStack
import net.minecraft.world.item.enchantment.EnchantmentHelper.hasVanishingCurse
import net.minecraftforge.fluids.capability.IFluidHandler
import ru.dbotthepony.mc.otm.container.util.IContainerSlot
import ru.dbotthepony.mc.otm.container.util.ItemStackHashStrategy
import ru.dbotthepony.mc.otm.container.util.containerSlot
import ru.dbotthepony.mc.otm.container.util.slotIterator
import ru.dbotthepony.mc.otm.core.addAll
import ru.dbotthepony.mc.otm.core.collect.filter
import ru.dbotthepony.mc.otm.core.collect.toList
import ru.dbotthepony.mc.otm.core.isNotEmpty
import ru.dbotthepony.mc.otm.core.map
import ru.dbotthepony.mc.otm.core.util.ItemStackSorter
import kotlin.math.roundToInt

@Suppress("nothing_to_inline")
inline operator fun Container.set(index: Int, value: ItemStack) = setItem(index, value)

@Suppress("nothing_to_inline")
inline operator fun Container.get(index: Int): ItemStack = getItem(index)

@Suppress("nothing_to_inline")
inline operator fun IFluidHandler.get(index: Int) = getFluidInTank(index)

val Container.slotRange: IntIterable get() {
	return IntIterable {
		val i = (0 until containerSize).iterator()

		object : IntIterator {
			override fun hasNext(): Boolean {
				return i.hasNext()
			}

			override fun remove() {
				throw UnsupportedOperationException()
			}

			override fun nextInt(): Int {
				return i.nextInt()
			}
		}
	}
}

fun Container.addItem(stack: ItemStack, simulate: Boolean, slots: IntIterable = slotRange): ItemStack {
	if (this is IMatteryContainer) {
		return this.addItem(stack, simulate, slots)
	}

	if (stack.isEmpty)
		return stack

	val copy = stack.copy()

	// двигаем в одинаковые слоты
	var i = slots.intIterator()

	while (i.hasNext()) {
		val slot = i.nextInt()

		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
				}
			}
		}
	}

	// двигаем в пустые слоты
	i = slots.intIterator()

	while (i.hasNext()) {
		val slot = i.nextInt()

		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.vanishCursedItems() {
	for (slot in slotIterator()) {
		if (hasVanishingCurse(slot.item)) {
			slot.remove()
		}
	}
}

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, endSlot: Int) {
	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)
}

fun Container.balance() {
	if (isEmpty)
		return

	balance(IntArraySet(containerSize).also { for (i in 0 until containerSize) it.add(i) }, false)
}

operator fun CraftingContainer.get(column: Int, row: Int): ItemStack {
	return getItem(column + row * width)
}

operator fun CraftingContainer.get(column: Int, row: Int, flop: Boolean): ItemStack {
	return if (flop)
		get(width - column - 1, row)
	else
		get(column, row)
}

fun Container.sort(comparator: Comparator<ItemStack> = ItemStackSorter.DEFAULT) {
	if (isEmpty)
		return

	val slots = slotIterator().filter { !it.isForbiddenForAutomation && it.getMaxStackSize() >= it.item.count }.toList()

	if (slots.isEmpty())
		return

	val items = Object2ObjectOpenCustomHashMap<ItemStack, ItemStack>(ItemStackHashStrategy)

	slots.forEach {
		val get = items[it.item]

		if (get == null) {
			items[it.item] = it.item.copy()
		} else {
			get.count += it.item.count
		}
	}

	val sortedItems = ObjectArrayList(items.values)
	sortedItems.sortWith(comparator)

	slots.forEach { it.remove() }
	sortedItems.forEach { addItem(it, false) }
}

fun Container.sortWithIndices(sortedSlots: IntCollection) {
	if (sortedSlots.isEmpty() || isEmpty)
		return

	val seen = IntOpenHashSet()
	val valid = ArrayList<IContainerSlot>()

	val iterator = sortedSlots.intIterator()

	while (iterator.hasNext()) {
		val value = iterator.nextInt()

		if (value in 0 until containerSize && seen.add(value)) {
			val slot = containerSlot(value)

			if (slot.isNotEmpty && !slot.isForbiddenForAutomation && slot.item.count <= slot.getMaxStackSize()) {
				valid.add(slot)
			}
		}
	}

	if (valid.isEmpty())
		return

	val items = valid.map { it.item.copy() }
	valid.forEach { it.remove() }
	items.forEach { addItem(it, false) }
}

fun Container.computeSortedIndices(comparator: Comparator<ItemStack> = ItemStackSorter.DEFAULT): IntList {
	if (isEmpty)
		return IntList.of()

	val slots = slotIterator().filter { !it.isForbiddenForAutomation && it.getMaxStackSize() >= it.item.count }.toList()

	if (slots.isEmpty())
		return IntList.of()

	val items = Object2ObjectOpenCustomHashMap<ItemStack, Pair<ItemStack, IntList>>(ItemStackHashStrategy)

	slots.forEach {
		val get = items[it.item]

		if (get == null) {
			items[it.item] = it.item.copy() to IntArrayList().also { s -> s.add(it.slot) }
		} else {
			get.first.count += it.item.count
			get.second.add(it.slot)
		}
	}

	val sortedItems = ObjectArrayList(items.values)
	sortedItems.sortWith(comparator.map { it.first })

	val result = IntArrayList()
	sortedItems.forEach { result.addAll(it.second) }
	return result
}
