package ru.dbotthepony.mc.otm.data.loot

import com.google.common.collect.ImmutableList
import com.google.gson.*
import com.mojang.serialization.Codec
import com.mojang.serialization.DataResult
import com.mojang.serialization.Dynamic
import com.mojang.serialization.JsonOps
import com.mojang.serialization.codecs.RecordCodecBuilder
import it.unimi.dsi.fastutil.objects.ObjectArrayList
import net.minecraft.resources.ResourceLocation
import net.minecraft.world.item.ItemStack
import net.minecraft.world.level.storage.loot.Deserializers
import net.minecraft.world.level.storage.loot.LootContext
import net.minecraft.world.level.storage.loot.LootPool
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition
import net.minecraftforge.common.ForgeHooks
import net.minecraftforge.common.loot.IGlobalLootModifier
import net.minecraftforge.common.loot.LootModifier
import java.util.Arrays
import java.util.Deque
import java.util.stream.Stream

class LootPoolAppender(conditions: Array<out LootItemCondition>, pools: Stream<LootPool>) : LootModifier(conditions) {
	private val pools = pools.collect(ImmutableList.toImmutableList())

	constructor(conditions: Array<out LootItemCondition>, pools: Collection<LootPool>) : this(conditions, pools.stream())
	constructor(conditions: Array<out LootItemCondition>, vararg pools: LootPool) : this(conditions, Arrays.stream(pools))

	override fun doApply(generatedLoot: ObjectArrayList<ItemStack>, context: LootContext): ObjectArrayList<ItemStack> {
		for (pool in pools) {
			pool.addRandomItems(generatedLoot::add, context)
		}

		return generatedLoot
	}

	override fun codec(): Codec<out IGlobalLootModifier> {
		return CODEC
	}

	// TODO: remove reflection once Forge implement a way to provide lootpool deserialization context in non-reflective way
	companion object {
		private val lootPoolCodec: Codec<List<LootPool>>

		private val forgeHooksLootContext by lazy {
			val field = ForgeHooks::class.java.getDeclaredField("lootContext")
			field.isAccessible = true
			field
		}

		private val lootTableContextConstructor by lazy {
			val clazz = ForgeHooks::class.java.declaredClasses.firstOrNull { it.name.contains("LootTableContext") } ?: throw NoSuchElementException("Unable to find ForgeHooks\$LootTableContext!")
			val constructor = clazz.getDeclaredConstructor(ResourceLocation::class.java, Boolean::class.java)
			constructor.isAccessible = true
			constructor
		}

		init {
			val serializer = Deserializers.createLootTableSerializer().create()
			val notExistingLocation = ResourceLocation("null", "null")

			lootPoolCodec = Codec.list(Codec.PASSTHROUGH.flatXmap({
				val dequeueHolder = forgeHooksLootContext.get(null) as ThreadLocal<Deque<Any>>
				val deque = dequeueHolder.get() ?: java.util.ArrayDeque()
				dequeueHolder.set(deque)

				try {
					deque.push(lootTableContextConstructor.newInstance(notExistingLocation, true))
					DataResult.success(serializer.fromJson(it.convert(JsonOps.INSTANCE).value, LootPool::class.java))
				} catch(err: JsonSyntaxException) {
					DataResult.error { err.message }
				} finally {
					deque.pop()
				}
			}, {
				try {
					DataResult.success(Dynamic(JsonOps.INSTANCE, serializer.toJsonTree(it)))
				} catch(err: JsonSyntaxException) {
					DataResult.error { err.message }
				}
			}))
		}

		val CODEC: Codec<LootPoolAppender> =
			RecordCodecBuilder.create {
				codecStart(it).and(
					lootPoolCodec.fieldOf("pools").forGetter(LootPoolAppender::pools)
				).apply(it, ::LootPoolAppender)
			}
	}
}
