package net.mcreator.customoregen.loot;

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.item.Items;
import net.minecraft.world.level.block.Block;
import net.minecraft.world.level.block.Blocks;
import net.minecraft.world.level.block.state.BlockState;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.level.storage.loot.parameters.LootContextParams;
import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
import net.neoforged.neoforge.common.loot.IGlobalLootModifier;
import net.neoforged.neoforge.common.loot.LootModifier;
import net.minecraft.core.registries.Registries;
import net.minecraft.world.item.enchantment.EnchantmentHelper;
import net.minecraft.world.item.enchantment.Enchantments;

import net.mcreator.customoregen.init.CustomOreGenModBlocks;
import net.mcreator.customoregen.init.CustomOreGenModItems;
import net.mcreator.customoregen.config.ModConfigs;
import net.mcreator.customoregen.CustomOreGenMod;

import javax.annotation.Nonnull;

public class CustomOreLootModifier extends LootModifier {
	public static final MapCodec<CustomOreLootModifier> CODEC = RecordCodecBuilder.mapCodec(inst -> 
		codecStart(inst).apply(inst, CustomOreLootModifier::new)
	);

	public CustomOreLootModifier(LootItemCondition[] conditionsIn) {
		super(conditionsIn);
	}

	@Override
	public MapCodec<? extends IGlobalLootModifier> codec() {
		return CODEC;
	}

	@Nonnull
	@Override
	protected ObjectArrayList<ItemStack> doApply(ObjectArrayList<ItemStack> generatedLoot, LootContext context) {
		try {
			BlockState state = context.getParamOrNull(LootContextParams.BLOCK_STATE);
			if (state == null) return generatedLoot;

			Block block = state.getBlock();
			String oreType = getOreType(block);

			if (oreType != null) {
				ItemStack tool = context.getParamOrNull(LootContextParams.TOOL);
				
				// Check for Silk Touch
				if (tool != null && !tool.isEmpty()) {
					var registry = context.getLevel().registryAccess().lookupOrThrow(Registries.ENCHANTMENT);
					int silkLevel = tool.getEnchantmentLevel(registry.getOrThrow(Enchantments.SILK_TOUCH));
					
					if (silkLevel > 0) {
						// For Diamond Shard, we want to keep the modded block itself
						if (oreType.equals("shard_diamond")) {
							ObjectArrayList<ItemStack> shardLoot = new ObjectArrayList<>();
							shardLoot.add(new ItemStack(block));
							return shardLoot;
						}
						// For others, convert to vanilla block equivalent
						ItemStack vanillaBlock = getVanillaBlockDrop(block);
						if (!vanillaBlock.isEmpty()) {
							ObjectArrayList<ItemStack> silkLoot = new ObjectArrayList<>();
							silkLoot.add(vanillaBlock);
							return silkLoot;
						}
						return generatedLoot;
					}
				}

				// If not silk touch, replace with custom drops
				ObjectArrayList<ItemStack> customDrops = new ObjectArrayList<>();
				int fortuneLevel = 0;
				if (tool != null && !tool.isEmpty()) {
					var registry = context.getLevel().registryAccess().lookupOrThrow(Registries.ENCHANTMENT);
					fortuneLevel = tool.getEnchantmentLevel(registry.getOrThrow(Enchantments.FORTUNE));
				}

				addCustomDrops(customDrops, oreType, fortuneLevel, context.getRandom());
				
				// Only return custom drops if we generated some, otherwise fallback to original
				if (!customDrops.isEmpty()) {
					return customDrops;
				}
			}
		} catch (Exception e) {
			CustomOreGenMod.LOGGER.error("Error applying CustomOreLootModifier", e);
		}
		
		return generatedLoot;
	}

	private String getOreType(Block block) {
		if (block == CustomOreGenModBlocks.SHARDDIAMONDBLOCKORE.get() || block == CustomOreGenModBlocks.DEEPSLATESHARDDIAMONDORE.get()) return "shard_diamond";
		if (block == CustomOreGenModBlocks.CONCENTRATEDCOALORE.get()) return "concentrated_coal";
		if (block == CustomOreGenModBlocks.PUREGOLDENORE.get() || block == CustomOreGenModBlocks.DEEPSLATEPUREGOLDENORE.get()) return "pure_golden";
		if (block == CustomOreGenModBlocks.IRONORE.get() || block == CustomOreGenModBlocks.DEEPSLATEIRONORE.get()) return "impure_iron";
		if (block == CustomOreGenModBlocks.DEEPSLATEDIAMONDORE.get()) return "concentrated_diamond";
		if (block == CustomOreGenModBlocks.LAPISORE.get() || block == CustomOreGenModBlocks.DEEPSLATELAPISORE.get()) return "lapis";
		if (block == CustomOreGenModBlocks.REDSTONEORE.get() || block == CustomOreGenModBlocks.DEEPSLATEREDSTONEORE.get()) return "redstone";
		if (block == CustomOreGenModBlocks.HIGHEMERALDORE.get() || block == CustomOreGenModBlocks.LOWEREMERALDORE.get()) return "emerald";
		if (block == CustomOreGenModBlocks.COPPERHIGHORE.get() || block == CustomOreGenModBlocks.COPPERLOWERORE.get()) return "copper";
		return null;
	}

	private ItemStack getVanillaBlockDrop(Block block) {
		if (block == CustomOreGenModBlocks.CONCENTRATEDCOALORE.get()) return new ItemStack(Blocks.COAL_ORE);
		if (block == CustomOreGenModBlocks.IRONORE.get()) return new ItemStack(Blocks.IRON_ORE);
		if (block == CustomOreGenModBlocks.DEEPSLATEIRONORE.get()) return new ItemStack(Blocks.DEEPSLATE_IRON_ORE);
		if (block == CustomOreGenModBlocks.PUREGOLDENORE.get()) return new ItemStack(Blocks.GOLD_ORE);
		if (block == CustomOreGenModBlocks.DEEPSLATEPUREGOLDENORE.get()) return new ItemStack(Blocks.DEEPSLATE_GOLD_ORE);
		if (block == CustomOreGenModBlocks.LAPISORE.get()) return new ItemStack(Blocks.LAPIS_ORE);
		if (block == CustomOreGenModBlocks.DEEPSLATELAPISORE.get()) return new ItemStack(Blocks.DEEPSLATE_LAPIS_ORE);
		if (block == CustomOreGenModBlocks.REDSTONEORE.get()) return new ItemStack(Blocks.REDSTONE_ORE);
		if (block == CustomOreGenModBlocks.DEEPSLATEREDSTONEORE.get()) return new ItemStack(Blocks.DEEPSLATE_REDSTONE_ORE);
		if (block == CustomOreGenModBlocks.HIGHEMERALDORE.get()) return new ItemStack(Blocks.EMERALD_ORE);
		if (block == CustomOreGenModBlocks.LOWEREMERALDORE.get()) return new ItemStack(Blocks.DEEPSLATE_EMERALD_ORE);
		if (block == CustomOreGenModBlocks.COPPERHIGHORE.get()) return new ItemStack(Blocks.COPPER_ORE);
		if (block == CustomOreGenModBlocks.COPPERLOWERORE.get()) return new ItemStack(Blocks.DEEPSLATE_COPPER_ORE);
		if (block == CustomOreGenModBlocks.DEEPSLATEDIAMONDORE.get()) return new ItemStack(Blocks.DEEPSLATE_DIAMOND_ORE);
		return ItemStack.EMPTY;
	}

	private void addCustomDrops(ObjectArrayList<ItemStack> drops, String oreType, int fortuneLevel, net.minecraft.util.RandomSource random) {
		int minDrops = 1;
		int maxDrops = 1;
		boolean enableFortune = true;
		ItemStack dropItem = ItemStack.EMPTY;

		switch (oreType) {
			case "shard_diamond":
				minDrops = ModConfigs.DROPS.shardDiamondOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.shardDiamondOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.shardDiamondOreEnableFortune.get();
				dropItem = new ItemStack(CustomOreGenModItems.DIAMONDSHARD.get());
				break;
			case "concentrated_diamond":
				minDrops = ModConfigs.DROPS.concentratedDiamondOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.concentratedDiamondOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.concentratedDiamondOreEnableFortune.get();
				dropItem = new ItemStack(Items.DIAMOND);
				break;
			case "concentrated_coal":
				minDrops = ModConfigs.DROPS.concentratedCoalOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.concentratedCoalOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.concentratedCoalOreEnableFortune.get();
				dropItem = new ItemStack(Items.COAL);
				break;
			case "pure_golden":
				minDrops = ModConfigs.DROPS.pureGoldenOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.pureGoldenOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.pureGoldenOreEnableFortune.get();
				dropItem = new ItemStack(Items.RAW_GOLD);
				break;
			case "impure_iron":
				minDrops = ModConfigs.DROPS.impureIronOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.impureIronOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.impureIronOreEnableFortune.get();
				dropItem = new ItemStack(Items.RAW_IRON);
				break;
			case "lapis":
				minDrops = ModConfigs.DROPS.lapisOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.lapisOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.lapisOreEnableFortune.get();
				dropItem = new ItemStack(Items.LAPIS_LAZULI);
				break;
			case "redstone":
				minDrops = ModConfigs.DROPS.redstoneOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.redstoneOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.redstoneOreEnableFortune.get();
				dropItem = new ItemStack(Items.REDSTONE);
				break;
			case "emerald":
				minDrops = ModConfigs.DROPS.emeraldOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.emeraldOreMaxDrops.get();
				dropItem = new ItemStack(Items.EMERALD);
				break;
			case "copper":
				minDrops = ModConfigs.DROPS.copperOreMinDrops.get();
				maxDrops = ModConfigs.DROPS.copperOreMaxDrops.get();
				enableFortune = ModConfigs.DROPS.copperOreEnableFortune.get();
				dropItem = new ItemStack(Items.RAW_COPPER);
				break;
		}

		if (dropItem.isEmpty()) return;

		int dropCount = minDrops + (maxDrops > minDrops ? random.nextInt(maxDrops - minDrops + 1) : 0);

		if (enableFortune && fortuneLevel > 0) {
			if (oreType.equals("lapis") || oreType.equals("copper") || oreType.equals("redstone")) {
				int multiplier = random.nextInt(fortuneLevel + 2) - 1;
				if (multiplier < 0) multiplier = 0;
				dropCount *= (multiplier + 1);
			} else {
				int fortuneBonus = random.nextInt(fortuneLevel + 2) - 1;
				if (fortuneBonus < 0) fortuneBonus = 0;
				dropCount += fortuneBonus;
			}
		}

		if (dropCount > 0) {
			dropItem.setCount(dropCount);
			drops.add(dropItem);
		}
	}
}
