package ru.dbotthepony.mc.otm.capability.drive

import it.unimi.dsi.fastutil.objects.Object2ObjectOpenCustomHashMap
import it.unimi.dsi.fastutil.objects.Object2ObjectOpenHashMap
import it.unimi.dsi.fastutil.objects.ObjectArraySet
import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet
import kotlin.jvm.JvmOverloads
import java.util.UUID
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.ListTag
import net.minecraft.nbt.Tag
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.mc.otm.core.*
import ru.dbotthepony.mc.otm.core.math.BigInteger
import ru.dbotthepony.mc.otm.core.math.isPositive
import ru.dbotthepony.mc.otm.core.math.serializeNBT
import ru.dbotthepony.mc.otm.core.nbt.map
import ru.dbotthepony.mc.otm.core.nbt.set
import ru.dbotthepony.mc.otm.storage.*
import java.math.BigInteger
import java.util.ArrayList
import java.util.stream.Stream

abstract class AbstractMatteryDrive<T : StorageStack<T>>(
	override var driveCapacity: BigInteger,
	override val uuid: UUID = UUID.randomUUID(),
	var maxDifferentStacks: Int = 0xFFFF
) : IMatteryDrive<T> {
	protected val stack2tuples = Object2ObjectOpenCustomHashMap<T, IStorageTuple.IMutable<T>>(StorageStack.Companion)
	protected val id2tuples = Object2ObjectOpenHashMap<UUID, IStorageTuple.IMutable<T>>()

	override var isDirty = false
		set(value) {
			if (value != field && value) {
				DrivePool.markDirty(uuid)
			}

			field = value
		}

	var storedDifferentStacks = 0
		protected set

	override var storedCount: BigInteger = BigInteger.ZERO
		protected set

	protected val listeners = ObjectLinkedOpenHashSet<IStorageEventConsumer<T>>()

	override fun addListener(listener: IStorageEventConsumer<T>): Boolean {
		return listeners.add(listener)
	}

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

	override fun equals(other: Any?): Boolean {
		return other is AbstractMatteryDrive<*> && storageType == other.storageType && uuid == other.uuid
	}

	override fun hashCode(): Int {
		return (uuid.hashCode() * 31) xor storageType.hashCode()
	}

	override fun insertStack(stack: T, simulate: Boolean): T {
		val maxInsert = driveCapacity.minus(storedCount).coerceAtMost(stack.count)
		if (maxInsert <= BigInteger.ZERO) return stack

		val tuple = stack2tuples[stack]

		if (tuple != null) {
			if (!simulate) {
				tuple.grow(maxInsert)
				storedCount += maxInsert

				for (listener in listeners) {
					listener.onStackChanged(tuple.stack, tuple.id)
				}

				isDirty = true
			}

			return stack.shrink(maxInsert)
		}

		if (storedDifferentStacks >= maxDifferentStacks) {
			return stack
		}

		if (!simulate) {
			storedDifferentStacks++
			storedCount = storedCount.plus(maxInsert)

			val state = IStorageTuple.M(UUID.randomUUID(), stack.copy(maxInsert))
			stack2tuples[stack] = state
			id2tuples[state.id] = state

			for (listener in listeners) {
				listener.onStackAdded(state.stack, state.id, this)
			}

			isDirty = true
		}

		return stack.shrink(maxInsert)
	}

	override fun extractStack(id: UUID, amount: BigInteger, simulate: Boolean): T {
		val get = id2tuples[id] ?: return storageType.empty

		@Suppress("NAME_SHADOWING")
		var amount = amount

		if (!amount.isPositive)
			amount = get.stack.maxStackSize ?: get.stack.count

		amount = amount.coerceAtMost(get.stack.count)

		if (amount <= BigInteger.ZERO)
			return storageType.empty

		val copy = get.stack.copy(amount)

		if (!simulate) {
			if (amount == get.stack.count) {
				stack2tuples.remove(get.stack) ?: throw IllegalStateException("Can't find storage key for ${get.stack}")

				storedDifferentStacks--

				for (listener in listeners) {
					listener.onStackRemoved(get.id)
				}
			}

			storedCount -= amount
			get.shrink(amount)

			if (get.stack.count.isPositive) {
				for (listener in listeners) {
					listener.onStackChanged(get.stack, get.id)
				}
			}

			isDirty = true
		}

		return copy
	}

	protected abstract fun serializeStack(item: T): CompoundTag?
	protected abstract fun deserializeStack(tag: CompoundTag): T?

	override fun serializeNBT(): CompoundTag {
		val compound = CompoundTag()

		compound["capacity"] = driveCapacity.serializeNBT()
		compound["max_different_stacks"] = maxDifferentStacks

		val list = ListTag()
		compound["items"] = list

		for (tuple in stack2tuples.values) {
			val serialized = serializeStack(tuple.stack)

			if (serialized != null) {
				list.add(serialized)
			}
		}

		return compound
	}

	override fun deserializeNBT(nbt: CompoundTag) {
		for (listener in listeners) {
			for (get in stack2tuples.values) {
				listener.onStackRemoved(get.id)
			}
		}

		stack2tuples.clear()
		id2tuples.clear()
		storedCount = BigInteger.ZERO
		storedDifferentStacks = 0
		// nextID = 0L

		driveCapacity = nbt.map("capacity", ::BigInteger) ?: driveCapacity
		maxDifferentStacks = nbt.getInt("max_different_stacks")

		for (entry in nbt.getList("items", Tag.TAG_COMPOUND.toInt())) {
			if (entry is CompoundTag) {
				val stack = deserializeStack(entry)

				if (stack != null && stack.isNotEmpty) {
					storedCount += stack.count
					storedDifferentStacks++
					val tuple = IStorageTuple.M(UUID.randomUUID(), stack)
					stack2tuples[tuple.stack] = tuple
					id2tuples[tuple.id] = tuple
				}
			}
		}
	}

	override fun get(id: UUID): T {
		return id2tuples[id]?.stack ?: storageType.empty
	}

	override val stacks: Stream<IStorageTuple<T>> get() {
		return ArrayList<IStorageTuple<T>>(stack2tuples.size).also { it.addAll(stack2tuples.values) }.stream()
	}
}
