package ru.dbotthepony.mc.otm.block.entity.storage

import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap
import it.unimi.dsi.fastutil.ints.IntAVLTreeSet
import it.unimi.dsi.fastutil.ints.IntArrayList
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenCustomHashMap
import net.minecraft.core.BlockPos
import net.minecraft.world.entity.player.Inventory
import net.minecraft.world.entity.player.Player
import net.minecraft.world.inventory.AbstractContainerMenu
import net.minecraft.world.item.ItemStack
import net.minecraft.world.level.Level
import net.minecraft.world.level.block.Block
import net.minecraft.world.level.block.state.BlockState
import net.minecraftforge.common.capabilities.ForgeCapabilities
import net.minecraftforge.items.IItemHandler
import ru.dbotthepony.mc.otm.*
import ru.dbotthepony.mc.otm.block.CableBlock
import ru.dbotthepony.mc.otm.block.entity.MatteryPoweredBlockEntity
import ru.dbotthepony.mc.otm.block.entity.storage.AbstractStorageImportExport.Companion.FILTER_KEY
import ru.dbotthepony.mc.otm.capability.FlowDirection
import ru.dbotthepony.mc.otm.capability.MatteryCapability
import ru.dbotthepony.mc.otm.capability.energy.ProfiledEnergyStorage
import ru.dbotthepony.mc.otm.capability.energy.WorkerEnergyStorage
import ru.dbotthepony.mc.otm.config.MachinesConfig
import ru.dbotthepony.mc.otm.container.ItemFilter
import ru.dbotthepony.mc.otm.core.*
import ru.dbotthepony.mc.otm.core.math.RelativeSide
import ru.dbotthepony.mc.otm.core.math.isPositive
import ru.dbotthepony.mc.otm.core.math.toIntSafe
import ru.dbotthepony.mc.otm.graph.storage.StorageNode
import ru.dbotthepony.mc.otm.menu.storage.StorageBusMenu
import ru.dbotthepony.mc.otm.registry.MBlockEntities
import ru.dbotthepony.mc.otm.storage.*
import java.math.BigInteger
import java.util.*
import java.util.stream.Stream

private data class SlotTuple(val slot: Int, val stack: ItemStack)

private class TrackedTuple(override var stack: ItemStorageStack, override val id: UUID) : IStorageTuple<ItemStorageStack> {
	// do not use hash map because we need keys to be iterated in natural order
	val children = Int2ObjectAVLTreeMap<SlotTuple>()

	override fun toString(): String {
		return "TrackedTuple[$stack $id]"
	}
}

class StorageBusBlockEntity(blockPos: BlockPos, blockState: BlockState) : MatteryPoweredBlockEntity(MBlockEntities.STORAGE_BUS, blockPos, blockState) {
	override fun createMenu(containerID: Int, inventory: Inventory, ply: Player): AbstractContainerMenu {
		return StorageBusMenu(containerID, inventory, this)
	}

	override val energy = ProfiledEnergyStorage(WorkerEnergyStorage(this::markDirtyFast, MachinesConfig.STORAGE_INTERFACES))
	val energyConfig = ConfigurableEnergy(energy, modesFront = FlowDirection.NONE)

	val cell: StorageNode = object : StorageNode(energy) {
		override fun onNeighbour(link: Link) {
			if (link is DirectionLink) {
				val newState = this@StorageBusBlockEntity.blockState.setValue(CableBlock.MAPPING_CONNECTION_PROP[link.direction]!!, true)

				if (newState !== this@StorageBusBlockEntity.blockState && SERVER_IS_LIVE)
					level?.setBlock(blockPos, newState, Block.UPDATE_CLIENTS)
			}
		}

		override fun onUnNeighbour(link: Link) {
			if (link is DirectionLink) {
				val newState = this@StorageBusBlockEntity.blockState.setValue(CableBlock.MAPPING_CONNECTION_PROP[link.direction]!!, false)

				if (newState !== this@StorageBusBlockEntity.blockState && SERVER_IS_LIVE)
					level?.setBlock(blockPos, newState, Block.UPDATE_CLIENTS)
			}
		}
	}

	private var component: ItemHandlerComponent? = null

	var insertPriority = 0
		set(value) {
			field = value
			markDirtyFast()
		}

	var extractPriority = 0
		set(value) {
			field = value
			markDirtyFast()
		}

	var mode: FlowDirection = FlowDirection.BI_DIRECTIONAL
		set(value) {
			field = value
			markDirtyFast()
		}

	init {
		savetables.stateful(::energy, ENERGY_KEY)
		savetables.int(::insertPriority)
		savetables.int(::extractPriority)
		exposeGlobally(MatteryCapability.STORAGE_NODE, cell) { it != RelativeSide.FRONT }

		side(RelativeSide.FRONT).track(ForgeCapabilities.ITEM_HANDLER).addListener {
			component?.let(cell::removeStorageComponent)
			component = if (it.isPresent) {
				ItemHandlerComponent(it.orThrow()).also { if (!redstoneControl.isBlockedByRedstone) cell.addStorageComponent(it) }
			} else {
				null
			}
		}

		redstoneControl.addListener {
			val component = component ?: return@addListener

			if (it) {
				cell.removeStorageComponent(component)
			} else {
				cell.addStorageComponent(component)
			}
		}
	}

	val filter = ItemFilter(MAX_FILTERS) {
		component?.scan()
		markDirtyFast()
	}

	init {
		savetables.stateful(::filter, FILTER_KEY)
	}

	override fun setLevel(level: Level) {
		super.setLevel(level)
		cell.discover(this)
	}

	override fun tick() {
		super.tick()
		component?.scan()
		cell.tickEnergyDemanding()
	}

	override fun setRemoved() {
		super.setRemoved()
		cell.isValid = false
	}

	companion object {
		const val MAX_FILTERS = 6 * 3
	}

	private inner class ItemHandlerComponent(private val parent: IItemHandler) : IStorageComponent<ItemStorageStack> {
		private fun updateTuple(tuple: TrackedTuple, diff: BigInteger) {
			val stack = tuple.stack
			tuple.stack = tuple.stack.grow(diff)

			if (tuple.stack.isNotEmpty) {
				for (listener in listeners) {
					listener.onStackChanged(tuple)
				}
			} else {
				for (listener in listeners) {
					listener.onStackRemoved(tuple)
				}

				id2tuples.remove(tuple.id)
				items2tuples.remove(stack) ?: throw IllegalStateException("Cross-reference integrity check failed for $tuple")
			}
		}

		private inner class PrioritizedSlotIterator(tuple: TrackedTuple?) : Iterator<Int> {
			private val tupleIndexes = IntAVLTreeSet()

			init {
				if (tuple != null) {
					for (child in tuple.children) {
						tupleIndexes.add(child.key)
					}
				}
			}

			private val seenIterator = tupleIndexes.iterator()
			private val regularIterator = (0 until parent.slots).iterator()

			private var storedValue = -1

			init {
				if (tuple == null) {
					storedValue = regularIterator.nextInt()
				}
			}

			override fun hasNext(): Boolean {
				return storedValue != -1 || seenIterator.hasNext()
			}

			private fun advance() {
				if (!regularIterator.hasNext()) {
					storedValue = -1
					return
				}

				var value = regularIterator.nextInt()
				var removed = tupleIndexes.remove(value)

				while (removed && regularIterator.hasNext()) {
					value = regularIterator.nextInt()
					removed = tupleIndexes.remove(value)
				}

				if (!removed) {
					storedValue = value
				}
			}

			override fun next(): Int {
				if (storedValue != -1) {
					val value = storedValue
					advance()
					return value
				} else if (seenIterator.hasNext()) {
					val value = seenIterator.nextInt()

					if (!seenIterator.hasNext()) {
						advance()
					}

					return value
				}

				throw NoSuchElementException("Reached end of iteration")
			}
		}

		override val storageType: StorageStack.Type<ItemStorageStack>
			get() = StorageStack.ITEMS

		private val listeners = ArrayList<IStorageEventConsumer<ItemStorageStack>>()

		override val insertPriority: Int
			get() = this@StorageBusBlockEntity.insertPriority
		override val extractPriority: Int
			get() = this@StorageBusBlockEntity.extractPriority

		override fun addListener(listener: IStorageEventConsumer<ItemStorageStack>): Boolean {
			if (!listeners.contains(listener)) {
				listeners.add(listener)
				return true
			}

			return false
		}

		override fun removeListener(listener: IStorageEventConsumer<ItemStorageStack>): Boolean {
			return listeners.remove(listener)
		}

		private var slot2itemStack = arrayOfNulls<ItemStack>(0)
		private var slot2tuple = arrayOfNulls<TrackedTuple>(0)
		private val items2tuples = Object2ObjectOpenCustomHashMap<ItemStorageStack, TrackedTuple>(StorageStack.Companion)
		private val id2tuples = HashMap<UUID, TrackedTuple>()

		private fun removeTracked(slot: Int) {
			slot2itemStack[slot] = null
			val tuple = checkNotNull(slot2tuple[slot]) { "Not tracking slot $slot" }
			slot2tuple[slot] = null

			val slotTuple = tuple.children[slot] ?: throw IllegalStateException("${tuple.id} does not track $slot")
			tuple.children.remove(slot)
			updateTuple(tuple, -slotTuple.stack.count.toBigInteger())
		}

		private fun diffTracked(slot: Int, diff: Int) {
			val tuple = checkNotNull(slot2tuple[slot]) { "Not tracking slot $slot" }
			val slotTuple = checkNotNull(tuple.children[slot]) { "${tuple.id} does not track $slot" }
			slotTuple.stack.count += diff
			updateTuple(tuple, diff.toBigInteger())
		}

		private fun addTracked(slot: Int, stack: ItemStack) {
			check(slot2tuple[slot] == null) { "Already tracking slot $slot" }

			val storageStack = ItemStorageStack(stack)
			var tuple = items2tuples[storageStack]
			val added = tuple == null

			if (tuple == null) {
				tuple = TrackedTuple(storageStack, UUID.randomUUID())
				id2tuples[tuple.id] = tuple
				items2tuples[storageStack] = tuple
			} else {
				tuple.stack = tuple.stack.grow(stack.count.toBigInteger())
			}

			tuple.children[slot] = SlotTuple(slot, stack.copy())
			slot2itemStack[slot] = tuple.children[slot].stack
			slot2tuple[slot] = tuple

			if (added) {
				for (listener in listeners) {
					listener.onStackAdded(tuple.stack, tuple.id, this)
				}
			} else {
				for (listener in listeners) {
					listener.onStackChanged(tuple.stack, tuple.id)
				}
			}
		}

		private fun sizeScan() {
			if (slot2itemStack.size != parent.slots) {
				val old = slot2itemStack
				val oldMap = slot2tuple

				if (slot2itemStack.size < parent.slots) {
					// grow
					slot2itemStack = arrayOfNulls(parent.slots)
					slot2tuple = arrayOfNulls(parent.slots)

					for ((i, item) in old.withIndex()) {
						slot2itemStack[i] = item
					}

					for ((i, item) in oldMap.withIndex()) {
						slot2tuple[i] = item
					}
				} else {
					// shrink
					for (i in parent.slots until slot2itemStack.size) {
						if (slot2tuple[i] != null)
							removeTracked(i)
					}

					slot2itemStack = arrayOfNulls(parent.slots)
					slot2tuple = arrayOfNulls(parent.slots)

					for (i in 0 until parent.slots) {
						slot2itemStack[i] = old[i]
						slot2tuple[i] = oldMap[i]
					}
				}
			}
		}

		fun scan(slot: Int) {
			val current = parent[slot].let { if (it.isEmpty || !filter.match(it)) null else it }
			val last = slot2itemStack[slot]

			if (current == null && last != null) {
				removeTracked(slot)
			} else if (current != null && last == null) {
				addTracked(slot, current)
			} else if (current != null && last != null) {
				if (!ItemStack.isSameItemSameTags(current, last)) {
					removeTracked(slot)
					addTracked(slot, current)
				} else if (current.count != last.count) {
					diffTracked(slot, current.count - last.count)
				}
			}
		}

		fun scan() {
			sizeScan()

			for (slot in 0 until parent.slots) {
				scan(slot)
			}
		}

		override fun insertStack(stack: ItemStorageStack, simulate: Boolean): ItemStorageStack {
			if (redstoneControl.isBlockedByRedstone || energy.batteryLevel.isZero || !filter.match(stack.toItemStack()) || !mode.input)
				return stack

			val required = StorageStack.ITEMS.energyPerInsert(stack)
			if (energy.extractEnergy(required, true) != required) return stack

			var leftover = stack

			for (slot in PrioritizedSlotIterator(items2tuples[stack])) {
				val oldCount = leftover.count
				leftover = ItemStorageStack.unsafe(parent.insertItem(slot, leftover.toItemStack(), simulate))

				if (oldCount != leftover.count && !simulate) {
					energy.extractEnergy(StorageStack.ITEMS.energyPerInsert(stack.copy(oldCount - leftover.count)), false)
					if (slot2itemStack.size <= slot) sizeScan()
					scan(slot)
				}

				if (leftover.isEmpty) {
					break
				}
			}

			return leftover
		}

		override fun get(id: UUID): ItemStorageStack {
			return id2tuples[id]?.stack ?: ItemStorageStack.EMPTY
		}

		override fun extractStack(id: UUID, amount: BigInteger, simulate: Boolean): ItemStorageStack {
			if (redstoneControl.isBlockedByRedstone || !amount.isPositive || !energy.batteryLevel.isPositive || !mode.output)
				return ItemStorageStack.EMPTY

			var total = BigInteger.ZERO
			val tuple = id2tuples[id] ?: return ItemStorageStack.EMPTY
			val slots = tuple.children.values.iterator()
			val lstack = tuple.stack
			val affectedSlots = IntArrayList()

			while (amount > total && slots.hasNext() && energy.batteryLevel.isPositive) {
				val (slot, stack) = slots.next()
				val extracted = parent.extractItem(slot, stack.count.coerceAtMost((amount - total).toIntSafe()), true)
				val required = StorageStack.ITEMS.energyPerExtract(ItemStorageStack.unsafe(extracted))

				if (extracted.isNotEmpty && lstack.equalsWithoutCount(ItemStorageStack.unsafe(extracted)) && energy.extractEnergy(required, true) == required) {
					if (simulate) {
						total += extracted.count.toBigInteger()
					} else {
						affectedSlots.add(slot)
						val extracted2 = parent.extractItem(slot, extracted.count, false)
						total += extracted2.count.toBigInteger()

						if (extracted2.count == extracted.count) {
							energy.extractEnergy(required, false)
						} else {
							energy.extractEnergy(StorageStack.ITEMS.energyPerExtract(ItemStorageStack.unsafe(extracted2)), false)
						}
					}
				}
			}

			val i = affectedSlots.intIterator()
			while (i.hasNext()) scan(i.nextInt())
			return lstack.copy(total)
		}

		override val stacks: Stream<IStorageTuple<ItemStorageStack>> get() {
			val listing = ArrayList<IStorageTuple<ItemStorageStack>>(id2tuples.size)

			for (tuple in id2tuples.values) {
				listing.add(IStorageTuple.I(tuple.id, tuple.stack))
			}

			return listing.stream()
		}
	}
}
