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

import it.unimi.dsi.fastutil.objects.ObjectArraySet
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 : IStorageStack> @JvmOverloads constructor(
	override var driveCapacity: BigInteger,
	override val uuid: UUID = UUID.randomUUID(),
	var maxDifferentStacks: Int = 0xFFFF
) : IMatteryDrive<T> {
	protected val tuples = HashMap<T, IStorageTuple<T>>()
	protected val tuplesByID: MutableMap<UUID, IStorageTuple<T>> = HashMap()

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

			field = value
		}

	var storedDifferentStacks = 0
		protected set

	override var storedCount: BigInteger = BigInteger.ZERO
		protected set

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

		val key = stack.key()
		val tuple = tuples[key]

		if (tuple != null) {
			if (!simulate) {
				val oldCount = tuple.stack.count
				tuple.stack.grow(maxInsert)
				storedCount += maxInsert

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

				isDirty = true
			}

			val copy = stack.copy() as T
			copy.shrink(maxInsert)
			return copy
		}

		if (storedDifferentStacks >= maxDifferentStacks) {
			return stack
		}

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

			val copy = stack.copy() as T
			copy.count = maxInsert

			val state = StorageTuple(UUID.randomUUID(), copy)
			tuples[key] = state
			tuplesByID[state.id] = state

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

			isDirty = true
		}

		val copy = stack.copy() as T
		copy.shrink(maxInsert)
		return copy
	}

	@Suppress("unchecked_cast")
	override fun extractStack(id: UUID, amount: BigInteger, simulate: Boolean): T {
		val get = tuplesByID[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() as T
		copy.count = amount

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

				storedDifferentStacks--

				for (listener in listeners) {
					listener.removeStack(get.stack, get.id)
				}
			}

			storedCount -= amount
			val oldCount = get.stack.count
			get.stack.shrink(amount)

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

			isDirty = true
		}

		return copy
	}

	protected abstract fun serializeStack(item: IStorageTuple<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 (stack in tuples.values) {
			val serialized = serializeStack(stack)

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

		return compound
	}

	override fun deserializeNBT(nbt: CompoundTag) {
		for (listener in listeners) {
			for (get in tuples.values) {
				listener.removeStack(get.stack, get.id)
			}
		}

		tuples.clear()
		tuplesByID.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) {
					storedCount += stack.count
					storedDifferentStacks++
					val tuple = StorageTuple(UUID.randomUUID(), stack)
					tuples[tuple.stack.key()] = tuple
					tuplesByID[tuple.id] = tuple
				}
			}
		}
	}

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

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

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

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

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

	companion object {
		private val LOGGER = LogManager.getLogger()
	}
}
