package ru.dbotthepony.mc.otm.container

import com.google.common.collect.ImmutableList
import com.google.common.collect.ImmutableMap
import com.google.common.collect.ImmutableSet
import it.unimi.dsi.fastutil.ints.IntAVLTreeSet
import it.unimi.dsi.fastutil.ints.IntArrayList
import it.unimi.dsi.fastutil.ints.IntSet
import it.unimi.dsi.fastutil.objects.Object2ObjectFunction
import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
import net.minecraft.world.Container
import net.minecraft.world.entity.player.Player
import net.minecraft.world.item.Item
import net.minecraft.world.item.ItemStack
import ru.dbotthepony.mc.otm.container.util.containerSlot
import ru.dbotthepony.mc.otm.container.util.iterator
import ru.dbotthepony.mc.otm.core.collect.concatIterators
import ru.dbotthepony.mc.otm.core.collect.filter
import ru.dbotthepony.mc.otm.core.collect.flatMap
import ru.dbotthepony.mc.otm.core.collect.map
import ru.dbotthepony.mc.otm.core.isNotEmpty
import ru.dbotthepony.mc.otm.core.stream
import java.util.stream.Stream

class CombinedContainer(containers: Stream<Pair<Container, Iterable<Int>>>) : IMatteryContainer {
	constructor(vararg containers: Container) : this(containers.stream().map { it to (0 until it.containerSize) })
	constructor(containers: Collection<Container>) : this(containers.stream().map { it to (0 until it.containerSize) })

	private inner class Slot(override val slot: Int, val outer: IContainerSlot) : IContainerSlot by outer {
		override val container: Container
			get() = this@CombinedContainer
	}

	private val slots: ImmutableList<Slot>
	private val slotsMap: ImmutableMap<Container, List<IContainerSlot>>
	private val containers: ImmutableSet<Container>
	private val fullCoverage: ImmutableList<Container>
	private val notFullCoverage: ImmutableMap<Container, List<IContainerSlot>>

	init {
		val list = ImmutableList.Builder<Slot>()
		var i = 0
		val validationMap = Reference2ObjectOpenHashMap<Container, IntSet>()
		val slotsMap = Reference2ObjectOpenHashMap<Container, ArrayList<IContainerSlot>>()

		for ((container, slots) in containers) {
			val validator = validationMap.computeIfAbsent(container, Object2ObjectFunction { IntAVLTreeSet() })
			val slotList = slotsMap.computeIfAbsent(container, Object2ObjectFunction { ArrayList() })

			for (slot in slots) {
				if (validator.add(slot)) {
					val slotObj = container.containerSlot(slot)
					list.add(Slot(i++, slotObj))
					slotList.add(slotObj)
				} else {
					throw IllegalArgumentException("Duplicate mapping for $container at $i for slot $slot")
				}
			}
		}

		slots = list.build()
		this.containers = ImmutableSet.copyOf(validationMap.keys)

		this.slotsMap = slotsMap.entries
			.stream()
			.map { it.key to ImmutableList.copyOf(it.value) }
			.collect(ImmutableMap.toImmutableMap({ it.first }, { it.second }))

		this.fullCoverage = this.slotsMap.entries
			.stream()
			.filter { it.value.size == it.key.containerSize }
			.map { it.key }
			.collect(ImmutableList.toImmutableList())

		this.notFullCoverage = this.slotsMap.entries
			.stream()
			.filter { it.key !in this.fullCoverage }
			.collect(ImmutableMap.toImmutableMap({ it.key }, { it.value }))
	}

	override fun clearContent() {
		for (container in fullCoverage) {
			container.clearContent()
		}

		for (slots in notFullCoverage.values) {
			for (slot in slots) {
				slot.item = ItemStack.EMPTY
			}
		}
	}

	override fun clearSlotFilters() {
		for (container in fullCoverage) {
			if (container is IMatteryContainer)
				container.clearSlotFilters()
		}

		for (slots in notFullCoverage.values) {
			for (slot in slots) {
				slot.setFilter()
			}
		}
	}

	override fun getContainerSize(): Int {
		return slots.size
	}

	override fun isEmpty(): Boolean {
		if (fullCoverage.isNotEmpty())
			for (container in fullCoverage)
				if (!container.isEmpty)
					return false

		if (notFullCoverage.isNotEmpty())
			for (slots in notFullCoverage.values)
				for (slot in slots)
					if (!slot.isEmpty)
						return false

		return true
	}

	override fun getItem(slot: Int): ItemStack {
		// do not violate contract of getItem not throwing exceptions when index is invalid
		return slots.getOrNull(slot)?.item ?: ItemStack.EMPTY
	}

	override fun removeItem(slot: Int, amount: Int): ItemStack {
		val data = slots.getOrNull(slot) ?: return ItemStack.EMPTY
		return data.outer.container.removeItem(data.outer.slot, amount)
	}

	override fun removeItemNoUpdate(slot: Int): ItemStack {
		val data = slots.getOrNull(slot) ?: return ItemStack.EMPTY
		return data.outer.container.removeItemNoUpdate(data.outer.slot)
	}

	override fun setItem(slot: Int, itemStack: ItemStack) {
		slots.getOrNull(slot)?.item = itemStack
	}

	override fun setChanged() {
		for (container in containers) {
			container.setChanged()
		}
	}

	override fun stillValid(player: Player): Boolean {
		for (container in containers)
			if (!container.stillValid(player))
				return false

		return true
	}

	override fun iterator(nonEmpty: Boolean): Iterator<ItemStack> {
		if (notFullCoverage.isEmpty())
			return fullCoverage.iterator().flatMap { it.iterator(nonEmpty) }

		return concatIterators(
			fullCoverage.iterator().flatMap { it.iterator(nonEmpty) },
			notFullCoverage.values.iterator().flatMap { it.iterator() }.map { it.item }.let { if (nonEmpty) it.filter { it.isNotEmpty } else it }
		)
	}

	override fun slotIterator(nonEmpty: Boolean): Iterator<IContainerSlot> {
		if (nonEmpty)
			return slots.iterator().filter { it.isNotEmpty }

		return slots.iterator()
	}

	override fun containerSlot(slot: Int): IContainerSlot {
		return slots[slot]
	}

	override fun getSlotFilter(slot: Int): Item? {
		return slots[slot].getFilter()
	}

	override fun setChanged(slot: Int) {
		slots[slot].setChanged()
	}

	override fun setSlotFilter(slot: Int, filter: Item?): Boolean {
		return slots[slot].setFilter(filter)
	}

	class Builder {
		private val values = ArrayList<Pair<Container, Iterable<Int>>>()

		fun add(container: Container): Builder {
			values.add(container to container.slotRange)
			return this
		}

		fun add(container: Container, slots: Iterator<Int>): Builder {
			values.add(container to IntArrayList(slots))
			return this
		}

		fun add(container: Container, slot: Int): Builder {
			values.add(container to intArrayOf(slot).asIterable())
			return this
		}

		fun add(container: Container, from: Int, to: Int): Builder {
			values.add(container to (from .. to))
			return this
		}

		fun add(container: Container, slots: Iterable<Int>): Builder {
			values.add(container to slots)
			return this
		}

		fun build(): CombinedContainer {
			return CombinedContainer(values.stream())
		}
	}

	companion object {
		fun fromMenuSlots(slots: Iterator<net.minecraft.world.inventory.Slot>): CombinedContainer {
			val builder = Builder()

			for (slot in slots) {
				builder.add(slot.container, slot.slotIndex)
			}

			return builder.build()
		}
	}
}
