package ru.dbotthepony.mc.otm.triggers

import com.google.gson.JsonObject
import com.google.gson.JsonParseException
import com.mojang.datafixers.util.Pair
import com.mojang.serialization.Codec
import com.mojang.serialization.DataResult
import com.mojang.serialization.DynamicOps
import com.mojang.serialization.JsonOps
import com.mojang.serialization.MapCodec
import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet
import it.unimi.dsi.fastutil.objects.Reference2ObjectFunction
import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap
import net.minecraft.advancements.Criterion
import net.minecraft.advancements.CriterionTrigger
import net.minecraft.advancements.CriterionTriggerInstance
import net.minecraft.advancements.critereon.ContextAwarePredicate
import net.minecraft.advancements.critereon.DeserializationContext
import net.minecraft.advancements.critereon.EntityPredicate
import net.minecraft.advancements.critereon.SerializationContext
import net.minecraft.resources.ResourceLocation
import net.minecraft.server.PlayerAdvancements
import net.minecraft.server.level.ServerPlayer
import ru.dbotthepony.mc.otm.core.collect.filter
import ru.dbotthepony.mc.otm.core.collect.toImmutableList
import ru.dbotthepony.mc.otm.core.set
import ru.dbotthepony.mc.otm.core.toJsonStrict
import java.util.Optional
import java.util.function.Predicate

// allows to support both 1.20.1 and 1.20.2 with ease
// and has slightly less memory footprint than vanilla SimpleCriterionTrigger
abstract class MCriterionTrigger<T : MCriterionTrigger<T>.AbstractInstance>(private val id: ResourceLocation) : CriterionTrigger<T> {
	override fun getId(): ResourceLocation {
		return id
	}

	private val listeners = Reference2ObjectOpenHashMap<PlayerAdvancements, ObjectOpenHashSet<CriterionTrigger.Listener<T>>>()

	override fun addPlayerListener(advancements: PlayerAdvancements, listener: CriterionTrigger.Listener<T>) {
		listeners.computeIfAbsent(advancements, Reference2ObjectFunction { ObjectOpenHashSet() }).add(listener)
	}

	override fun removePlayerListener(advancements: PlayerAdvancements, listener: CriterionTrigger.Listener<T>) {
		listeners[advancements]?.remove(listener)
	}

	override fun removePlayerListeners(advancements: PlayerAdvancements) {
		listeners.remove(advancements)
	}

	protected abstract val codec: Codec<T>

	override fun createInstance(data: JsonObject, context: DeserializationContext): T {
		return try {
			deserializationContext.get().addLast(context)
			codec.decode(JsonOps.INSTANCE, data).getOrThrow(false) { throw JsonParseException("Failed to decode data: $it") }.first
		} finally {
			deserializationContext.get().removeLast()
		}
	}

	fun trigger(player: ServerPlayer, predicate: Predicate<T> = Predicate { true }) {
		val advancements = player.advancements
		val listeners = listeners[advancements] ?: return

		val context = EntityPredicate.createContext(player, player)

		listeners.iterator()
			.filter { predicate.test(it.triggerInstance) && it.triggerInstance.playerPredicate.map { it.matches(context) }.orElse(true) }
			.toImmutableList()
			.forEach { it.run(advancements) }
	}

	abstract inner class AbstractInstance(val playerPredicate: Optional<ContextAwarePredicate>) : CriterionTriggerInstance {
		fun criterion() = Criterion(this as T)

		override fun getCriterion(): ResourceLocation {
			return id
		}

		final override fun serializeToJson(context: SerializationContext): JsonObject {
			return try {
				serializationContext.get().addLast(context)
				codec.toJsonStrict(this as T) as JsonObject
			} finally {
				serializationContext.get().removeLast()
			}
		}
	}

	companion object {
		protected val deserializationContext = object : ThreadLocal<ArrayDeque<DeserializationContext>>() {
			override fun initialValue(): ArrayDeque<DeserializationContext> {
				return ArrayDeque()
			}
		}

		protected val serializationContext = object : ThreadLocal<ArrayDeque<SerializationContext>>() {
			override fun initialValue(): ArrayDeque<SerializationContext> {
				return ArrayDeque()
			}
		}

		@JvmStatic
		protected val predicateCodec: Codec<ContextAwarePredicate> = object : Codec<ContextAwarePredicate> {
			override fun <T : Any?> encode(input: ContextAwarePredicate, ops: DynamicOps<T>, prefix: T): DataResult<T> {
				return DataResult.success(JsonOps.INSTANCE.convertTo(ops, input.toJson(serializationContext.get().lastOrNull() ?: return DataResult.error { "Not serializing trigger instance" })))
			}

			override fun <T : Any?> decode(ops: DynamicOps<T>, input: T): DataResult<Pair<ContextAwarePredicate, T>> {
				val context = deserializationContext.get().lastOrNull() ?: return DataResult.error { "Not current deserializing trigger instance" }

				return try {
					DataResult.success(Pair.of(EntityPredicate.fromJson(JsonObject().also { it["a"] = ops.convertTo(JsonOps.INSTANCE, input) }, "a", context), ops.empty()))
				} catch (err: Exception) {
					DataResult.error { "Failed to deserialize ContextAwarePredicate: " + err.message }
				}
			}
		}

		@JvmStatic
		protected val playerPredicateCodec: MapCodec<Optional<ContextAwarePredicate>> = predicateCodec.optionalFieldOf("player")
	}
}
