package ru.dbotthepony.mc.otm.datagen.tags

import it.unimi.dsi.fastutil.objects.ObjectArraySet
import net.minecraft.core.HolderLookup
import net.minecraft.core.Registry
import net.minecraft.core.registries.BuiltInRegistries
import net.minecraft.core.registries.Registries
import net.minecraft.resources.ResourceKey
import net.minecraft.resources.ResourceLocation
import net.minecraft.tags.BlockTags
import net.minecraft.tags.GameEventTags
import net.minecraft.tags.TagKey
import net.minecraft.world.item.Item
import net.minecraft.world.item.Tier
import net.minecraft.world.item.Tiers
import net.minecraft.world.level.block.Block
import net.minecraftforge.data.event.GatherDataEvent
import net.minecraftforge.registries.ForgeRegistries
import net.minecraftforge.registries.IForgeRegistry
import ru.dbotthepony.mc.otm.capability.MatteryPlayerCapability
import ru.dbotthepony.mc.otm.datagen.DataGen
import java.util.stream.Stream
import net.minecraft.data.tags.TagsProvider as MinecraftTagsProvider

private fun <T : Any> vanillaLookup(key: ResourceKey<Registry<T>>): (T) -> ResourceLocation {
	val registry by lazy { (BuiltInRegistries.REGISTRY.get(key.location()) ?: throw NoSuchElementException("No such registry $key")) as Registry<T> }
	return { registry.getKey(it) ?: throw NoSuchElementException("Registry $key does not contain $it") }
}

class TagsProvider(private val event: GatherDataEvent) {
	inner class Delegate<T : Any> private constructor(key: ResourceKey<Registry<T>>, val lookup: (T) -> ResourceLocation) : MinecraftTagsProvider<T>(event.generator.packOutput, key, event.lookupProvider, DataGen.MOD_ID, event.existingFileHelper) {
		constructor(registry: IForgeRegistry<T>) : this(registry.registryKey, { registry.getKey(it) ?: throw NoSuchElementException("Registry $registry does not contain $it") })
		constructor(key: ResourceKey<Registry<T>>) : this(key, vanillaLookup(key))

		init {
			if (isRegistered)
				event.generator.addProvider(event.includeServer(), this)
			else
				delegates.add(this)
		}

		private val tags = HashMap<TagKey<T>, ObjectArraySet<ResourceLocation>>()
		private val rigidTags = HashMap<TagKey<T>, ObjectArraySet<ResourceKey<T>>>()
		private val tagInTag = HashMap<TagKey<T>, ObjectArraySet<TagKey<T>>>()

		inner class Appender(val tag: TagKey<T>) {
			constructor(tag: ResourceLocation) : this(TagKey.create(registryKey, tag))

			init {
				require(tag.registry == registryKey) { "Trying to create appender for $tag inside registry $registryKey" }
			}

			private val locations by lazy { tags.computeIfAbsent(tag) { ObjectArraySet() } }
			private val rigidLocations by lazy { rigidTags.computeIfAbsent(tag) { ObjectArraySet() } }
			private val tagsInTags by lazy { tagInTag.computeIfAbsent(tag) { ObjectArraySet() } }

			fun add(value: ResourceLocation): Appender {
				check(locations.add(value)) { "Tag ${tag.location} of registry ${registryKey.location()} already contains $value" }
				return this
			}

			fun add(value: ResourceKey<T>): Appender {
				require(value.registry() == registryKey.location()) { "Invalid registry in provided ResourceKey: ${value.registry()} (this tag appender is for ${registryKey.location()})" }
				// check(rigidLocations.add(value)) { "Tag ${tag.location} of registry ${registryKey.location()} already contains $value" }
				// return this
				return add(value.location())
			}

			fun add(value: TagKey<T>): Appender {
				require(value.registry() == registryKey) { "Invalid registry in provided ResourceKey: ${value.registry().location()} (this tag appender is for ${registryKey.location()})" }
				check(tagsInTags.add(value)) { "Tag ${tag.location} of registry ${registryKey.location()} already contains $value" }
				return this
			}

			fun add(value: T): Appender {
				return add(ResourceKey.create(registryKey, lookup.invoke(value)))
			}

			fun add(values: Collection<T>): Appender {
				for (value in values) add(value)
				return this
			}

			fun add(vararg values: T): Appender {
				for (value in values) add(value)
				return this
			}

			fun add(vararg values: ResourceKey<T>): Appender {
				for (value in values) add(value)
				return this
			}

			fun add(vararg values: TagKey<T>): Appender {
				for (value in values) add(value)
				return this
			}

			fun add(vararg values: ResourceLocation): Appender {
				for (value in values) add(value)
				return this
			}

			fun add(values: Stream<T>): Appender {
				values.forEach { add(it) }
				return this
			}

			fun leaf(name: String) = Appender(TagKey.create(tag.registry, ResourceLocation(tag.location.namespace, tag.location.path + "/$name")))

			fun add(leaf: String, value: T) = also { leaf(leaf).add(value) }
			fun add(leaf: String, value: TagKey<T>) = also { leaf(leaf).add(value) }
			fun add(leaf: String, vararg value: T) = also { leaf(leaf).add(*value) }
			fun add(leaf: String, vararg value: TagKey<T>) = also { leaf(leaf).add(*value) }
			fun add(leaf: String, value: Stream<T>) = also { leaf(leaf).add(value) }
			fun add(leaf: String, value: ResourceKey<T>) = also { leaf(leaf).add(value) }
			fun add(leaf: String, vararg value: ResourceKey<T>) = also { leaf(leaf).add(*value) }
			fun add(leaf: String, value: ResourceLocation) = also { leaf(leaf).add(value) }
			fun add(leaf: String, vararg value: ResourceLocation) = also { leaf(leaf).add(*value) }
		}

		fun forge(path: String) = Appender(ResourceLocation("forge", path))
		fun minecraft(path: String) = Appender(ResourceLocation("minecraft", path))

		override fun addTags(provider: HolderLookup.Provider) {
			for ((tag, values) in tags) {
				val appender = tag(tag)

				for (value in values) {
					appender.addOptional(value)
				}
			}

			for ((tag, values) in rigidTags) {
				val appender = tag(tag)

				for (value in values) {
					appender.add(value)
				}
			}

			for ((tag, values) in tagInTag) {
				val appender = tag(tag)

				for (value in values) {
					appender.addTag(value)
				}
			}
		}
	}

	private val delegates = ArrayList<Delegate<*>>()
	private var isRegistered = false

	fun register() {
		if (!isRegistered) {
			isRegistered = true

			for (value in delegates) {
				event.generator.addProvider(event.includeServer(), value)
			}

			delegates.clear()
		}
	}

	val blocks = Delegate(ForgeRegistries.BLOCKS)
	val items = Delegate(ForgeRegistries.ITEMS)
	val mobEffects = Delegate(ForgeRegistries.MOB_EFFECTS)
	val damageTypes = Delegate(Registries.DAMAGE_TYPE)

	val androidImmuneEffects = mobEffects.Appender(MatteryPlayerCapability.ANDROID_IMMUNE_EFFECTS)

	val requiresShovel = blocks.Appender(BlockTags.MINEABLE_WITH_SHOVEL)
	val requiresAxe = blocks.Appender(BlockTags.MINEABLE_WITH_AXE)
	val requiresHoe = blocks.Appender(BlockTags.MINEABLE_WITH_HOE)
	val requiresPickaxe = blocks.Appender(BlockTags.MINEABLE_WITH_PICKAXE)

	val requiresStoneTool = blocks.Appender(BlockTags.NEEDS_STONE_TOOL)
	val requiresIronTool = blocks.Appender(BlockTags.NEEDS_IRON_TOOL)
	val requiresDiamondTool = blocks.Appender(BlockTags.NEEDS_DIAMOND_TOOL)

	val witherImmune = blocks.Appender(BlockTags.WITHER_IMMUNE)
	val dragonImmune = blocks.Appender(BlockTags.DRAGON_IMMUNE)

	fun stoneOre(key: String, block: Block): TagsProvider {
		ore(key, block)

		itemGroundOresStone.add(block.asItem())
		blockGroundOresStone.add(block)

		return this
	}

	fun deepslateOre(key: String, block: Block): TagsProvider {
		ore(key, block)

		itemGroundOresDeepslate.add(block.asItem())
		blockGroundOresDeepslate.add(block)

		return this
	}

	fun singleDropOre(vararg blocks: Block): TagsProvider {
		for (block in blocks) {
			itemOreRatesSingular.add(block.asItem())
			blockOreRatesSingular.add(block)
		}

		return this
	}

	fun ore(key: String, block: Block): TagsProvider {
		val forgeKey = ResourceLocation("forge", "ores/$key")
		val b = TagKey.create(Registries.BLOCK, forgeKey)
		val i = TagKey.create(Registries.ITEM, forgeKey)

		items.Appender(i).add(block.asItem())
		itemOres.add(block.asItem())
		blocks.Appender(b).add(block)
		blockOres.add(block)

		return this
	}

	fun clump(key: String, clump: Item): TagsProvider {
		rawMaterials.add(key, clump)
		rawMaterials.add(clump)
		return this
	}

	fun clump(key: String, clump: Item, rawBlock: Block): TagsProvider {
		val rawKey = "raw_$key"

		rawMaterials.add(key, clump)
		rawMaterials.add(clump)
		storageBlocksAsItem.add(rawKey, rawBlock.asItem())
		storageBlocksAsItem.add(rawBlock.asItem())
		storageBlocksAsBlock.add(rawKey, rawBlock)
		storageBlocksAsBlock.add(rawBlock)

		return this
	}

	val circuits = items.forge("circuits")
	val dusts = items.forge("dusts")
	val ingots = items.forge("ingots")
	val itemOres = items.forge("ores")
	val blockOres = blocks.forge("ores")
	val plates = items.forge("plates")
	val storageBlocksAsItem = items.forge("storage_blocks")
	val storageBlocksAsBlock = blocks.forge("storage_blocks")
	val rawMaterials = items.forge("raw_materials")
	val wires = items.forge("wires")

	val blockGroundOresStone = blocks.forge("ores_in_ground/stone")
	val blockGroundOresDeepslate = blocks.forge("ores_in_ground/deepslate")
//	val blockGroundOresNetherrack = blocks.forge("ores_in_ground/netherrack")
	val itemGroundOresStone = items.forge("ores_in_ground/stone")
	val itemGroundOresDeepslate = items.forge("ores_in_ground/deepslate")
//	val itemGroundOresNetherrack = items.forge("ores_in_ground/netherrack")

//	val blockOreRatesSparse = blocks.forge("ore_rates/sparse")
	val blockOreRatesSingular = blocks.forge("ore_rates/singular")
//	val blockOreRatesDense = blocks.forge("ore_rates/dense")
//	val itemOreRatesSparse = items.forge("ore_rates/sparse")
	val itemOreRatesSingular = items.forge("ore_rates/singular")
//	val itemOreRatesDense = items.forge("ore_rates/dense")

	val gameEvents = Delegate(Registries.GAME_EVENT)
	val vibrations = gameEvents.Appender(GameEventTags.VIBRATIONS)

	fun requiresPickaxe(block: Block, tier: Tier? = null): TagsProvider {
		requiresPickaxe.add(block)

		when (tier) {
			Tiers.STONE -> requiresStoneTool.add(block)
			Tiers.IRON -> requiresIronTool.add(block)
			Tiers.DIAMOND, Tiers.NETHERITE -> requiresDiamondTool.add(block)
		}

		return this
	}

	fun requiresPickaxe(blocks: Collection<Block>, tier: Tier? = null): TagsProvider {
		for (block in blocks) requiresPickaxe(block, tier)
		return this
	}

	fun requiresPickaxe(vararg blocks: Block, tier: Tier? = null): TagsProvider {
		for (block in blocks) requiresPickaxe(block, tier)
		return this
	}

	fun requiresShovel(block: Block, tier: Tier? = null): TagsProvider {
		requiresShovel.add(block)

		when (tier) {
			Tiers.STONE -> requiresStoneTool.add(block)
			Tiers.IRON -> requiresIronTool.add(block)
			Tiers.DIAMOND, Tiers.NETHERITE -> requiresDiamondTool.add(block)
		}

		return this
	}

	fun requiresShovel(blocks: Collection<Block>, tier: Tier? = null): TagsProvider {
		for (block in blocks) requiresShovel(block, tier)
		return this
	}

	fun requiresAxe(block: Block, tier: Tier? = null): TagsProvider {
		requiresAxe.add(block)

		when (tier) {
			Tiers.STONE -> requiresStoneTool.add(block)
			Tiers.IRON -> requiresIronTool.add(block)
			Tiers.DIAMOND, Tiers.NETHERITE -> requiresDiamondTool.add(block)
		}

		return this
	}

	fun requiresAxe(blocks: Collection<Block>, tier: Tier? = null): TagsProvider {
		for (block in blocks) requiresAxe(block, tier)
		return this
	}
}

