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

import com.google.common.collect.ImmutableList
import com.mojang.serialization.Codec
import com.mojang.serialization.MapCodec
import com.mojang.serialization.codecs.RecordCodecBuilder
import net.minecraft.core.component.DataComponents
import net.minecraft.nbt.CompoundTag
import net.minecraft.network.chat.ComponentSerialization
import net.minecraft.world.item.ItemStack
import net.minecraft.world.item.component.CustomData
import net.minecraft.world.level.block.entity.BlockEntity
import net.minecraft.world.level.storage.loot.LootContext
import net.minecraft.world.level.storage.loot.functions.LootItemFunction
import net.minecraft.world.level.storage.loot.functions.LootItemFunctionType
import net.minecraft.world.level.storage.loot.parameters.LootContextParams
import ru.dbotthepony.mc.otm.block.entity.MatteryBlockEntity
import ru.dbotthepony.mc.otm.core.fromNbt
import ru.dbotthepony.mc.otm.core.nbt.set
import ru.dbotthepony.mc.otm.core.stream
import ru.dbotthepony.mc.otm.registry.data.MItemFunctionTypes
import java.util.*
import java.util.stream.Stream

class CopyTileNbtFunction(filter: Stream<out String> = Stream.empty()) : LootItemFunction, LootItemFunction.Builder {
	constructor(filter: Collection<String>) : this(filter.stream())
	constructor(vararg filter: String) : this(filter.stream())

	val filter: ImmutableList<String> = /*immutableList {
		filter.forEach(this)
	}*/ filter.collect(ImmutableList.toImmutableList())

	override fun apply(itemStack: ItemStack, context: LootContext): ItemStack {
		val blockEntity = context.getParamOrNull(LootContextParams.BLOCK_ENTITY) ?: return itemStack
		val result = itemStack[DataComponents.BLOCK_ENTITY_DATA] ?: CustomData.EMPTY
		val data: CompoundTag

		if (blockEntity is MatteryBlockEntity) {
			data = CompoundTag()
			// TODO: 1.21 added components, which are present even on blockentities; these are lost
			blockEntity.saveShared(data, context.level.registryAccess())
		} else {
			data = blockEntity.saveWithoutMetadata(context.level.registryAccess())
		}

		for (k in filter) {
			data.remove(k)
		}

		BlockEntity.addEntityType(data, blockEntity.type)

		if (result.isEmpty) {
			itemStack[DataComponents.BLOCK_ENTITY_DATA] = CustomData.of(data)
		} else {
			val copy = result.copyTag()

			for (k in data.allKeys) {
				copy[k] = data[k]!!
			}

			itemStack[DataComponents.BLOCK_ENTITY_DATA] = CustomData.of(copy)
		}

		if ("Name" in data && !itemStack.has(DataComponents.CUSTOM_NAME)) {
			val deserialize = ComponentSerialization.CODEC.fromNbt(data["Name"]!!)

			if (deserialize != null) {
				itemStack[DataComponents.CUSTOM_NAME] = deserialize
			}

			data.remove("Name")
		}

		return itemStack
	}

	override fun getType(): LootItemFunctionType<CopyTileNbtFunction> {
		return MItemFunctionTypes.COPY_TILE_NBT
	}

	override fun build(): LootItemFunction {
		return this
	}

	companion object {
		val CODEC: MapCodec<CopyTileNbtFunction> by lazy {
			RecordCodecBuilder.mapCodec {
				it.group(
					Codec.STRING.listOf()
						.optionalFieldOf("filter")
						.xmap({ it.orElseGet { listOf() } }, { if (it.isEmpty()) Optional.empty() else Optional.of(it) })
						.forGetter(CopyTileNbtFunction::filter),
				).apply(it, ::CopyTileNbtFunction)
			}
		}
	}
}
