
@file:Suppress("unused")

package ru.dbotthepony.mc.otm.datagen.loot

import it.unimi.dsi.fastutil.objects.Reference2ObjectArrayMap
import it.unimi.dsi.fastutil.objects.Reference2ObjectFunction
import net.minecraft.advancements.critereon.StatePropertiesPredicate
import net.minecraft.data.DataGenerator
import net.minecraft.data.loot.LootTableProvider
import net.minecraft.data.loot.LootTableSubProvider
import net.minecraft.resources.ResourceLocation
import net.minecraft.world.level.ItemLike
import net.minecraft.world.level.block.Block
import net.minecraft.world.level.block.SlabBlock
import net.minecraft.world.level.block.state.properties.SlabType
import net.minecraft.world.level.storage.loot.LootPool
import net.minecraft.world.level.storage.loot.LootTable
import net.minecraft.world.level.storage.loot.ValidationContext
import net.minecraft.world.level.storage.loot.entries.LootItem
import net.minecraft.world.level.storage.loot.entries.LootPoolSingletonContainer
import net.minecraft.world.level.storage.loot.functions.CopyNbtFunction
import net.minecraft.world.level.storage.loot.functions.SetItemCountFunction
import net.minecraft.world.level.storage.loot.parameters.LootContextParamSet
import net.minecraft.world.level.storage.loot.parameters.LootContextParamSets
import net.minecraft.world.level.storage.loot.predicates.LootItemBlockStatePropertyCondition
import net.minecraft.world.level.storage.loot.providers.number.ConstantValue
import ru.dbotthepony.mc.otm.block.entity.MatteryBlockEntity
import ru.dbotthepony.mc.otm.block.entity.MatteryWorkerBlockEntity
import ru.dbotthepony.mc.otm.core.stream
import ru.dbotthepony.mc.otm.data.loot.CopyTileNbtFunction

data class NbtCopy(val source: String, val destination: String, val strategy: CopyNbtFunction.MergeStrategy = CopyNbtFunction.MergeStrategy.REPLACE)

fun TileNbtCopy(source: String, strategy: CopyNbtFunction.MergeStrategy = CopyNbtFunction.MergeStrategy.REPLACE): NbtCopy {
	return NbtCopy(source, "BlockEntityTag.$source", strategy)
}

class LootTables(generator: DataGenerator) : LootTableProvider(generator.packOutput, setOf() /* because we don't fucking validate you fuck */, listOf() /* because we attach everything after class is constructed duh */) {
	private val providersTable = Reference2ObjectArrayMap<LootContextParamSet, HashMap<ResourceLocation, () -> LootTable.Builder>>()

	fun builder(context: LootContextParamSet, id: ResourceLocation, provider: LootTable.Builder.() -> Unit) {
		provider(context, id) {
			LootTable.lootTable().also(provider)
		}
	}

	fun provider(context: LootContextParamSet, id: ResourceLocation, provider: () -> LootTable.Builder) {
		check(providersTable
			.computeIfAbsent(context, Reference2ObjectFunction { HashMap() })
			.put(id, provider) == null) { "Duplicate loot pool entry for $id" }
	}

	override fun getTables(): List<SubProviderEntry> {
		return providersTable.entries.stream().map { entry ->
			SubProviderEntry({
				LootTableSubProvider {
					for ((id, callback) in entry.value) {
						it.accept(id, callback.invoke())
					}
				}
			}, entry.key)
		}.toList()
	}

	override fun validate(map: MutableMap<ResourceLocation, LootTable>, validationtracker: ValidationContext) {}

	fun createSlabItemTable(block: Block, configurator: LootPoolSingletonContainer.Builder<*>.() -> Unit = {}) {
		builder(LootContextParamSets.BLOCK, block.lootTable) {
			withPool(
				LootPool.lootPool().setRolls(ConstantValue.exactly(1.0f)).add(
					LootItem.lootTableItem(block).apply(
						SetItemCountFunction.setCount(
							ConstantValue.exactly(2.0f)
						).`when`(
							LootItemBlockStatePropertyCondition.hasBlockStateProperties(block).setProperties(
								StatePropertiesPredicate.Builder.properties()
									.hasProperty(SlabBlock.TYPE, SlabType.DOUBLE)
							)
						)
					).also(configurator)
				)
			)
		}
	}

	fun createSlabItemTable(blocks: Collection<Block>, configurator: LootPoolSingletonContainer.Builder<*>.() -> Unit = {}) {
		for (block in blocks) {
			createSlabItemTable(block, configurator)
		}
	}

	fun blockProvider(block: Block, provider: () -> LootTable.Builder) {
		provider(LootContextParamSets.BLOCK, block.lootTable, provider)
	}

	fun block(block: Block, provider: LootPool.Builder.() -> Unit) {
		builder(LootContextParamSets.BLOCK, block.lootTable) {
			lootPool(provider)
		}
	}

	fun singleLootPool(context: LootContextParamSet, id: ResourceLocation, block: LootPool.Builder.() -> Unit) {
		builder(context, id) {
			withPool(LootPool.lootPool().also(block))
		}
	}

	fun dropsOther(block: Block, other: ItemLike) {
		singleLootPool(LootContextParamSets.BLOCK, block.lootTable) {
			add(LootItem.lootTableItem(other))
		}
	}

	fun dropsSelf(vararg blocks: Block) {
		for (block in blocks) {
			singleLootPool(LootContextParamSets.BLOCK, block.lootTable) {
				item(block) {}
			}
		}
	}

	fun dropsSelf(block: Block, configurator: LootPoolSingletonContainer.Builder<*>.() -> Unit = {}) {
		singleLootPool(LootContextParamSets.BLOCK, block.lootTable) {
			item(block, configurator)
		}
	}

	fun dropsSelf(blocks: Collection<Block>, configurator: LootPoolSingletonContainer.Builder<*>.() -> Unit = {}) {
		for (block in blocks) {
			singleLootPool(LootContextParamSets.BLOCK, block.lootTable) {
				item(block, configurator)
			}
		}
	}

	fun tile(block: Block, vararg filterTags: String) {
		singleLootPool(LootContextParamSets.BLOCK, block.lootTable) {
			add(LootItem.lootTableItem(block).also {
				it.apply(CopyTileNbtFunction(filterTags.stream()))
			})
		}
	}
}
