package ru.dbotthepony.mc.otm.saveddata

import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap
import it.unimi.dsi.fastutil.ints.Int2ObjectFunction
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.IntTag
import net.minecraft.nbt.ListTag
import net.minecraft.nbt.Tag
import net.minecraft.world.level.saveddata.SavedData
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.mc.otm.core.collect.ProxiedMap
import ru.dbotthepony.mc.otm.core.nbt.set
import ru.dbotthepony.mc.otm.core.util.WriteOnce

class SavedMapDelegate<V>(val parent: SavedCountingMap<SavedMapDelegate<V>>?, val index: Int, value: V) {
	constructor(value: V) : this(null, -1, value)

	var value: V = value
		set(value) {
			if (field != value) {
				field = value
				parent?.isDirty = true
			}
		}

	companion object {
		fun <V> makeMap(
			serializer: (value: V) -> Tag,
			deserializer: (nbt: Tag) -> V,
			defaultValue: V
		): SavedCountingMap<SavedMapDelegate<V>> {
			return SavedCountingMap(
				serializer = { _, value, _ -> serializer.invoke(value.value) },
				deserializer = { map, nbt, index -> SavedMapDelegate(map, index, deserializer.invoke(nbt)) },
				factory = { map, index -> SavedMapDelegate(map, index, defaultValue) }
			)
		}
	}
}

private class ProxiedCountingMap<T> : ProxiedMap<Int, T>(Int2ObjectAVLTreeMap()) {
	var parent: SavedCountingMap<T> by WriteOnce()

	override fun onClear() {
		parent.isDirty = true
	}

	override fun onValueAdded(key: Int, value: T) {
		parent.isDirty = true
	}

	override fun onValueRemoved(key: Int, value: T) {
		parent.isDirty = true
	}
}

class SavedCountingMap<T> private constructor(
	val serializer: (map: SavedCountingMap<T>, value: T, index: Int) -> Tag,
	val deserializer: (map: SavedCountingMap<T>, nbt: Tag, index: Int) -> T,
	val factory: (map: SavedCountingMap<T>, index: Int) -> T,
	private val map: ProxiedCountingMap<T>
) : SavedData(), MutableMap<Int, T> by map {
	constructor(
		serializer: (map: SavedCountingMap<T>, value: T, index: Int) -> Tag,
		deserializer: (map: SavedCountingMap<T>, nbt: Tag, index: Int) -> T,
		factory: (map: SavedCountingMap<T>, index: Int) -> T,
	) : this(serializer, deserializer, factory, ProxiedCountingMap())

	init {
		map.parent = this
	}

	var nextIndex = 0
		private set

	fun computeIfAbsent(index: Int): T {
		return map.computeIfAbsent(index, Int2ObjectFunction {
			isDirty = true
			return@Int2ObjectFunction factory.invoke(this, it)
		})
	}

	fun factorize(): T {
		return computeIfAbsent(nextIndex++)
	}

	override fun save(output: CompoundTag): CompoundTag {
		output["map"] = ListTag().also {
			for ((key, value) in this.map) {
				val compound = CompoundTag()

				try {
					compound["value"] = this.serializer.invoke(this, value, key)
					it.add(compound)
				} catch(err: Exception) {
					LOGGER.error("Unable to serialize $value at $key", err)
					continue
				}

				compound["key"] = key
			}
		}

		output["next_index"] = nextIndex

		return output
	}

	fun load(input: CompoundTag): SavedCountingMap<T> {
		this.map.clear()

		val map = input["map"] as? ListTag ?: ListTag()
		var maximal = -1

		for (value in map) {
			if (value !is CompoundTag) {
				continue
			}

			val key = (value["key"] as IntTag?)?.asInt ?: continue
			val nbt = value["value"] ?: continue

			try {
				val deserialized = this.deserializer.invoke(this, nbt, key)
				this.map.put(key, deserialized)
			} catch(err: Exception) {
				LOGGER.error("Unable to deserialize value in $this at $key", err)
			}

			maximal = maximal.coerceAtLeast(key)
		}

		nextIndex = (input["next_index"] as IntTag?)?.asInt ?: (maximal + 1)

		isDirty = false
		return this
	}

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