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

import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap
import it.unimi.dsi.fastutil.ints.IntAVLTreeSet
import net.minecraft.core.BlockPos
import net.minecraft.core.Direction
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.common.util.LazyOptional
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.MatteryCapability
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.BlockRotationFreedom
import ru.dbotthepony.mc.otm.core.math.RelativeSide
import ru.dbotthepony.mc.otm.core.math.getCapability
import ru.dbotthepony.mc.otm.core.math.isPositive
import ru.dbotthepony.mc.otm.core.math.isZero
import ru.dbotthepony.mc.otm.core.math.plus
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.lang.ref.WeakReference
import java.math.BigInteger
import java.util.*
import java.util.stream.Stream

private class SlotTuple(val slot: Int, val stack: ItemStack)
private class TrackedTuple(override val stack: ItemStackWrapper, override val id: UUID) : IStorageTuple<ItemStackWrapper> {
	// 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]"
	}
}

private fun Long.clamp(): Int {
	if (this > Int.MAX_VALUE) {
		return Int.MAX_VALUE
	}

	return this.toInt()
}

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

	val energy = WorkerEnergyStorage(this::setChangedLight, MachinesConfig.STORAGE_INTERFACES)

	init {
		exposeEnergyGlobally(energy)
		savetable(::energy, ENERGY_KEY)
	}

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

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

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

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

	init {
		exposeSideless(MatteryCapability.STORAGE_NODE, cell)
		exposeAllSidesExcept(RelativeSide.FRONT, MatteryCapability.STORAGE_NODE, cell)
	}

	val filter = ItemFilter(MAX_FILTERS) { _, _, _ ->
		component?.scan()
		setChangedLight()
	}

	init {
		savetable(::filter, FILTER_KEY)
	}

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

	private var neighbour: LazyOptional<IItemHandler>? = null
	private var component: ItemHandlerComponent? = null

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

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

	fun checkSurroundings() {
		if (isRemoved)
			return

		val front = blockPos + blockState.getValue(BlockRotationFreedom.TWO.property)
		val storage = level?.getBlockEntity(front)?.getCapability(ForgeCapabilities.ITEM_HANDLER, -blockState.getValue(BlockRotationFreedom.TWO.property))?.let { if (it.isPresent) it else null }

		if (neighbour != storage) {
			neighbour = storage

			if (storage != null) {
				val ref = WeakReference(this)

				storage.addListener {
					val self = ref.get() ?: return@addListener

					if (self.neighbour === it && SERVER_IS_LIVE) {
						self.checkSurroundings()
					}
				}

				component?.let(cell::removeStorageComponent)
				component = ItemHandlerComponent(storage.orThrow())
				component!!.let(cell::addStorageComponent)
			} else {
				component?.let(cell::removeStorageComponent)
				component = null
			}
		}
	}

	companion object {
		const val MAX_FILTERS = 6 * 3
	}

	private inner class ItemHandlerComponent(private val parent: IItemHandler) : IStorageComponent<ItemStackWrapper> {
		private inner class EventsSnapshot {
			val index = HashMap<UUID, BigInteger>()

			fun change(key: UUID, diff: BigInteger) {
				if (diff.isZero)
					return

				val value = index[key]

				if (value == null) {
					index[key] = diff
				} else {
					val newvalue = value + diff

					if (newvalue.isZero) {
						index.remove(key)
					} else {
						index[key] = newvalue
					}
				}
			}

			fun apply() {
				for ((key, value) in index) {
					val tuple = checkNotNull(this@ItemHandlerComponent.index[key]) { "Tuple with ID $key is missing!" }

					val count = tuple.stack.count
					tuple.stack.count += value

					if (tuple.stack.count.isPositive) {
						for (listener in listeners) {
							listener.changeStack(tuple, count)
						}
					} else {
						for (listener in listeners) {
							listener.removeStack(tuple)
						}

						this@ItemHandlerComponent.index.remove(tuple.id)

						val tuplekey = tuple.stack.key()
						tuples.remove(tuplekey) ?: 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")
			}
		}

		private var snapshot: EventsSnapshot? = null

		override val storageType: StorageStackType<ItemStackWrapper>
			get() = ITEM_STORAGE

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

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

			return false
		}

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

		private var scanned = arrayOfNulls<ItemStack>(0)
		private var scannedMap = arrayOfNulls<TrackedTuple>(0)
		private val tuples = HashMap<ItemStackWrapper, TrackedTuple>()
		private val index = HashMap<UUID, TrackedTuple>()

		private fun removeTracked(slot: Int) {
			scanned[slot] = null
			val scannedMap = checkNotNull(scannedMap[slot]) { "Not tracking slot $slot" }
			this.scannedMap[slot] = null

			val item = scannedMap.children[slot] ?: throw IllegalStateException("${scannedMap.id} does not track $slot")
			scannedMap.children.remove(slot)
			val count = scannedMap.stack.count

			val snapshot = snapshot

			if (snapshot != null) {
				snapshot.change(scannedMap.id, -item.stack.count.toBigInteger())
			} else {
				scannedMap.stack.count -= item.stack.count.toBigInteger()

				if (scannedMap.stack.count.isPositive) {
					for (listener in listeners) {
						listener.changeStack(scannedMap, count)
					}
				} else {
					for (listener in listeners) {
						listener.removeStack(scannedMap)
					}

					index.remove(scannedMap.id)

					val key = scannedMap.stack.key()
					tuples.remove(key) ?: throw IllegalStateException("Item tuple is not present for slot $slot at ${scannedMap.stack}")
				}
			}
		}

		private fun diffTracked(slot: Int, diff: Int) {
			val scannedMap = checkNotNull(scannedMap[slot]) { "Not tracking slot $slot" }
			val item = checkNotNull(scannedMap.children[slot]) { "${scannedMap.id} does not track $slot" }

			val oldCount = scannedMap.stack.count
			item.stack.count += diff

			val snapshot = snapshot

			if (snapshot != null) {
				snapshot.change(scannedMap.id, diff.toBigInteger())
			} else {
				scannedMap.stack.count += diff.toBigInteger()

				for (listener in listeners) {
					listener.changeStack(scannedMap.stack, scannedMap.id, oldCount)
				}
			}
		}

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

			val storageStack = ItemStackWrapper(stack)
			val key = storageStack.key()
			var tuple: TrackedTuple? = tuples[key]
			val added = tuple == null
			var oldCount = BigInteger.ZERO

			if (added) {
				tuple = TrackedTuple(storageStack, UUID.randomUUID())
				index[tuple.id] = tuple
				tuples[key] = tuple
			} else {
				oldCount = tuple!!.stack.count

				if (snapshot == null)
					tuple.stack.count += stack.count.toBigInteger()
			}

			tuple.children[slot] = SlotTuple(slot, stack.copy())
			scanned[slot] = tuple.children[slot].stack
			scannedMap[slot] = tuple

			if (added) {
				for (listener in listeners) {
					listener.addStack(tuple.stack, tuple.id, this)
				}
			} else {
				val snapshot = snapshot

				if (snapshot != null) {
					snapshot.change(tuple.id, stack.count.toBigInteger())
				} else {
					for (listener in listeners) {
						listener.changeStack(tuple.stack, tuple.id, oldCount)
					}
				}
			}
		}

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

				if (scanned.size < parent.slots) {
					// grow
					scanned = arrayOfNulls(parent.slots)
					scannedMap = arrayOfNulls(parent.slots)

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

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

					scanned = arrayOfNulls(parent.slots)
					scannedMap = arrayOfNulls(parent.slots)

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

		fun scan(slot: Int) {
			val current = parent[slot].let { if (it.isEmpty || !filter.match(it)) null else it }
			val last = scanned[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()

			snapshot = EventsSnapshot()

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

			snapshot!!.apply()
			snapshot = null
		}

		override fun insertStack(stack: ItemStackWrapper, simulate: Boolean): ItemStackWrapper {
			if (redstoneControl.isBlockedByRedstone || energy.batteryLevel.isZero || !filter.match(stack.item))
				return stack

			val maxPossibleDemand = ITEM_STORAGE.energyPerOperation * stack.count
			val maxExtractEnergy = energy.extractEnergy(maxPossibleDemand, true)

			var leftover: ItemStackWrapper = stack.copy()
			var additional: BigInteger = BigInteger.ZERO

			if (maxExtractEnergy != maxPossibleDemand) {
				additional = (maxExtractEnergy / ITEM_STORAGE.energyPerOperation).whole
				leftover.count -= additional
			}

			for (slot in PrioritizedSlotIterator(tuples[stack.key()])) {
				val oldCount = leftover.count
				leftover = ItemStackWrapper(parent.insertItem(slot, leftover.stack, simulate))

				if (oldCount != leftover.count && !simulate) {
					energy.extractEnergy(ITEM_STORAGE.energyPerOperation * (oldCount - leftover.count), false)

					if (scanned.size <= slot) {
						sizeScan()
					}

					scan(slot)
				}

				if (leftover.isEmpty) {
					break
				}
			}

			return leftover.also {
				if (additional != BigInteger.ZERO) {
					it.count += additional
				}
			}
		}

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

		override fun extractStack(id: UUID, amount: BigInteger, simulate: Boolean): ItemStackWrapper {
			if (redstoneControl.isBlockedByRedstone || !amount.isPositive)
				return ItemStackWrapper.EMPTY

			val maxPossibleDemand = ITEM_STORAGE.energyPerOperation * amount
			val maxExtractEnergy = energy.extractEnergy(maxPossibleDemand, true)

			@Suppress("NAME_SHADOWING")
			var amount = amount

			if (maxPossibleDemand != maxExtractEnergy) {
				amount = (maxExtractEnergy / ITEM_STORAGE.energyPerOperation).whole
			}

			val intAmount = amount.toLong()
			val tuple = index[id] ?: return ItemStackWrapper.EMPTY
			var totalExtracted = 0L

			val copy = tuple.stack.copy()

			for (stack in tuple.children.values) {
				val extracted = parent.extractItem(stack.slot, stack.stack.count.coerceAtMost((intAmount - totalExtracted).clamp()), true)

				if (!extracted.isEmpty && tuple.stack.sameItem(extracted)) {
					if (!simulate) {
						parent.extractItem(stack.slot, extracted.count, false)
					}

					totalExtracted += extracted.count

					if (extracted.count != 0 && !simulate) {
						energy.extractEnergy(ITEM_STORAGE.energyPerOperation * extracted.count, false)
					}

					if (totalExtracted >= intAmount) {
						break
					}
				}
			}

			if (totalExtracted == 0L) {
				return ItemStackWrapper.EMPTY
			}

			if (!simulate) {
				scan()
			}

			copy.count = totalExtracted.toBigInteger()
			return copy
		}

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

			for (tuple in index.values) {
				listing.add(StorageTuple(tuple.id, tuple.stack))
			}

			return listing.stream()
		}
	}
}
