package ru.dbotthepony.mc.otm.core.util

import com.google.common.collect.ImmutableList
import com.mojang.serialization.Codec
import net.minecraft.core.HolderLookup
import net.minecraft.core.HolderLookup.Provider
import net.minecraft.nbt.ByteTag
import net.minecraft.nbt.CompoundTag
import net.minecraft.nbt.DoubleTag
import net.minecraft.nbt.FloatTag
import net.minecraft.nbt.IntTag
import net.minecraft.nbt.LongTag
import net.minecraft.nbt.NbtOps
import net.minecraft.nbt.NumericTag
import net.minecraft.nbt.StringTag
import net.minecraft.nbt.Tag
import net.minecraft.resources.ResourceLocation
import net.neoforged.neoforge.common.util.INBTSerializable
import org.apache.logging.log4j.LogManager
import ru.dbotthepony.kommons.util.Delegate
import ru.dbotthepony.mc.otm.core.get
import ru.dbotthepony.mc.otm.core.immutableList
import ru.dbotthepony.mc.otm.core.math.Decimal
import ru.dbotthepony.mc.otm.core.math.Vector
import ru.dbotthepony.mc.otm.core.nbt.set
import java.util.function.Supplier
import kotlin.jvm.optionals.getOrNull
import kotlin.reflect.KMutableProperty0
import kotlin.reflect.KProperty0

/**
 * Utility class to manage serialization and deserialization of properties
 * of mutable class
 */
class Savetables : INBTSerializable<CompoundTag?> {
	private val entries = ArrayList<Entry<*, *>>()

	sealed interface Entry<V : Any?, T : Tag?> : INBTSerializable<T?> {
		val name: String
		val type: Class<T>
		fun validate()
	}

	inline fun <V : INBTSerializable<T>, reified T : Tag?> stateful(getter: Supplier<V>, name: String): Stateful<V, T> {
		return stateful(getter, name, T::class.java)
	}

	inline fun <V : INBTSerializable<T>, reified T : Tag?> stateful(getter: V, name: String): Stateful<V, T> {
		return stateful(getter, name, T::class.java)
	}

	inline fun <V : INBTSerializable<T>, reified T : Tag?> stateful(getter: KProperty0<V>, name: String = getter.name): Stateful<V, T> {
		return stateful(getter, name, T::class.java)
	}

	inline fun <V : INBTSerializable<T>, reified T : Tag?> stateful(values: Collection<V>, name: String): ImmutableList<Stateful<V, T>> {
		val iterator = values.iterator()
		return immutableList(values.size) {
			stateful(iterator.next(), "${name}_$it", T::class.java)
		}
	}

	fun <V : INBTSerializable<T>, T : Tag?> stateful(getter: Supplier<V>, name: String, type: Class<T>): Stateful<V, T> {
		return Stateful(getter, name, type)
			.withSerializer { it, l -> it.serializeNBT(l) }
			.withDeserializer { v, t, l -> v.deserializeNBT(l, t) }
	}

	fun <V : INBTSerializable<T>, T : Tag?> stateful(getter: V, name: String, type: Class<T>): Stateful<V, T> {
		return Stateful({ getter }, name, type)
			.withSerializer { it, l -> it.serializeNBT(l) }
			.withDeserializer { v, t, l -> v.deserializeNBT(l, t) }
	}

	fun <V : INBTSerializable<T>, T : Tag?> stateful(getter: KProperty0<V>, name: String = getter.name, type: Class<T>): Stateful<V, T> {
		return Stateful(getter, name, type)
			.withSerializer { it, l -> it.serializeNBT(l) }
			.withDeserializer { v, t, l -> v.deserializeNBT(l, t) }
	}

	fun decimal(prop: Delegate<Decimal>, name: String, default: Decimal = Decimal.ZERO): Stateless<Decimal, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withDeserializer { Decimal.deserializeNBT(it) }
			.withSerializer { it.serializeNBT() }
			.withDefault { default }
	}

	fun decimalNullable(prop: Delegate<Decimal?>, name: String, default: Decimal = Decimal.ZERO): Stateless<Decimal?, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withDeserializer { Decimal.deserializeNBT(it) }
			.withSerializer { it?.serializeNBT() }
			.withDefault { default }
	}

	fun decimal(prop: KMutableProperty0<Decimal>, name: String = prop.name, default: Decimal = Decimal.ZERO): Stateless<Decimal, Tag> {
		return decimal(Delegate.Of(prop), name, default)
	}

	fun decimalNullable(prop: KMutableProperty0<Decimal?>, name: String = prop.name, default: Decimal = Decimal.ZERO): Stateless<Decimal?, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withDeserializer { Decimal.deserializeNBT(it) }
			.withSerializer { it?.serializeNBT() }
			.withDefault { default }
	}

	fun float(prop: Delegate<Float>, name: String): Stateless<Float, NumericTag> {
		return Stateless(prop, name, NumericTag::class.java)
			.withSerializer { FloatTag.valueOf(it) }
			.withDeserializer { it.asFloat }
	}

	fun float(prop: KMutableProperty0<Float>, name: String = prop.name): Stateless<Float, NumericTag> {
		return float(Delegate.Of(prop), name)
	}

	fun double(prop: Delegate<Double>, name: String): Stateless<Double, NumericTag> {
		return Stateless(prop, name, NumericTag::class.java)
			.withSerializer { DoubleTag.valueOf(it) }
			.withDeserializer { it.asDouble }
	}

	fun double(prop: KMutableProperty0<Double>, name: String = prop.name): Stateless<Double, NumericTag> {
		return double(Delegate.Of(prop), name)
	}

	fun int(prop: Delegate<Int>, name: String): Stateless<Int, NumericTag> {
		return Stateless(prop, name, NumericTag::class.java)
			.withSerializer { IntTag.valueOf(it) }
			.withDeserializer { it.asInt }
	}

	fun int(prop: KMutableProperty0<Int>, name: String = prop.name): Stateless<Int, NumericTag> {
		return int(Delegate.Of(prop), name)
	}

	fun long(prop: Delegate<Long>, name: String): Stateless<Long, NumericTag> {
		return Stateless(prop, name, NumericTag::class.java)
			.withSerializer { LongTag.valueOf(it) }
			.withDeserializer { it.asLong }
	}

	fun long(prop: KMutableProperty0<Long>, name: String = prop.name): Stateless<Long, NumericTag> {
		return long(Delegate.Of(prop), name)
	}

	fun bool(prop: Delegate<Boolean>, name: String): Stateless<Boolean, NumericTag> {
		return Stateless(prop, name, NumericTag::class.java)
			.withSerializer { ByteTag.valueOf(it) }
			.withDeserializer { it.asByte > 0 }
	}

	fun bool(prop: KMutableProperty0<Boolean>, name: String = prop.name): Stateless<Boolean, NumericTag> {
		return bool(Delegate.Of(prop), name)
	}

	fun string(prop: Delegate<String>, name: String): Stateless<String, StringTag> {
		return Stateless(prop, name, StringTag::class.java)
			.withSerializer { StringTag.valueOf(it) }
			.withDeserializer { it.asString }
	}

	fun string(prop: KMutableProperty0<String>, name: String = prop.name): Stateless<String, StringTag> {
		return string(Delegate.Of(prop), name)
	}

	fun <E : Enum<E>> enum(prop: Delegate<E>, name: String, map: (String) -> E): Stateless<E, StringTag> {
		return Stateless(prop, name, StringTag::class.java)
			.withSerializer { StringTag.valueOf(it.name) }
			.withDeserializer { map.invoke(it.asString) }
	}

	fun <E : Enum<E>> enum(prop: KMutableProperty0<E>, name: String = prop.name, map: (String) -> E): Stateless<E, StringTag> {
		return enum(Delegate.Of(prop), name, map)
	}

	fun <T : Any> codecNullable(prop: Delegate<T?>, codec: Codec<T>, name: String): Stateless<T?, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withSerializer { prop.get()?.let { codec.encode(it, NbtOps.INSTANCE, NbtOps.INSTANCE.empty()).getOrThrow { throw IllegalStateException("Failed to write save data for '$name': $it") } } }
			.withDeserializer { codec.decode(NbtOps.INSTANCE, it).resultOrPartial { throw IllegalStateException("Failed to read save data for '$name'", RuntimeException(it)) }.getOrNull()?.first }
	}

	fun <T : Any> codecNullable(prop: KMutableProperty0<T?>, codec: Codec<T>, name: String = prop.name): Stateless<T?, Tag> {
		return codecNullable(Delegate.Of(prop), codec, name)
	}

	fun <T : Any> codecNullable(prop: Delegate<T?>, codec: Codec<T>, name: String, default: T?): Stateless<T?, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withSerializer { prop.get()?.let { codec.encode(it, NbtOps.INSTANCE, NbtOps.INSTANCE.empty()).getOrThrow { throw IllegalStateException("Failed to write save data for '$name': $it") } } }
			.withDeserializer { codec.decode(NbtOps.INSTANCE, it).resultOrPartial { LOGGER.error("Failed to read save data for '$name'", RuntimeException(it)) }.map { it.first }.orElse(default) }
			.withDefault { default }
	}

	fun <T : Any> codecNullable(prop: KMutableProperty0<T?>, codec: Codec<T>, name: String = prop.name, default: T?): Stateless<T?, Tag> {
		return codecNullable(Delegate.Of(prop), codec, name, default)
	}

	fun <T : Any> codec(prop: Delegate<T>, codec: Codec<T>, name: String): Stateless<T, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withSerializer { codec.encode(prop.get(), NbtOps.INSTANCE, NbtOps.INSTANCE.empty()).getOrThrow { throw IllegalStateException("Failed to write save data for '$name': $it") } }
			.withDeserializer {
				codec.decode(NbtOps.INSTANCE, it).let { res ->
					res
						.resultOrPartial { LOGGER.error("Failed to read save data for '$name'", RuntimeException(it)) }
						.map { it.first }
						.orElseThrow { throw RuntimeException("Failed to read save data for '$name': ${res.error().get().message()}") }
				}
			}
	}

	fun <T : Any> codec(prop: Delegate<T>, codec: Codec<T>, name: String, default: Supplier<T>): Stateless<T, Tag> {
		return Stateless(prop, name, Tag::class.java)
			.withSerializer { codec.encode(prop.get(), NbtOps.INSTANCE, NbtOps.INSTANCE.empty()).getOrThrow { throw IllegalStateException("Failed to write save data for '$name': $it") } }
			.withDeserializer {
				codec
					.decode(NbtOps.INSTANCE, it)
					.resultOrPartial { LOGGER.error("Failed to read save data for '$name'", RuntimeException(it)) }
					.map { it.first }
					.orElseGet(default)
			}
			.withDefault { default.get() }
	}

	fun <T : Any> codec(prop: KMutableProperty0<T>, codec: Codec<T>, name: String = prop.name): Stateless<T, Tag> {
		return codec(Delegate.Of(prop), codec, name)
	}

	fun <T : Any> codec(prop: KMutableProperty0<T>, codec: Codec<T>, name: String = prop.name, default: T): Stateless<T, Tag> {
		return codec(Delegate.Of(prop), codec, name, Supplier { default })
	}

	fun <T : Any> codec(prop: KMutableProperty0<T>, codec: Codec<T>, name: String = prop.name, default: Supplier<T>): Stateless<T, Tag> {
		return codec(Delegate.Of(prop), codec, name, default)
	}

	fun <T : Any> codec(prop: KMutableProperty0<T>, codec: Codec<T>, default: Supplier<T>): Stateless<T, Tag> {
		return codec(Delegate.Of(prop), codec, prop.name, default)
	}

	fun vector(prop: Delegate<Vector>, name: String, default: Vector = Vector.ZERO): Stateless<Vector, Tag> {
		return codec(prop, Vector.CODEC, name, Supplier { default })
	}

	fun vector(prop: Delegate<Vector>, name: String, default: Supplier<Vector>): Stateless<Vector, Tag> {
		return codec(prop, Vector.CODEC, name, default)
	}

	fun vector(prop: KMutableProperty0<Vector>, name: String = prop.name, default: Vector = Vector.ZERO): Stateless<Vector, Tag> {
		return vector(Delegate.Of(prop), name, default)
	}

	fun location(prop: Delegate<ResourceLocation>, name: String): Stateless<ResourceLocation, Tag> {
		return codec(prop, ResourceLocation.CODEC, name)
	}

	fun location(prop: KMutableProperty0<ResourceLocation>, name: String = prop.name): Stateless<ResourceLocation, Tag> {
		return location(Delegate.Of(prop), name)
	}

	override fun serializeNBT(registry: Provider): CompoundTag {
		return CompoundTag().also { this.serializeNBT(it, registry) }
	}

	private var validated = false

	fun validate() {
		if (validated) return

		for (entry in entries)
			entry.validate()

		validated = true
	}

	override fun deserializeNBT(registry: Provider, nbt: CompoundTag?) {
		validate()

		if (nbt == null) {
			for (entry in entries) {
				entry.deserializeNBT(registry, null)
			}
		} else {
			for (entry in entries) {
				val value = nbt[entry.name]

				if (value != null && entry.type.isAssignableFrom(value.javaClass)) {
					(entry as INBTSerializable<Tag>).deserializeNBT(registry, value)
				} else {
					entry.deserializeNBT(registry, null)
				}
			}
		}
	}

	fun serializeNBT(nbt: CompoundTag, registry: Provider) {
		validate()

		for (entry in entries) {
			val value = entry.serializeNBT(registry)

			if (value != null)
				nbt[entry.name] = value
		}
	}

	inner class Stateful<V : Any?, T : Tag?>(private val prop: Supplier<V>, override val name: String, override val type: Class<T>) : Entry<V, T> {
		constructor(field: KProperty0<V>, name: String = field.name, type: Class<T>) : this(field::get, name, type)

		init {
			check(!entries.any { it.name == name }) { "Already has entry with name $name!" }
			entries.add(this)
			validated = false
		}

		private var serializer: ((V, HolderLookup.Provider) -> T?)? = null
		private var deserializer: ((V, T & Any, HolderLookup.Provider) -> Unit)? = null
		private var resetter: ((V) -> Unit)? = null

		fun withSerializer(serializer: (V, HolderLookup.Provider) -> T?): Stateful<V, T> {
			this.serializer = serializer
			return this
		}

		fun withDeserializer(deserializer: (V, T & Any, HolderLookup.Provider) -> Unit): Stateful<V, T> {
			this.deserializer = deserializer
			return this
		}

		fun withResetter(resetter: (V) -> Unit): Stateful<V, T> {
			this.resetter = resetter
			return this
		}

		override fun serializeNBT(registry: HolderLookup.Provider): T? {
			return checkNotNull(serializer) { "No serializer specified for $name" }.invoke(prop.get(), registry)
		}

		override fun deserializeNBT(registry: Provider, nbt: T?) {
			if (nbt == null) {
				resetter?.invoke(prop.get())
			} else {
				checkNotNull(deserializer) { "No deserializer specified for $name" }.invoke(prop.get(), nbt, registry)
			}
		}

		override fun validate() {
			checkNotNull(serializer) { "No serializer specified for $name" }
			checkNotNull(deserializer) { "No deserializer specified for $name" }
		}
	}

	inner class Stateless<V : Any?, T : Tag>(private val prop: Delegate<V>, override val name: String, override val type: Class<T>) : Entry<V, T> {
		constructor(field: KMutableProperty0<V>, name: String = field.name, type: Class<T>) : this(Delegate.Of(field), name, type)

		init {
			check(!entries.any { it.name == name }) { "Already has entry with name $name!" }
			entries.add(this)
			validated = false
		}

		private var serializer: ((V) -> T?)? = null
		private var deserializer: ((T) -> V)? = null
		private var default: (() -> V)? = null

		override fun validate() {
			checkNotNull(serializer) { "No serializer specified for $name" }
			checkNotNull(deserializer) { "No deserializer specified for $name" }
		}

		fun withSerializer(serializer: (V) -> T?): Stateless<V, T> {
			this.serializer = serializer
			return this
		}

		fun withDeserializer(deserializer: (T) -> V): Stateless<V, T> {
			this.deserializer = deserializer
			return this
		}

		fun withDefault(default: () -> V): Stateless<V, T> {
			this.default = default
			return this
		}

		override fun serializeNBT(registry: Provider): T? {
			return checkNotNull(serializer) { "No serializer specified for $name" }.invoke(prop.get())
		}

		override fun deserializeNBT(registry: Provider, nbt: T?) {
			if (nbt == null) {
				if (default != null) {
					prop.accept(default!!.invoke())
				}
			} else {
				prop.accept(checkNotNull(deserializer) { "No deserializer specified for $name" }.invoke(nbt))
			}
		}
	}

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