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

import com.google.common.collect.ImmutableList
import com.mojang.serialization.MapCodec
import com.mojang.serialization.codecs.RecordCodecBuilder
import it.unimi.dsi.fastutil.objects.ObjectArrayList
import net.minecraft.world.item.ItemStack
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.neoforged.neoforge.common.loot.IGlobalLootModifier
import net.neoforged.neoforge.common.loot.LootModifier
import java.util.*
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(): MapCodec<out IGlobalLootModifier> {
		return CODEC
	}

	companion object {
		val CODEC: MapCodec<LootPoolAppender> =
			RecordCodecBuilder.mapCodec {
				codecStart(it).and(
					LootPool.CODEC.listOf().fieldOf("pools").forGetter(LootPoolAppender::pools)
				).apply(it, ::LootPoolAppender)
			}
	}
}
