package net.minecraft.server.commands;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.mojang.brigadier.CommandDispatcher;
import com.mojang.brigadier.arguments.IntegerArgumentType;
import com.mojang.brigadier.exceptions.CommandSyntaxException;
import com.mojang.brigadier.exceptions.Dynamic2CommandExceptionType;
import com.mojang.brigadier.exceptions.Dynamic3CommandExceptionType;
import com.mojang.brigadier.exceptions.DynamicCommandExceptionType;
import com.mojang.brigadier.suggestion.SuggestionProvider;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Map.Entry;
import net.minecraft.commands.CommandBuildContext;
import net.minecraft.commands.CommandSourceStack;
import net.minecraft.commands.Commands;
import net.minecraft.commands.SharedSuggestionProvider;
import net.minecraft.commands.arguments.EntityArgument;
import net.minecraft.commands.arguments.ResourceOrIdArgument;
import net.minecraft.commands.arguments.SlotArgument;
import net.minecraft.commands.arguments.coordinates.BlockPosArgument;
import net.minecraft.commands.arguments.item.ItemArgument;
import net.minecraft.core.BlockPos;
import net.minecraft.core.Holder;
import net.minecraft.core.registries.Registries;
import net.minecraft.network.chat.Component;
import net.minecraft.server.ReloadableServerRegistries;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.Container;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.SlotAccess;
import net.minecraft.world.item.ItemStack;
import net.minecraft.world.level.block.entity.BlockEntity;
import net.minecraft.world.level.storage.loot.LootContext;
import net.minecraft.world.level.storage.loot.LootParams;
import net.minecraft.world.level.storage.loot.functions.LootItemFunction;
import net.minecraft.world.level.storage.loot.parameters.LootContextParamSets;
import net.minecraft.world.level.storage.loot.parameters.LootContextParams;

public class ItemCommands {
	static final Dynamic3CommandExceptionType ERROR_TARGET_NOT_A_CONTAINER = new Dynamic3CommandExceptionType(
		(object, object2, object3) -> Component.translatableEscape("commands.item.target.not_a_container", object, object2, object3)
	);
	static final Dynamic3CommandExceptionType ERROR_SOURCE_NOT_A_CONTAINER = new Dynamic3CommandExceptionType(
		(object, object2, object3) -> Component.translatableEscape("commands.item.source.not_a_container", object, object2, object3)
	);
	static final DynamicCommandExceptionType ERROR_TARGET_INAPPLICABLE_SLOT = new DynamicCommandExceptionType(
		object -> Component.translatableEscape("commands.item.target.no_such_slot", object)
	);
	private static final DynamicCommandExceptionType ERROR_SOURCE_INAPPLICABLE_SLOT = new DynamicCommandExceptionType(
		object -> Component.translatableEscape("commands.item.source.no_such_slot", object)
	);
	private static final DynamicCommandExceptionType ERROR_TARGET_NO_CHANGES = new DynamicCommandExceptionType(
		object -> Component.translatableEscape("commands.item.target.no_changes", object)
	);
	private static final Dynamic2CommandExceptionType ERROR_TARGET_NO_CHANGES_KNOWN_ITEM = new Dynamic2CommandExceptionType(
		(object, object2) -> Component.translatableEscape("commands.item.target.no_changed.known_item", object, object2)
	);
	private static final SuggestionProvider<CommandSourceStack> SUGGEST_MODIFIER = (commandContext, suggestionsBuilder) -> {
		ReloadableServerRegistries.Holder holder = commandContext.getSource().getServer().reloadableRegistries();
		return SharedSuggestionProvider.suggestResource(holder.getKeys(Registries.ITEM_MODIFIER), suggestionsBuilder);
	};

	public static void register(CommandDispatcher<CommandSourceStack> dispatcher, CommandBuildContext context) {
		dispatcher.register(
			Commands.literal("item")
				.requires(commandSourceStack -> commandSourceStack.hasPermission(2))
				.then(
					Commands.literal("replace")
						.then(
							Commands.literal("block")
								.then(
									Commands.argument("pos", BlockPosArgument.blockPos())
										.then(
											Commands.argument("slot", SlotArgument.slot())
												.then(
													Commands.literal("with")
														.then(
															Commands.argument("item", ItemArgument.item(context))
																.executes(
																	commandContext -> setBlockItem(
																		commandContext.getSource(),
																		BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																		SlotArgument.getSlot(commandContext, "slot"),
																		ItemArgument.getItem(commandContext, "item").createItemStack(1, false)
																	)
																)
																.then(
																	Commands.argument("count", IntegerArgumentType.integer(1, 99))
																		.executes(
																			commandContext -> setBlockItem(
																				commandContext.getSource(),
																				BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																				SlotArgument.getSlot(commandContext, "slot"),
																				ItemArgument.getItem(commandContext, "item").createItemStack(IntegerArgumentType.getInteger(commandContext, "count"), true)
																			)
																		)
																)
														)
												)
												.then(
													Commands.literal("from")
														.then(
															Commands.literal("block")
																.then(
																	Commands.argument("source", BlockPosArgument.blockPos())
																		.then(
																			Commands.argument("sourceSlot", SlotArgument.slot())
																				.executes(
																					commandContext -> blockToBlock(
																						commandContext.getSource(),
																						BlockPosArgument.getLoadedBlockPos(commandContext, "source"),
																						SlotArgument.getSlot(commandContext, "sourceSlot"),
																						BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																						SlotArgument.getSlot(commandContext, "slot")
																					)
																				)
																				.then(
																					Commands.argument("modifier", ResourceOrIdArgument.lootModifier(context))
																						.suggests(SUGGEST_MODIFIER)
																						.executes(
																							commandContext -> blockToBlock(
																								(CommandSourceStack)commandContext.getSource(),
																								BlockPosArgument.getLoadedBlockPos(commandContext, "source"),
																								SlotArgument.getSlot(commandContext, "sourceSlot"),
																								BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																								SlotArgument.getSlot(commandContext, "slot"),
																								ResourceOrIdArgument.getLootModifier(commandContext, "modifier")
																							)
																						)
																				)
																		)
																)
														)
														.then(
															Commands.literal("entity")
																.then(
																	Commands.argument("source", EntityArgument.entity())
																		.then(
																			Commands.argument("sourceSlot", SlotArgument.slot())
																				.executes(
																					commandContext -> entityToBlock(
																						commandContext.getSource(),
																						EntityArgument.getEntity(commandContext, "source"),
																						SlotArgument.getSlot(commandContext, "sourceSlot"),
																						BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																						SlotArgument.getSlot(commandContext, "slot")
																					)
																				)
																				.then(
																					Commands.argument("modifier", ResourceOrIdArgument.lootModifier(context))
																						.suggests(SUGGEST_MODIFIER)
																						.executes(
																							commandContext -> entityToBlock(
																								(CommandSourceStack)commandContext.getSource(),
																								EntityArgument.getEntity(commandContext, "source"),
																								SlotArgument.getSlot(commandContext, "sourceSlot"),
																								BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																								SlotArgument.getSlot(commandContext, "slot"),
																								ResourceOrIdArgument.getLootModifier(commandContext, "modifier")
																							)
																						)
																				)
																		)
																)
														)
												)
										)
								)
						)
						.then(
							Commands.literal("entity")
								.then(
									Commands.argument("targets", EntityArgument.entities())
										.then(
											Commands.argument("slot", SlotArgument.slot())
												.then(
													Commands.literal("with")
														.then(
															Commands.argument("item", ItemArgument.item(context))
																.executes(
																	commandContext -> setEntityItem(
																		commandContext.getSource(),
																		EntityArgument.getEntities(commandContext, "targets"),
																		SlotArgument.getSlot(commandContext, "slot"),
																		ItemArgument.getItem(commandContext, "item").createItemStack(1, false)
																	)
																)
																.then(
																	Commands.argument("count", IntegerArgumentType.integer(1, 99))
																		.executes(
																			commandContext -> setEntityItem(
																				commandContext.getSource(),
																				EntityArgument.getEntities(commandContext, "targets"),
																				SlotArgument.getSlot(commandContext, "slot"),
																				ItemArgument.getItem(commandContext, "item").createItemStack(IntegerArgumentType.getInteger(commandContext, "count"), true)
																			)
																		)
																)
														)
												)
												.then(
													Commands.literal("from")
														.then(
															Commands.literal("block")
																.then(
																	Commands.argument("source", BlockPosArgument.blockPos())
																		.then(
																			Commands.argument("sourceSlot", SlotArgument.slot())
																				.executes(
																					commandContext -> blockToEntities(
																						commandContext.getSource(),
																						BlockPosArgument.getLoadedBlockPos(commandContext, "source"),
																						SlotArgument.getSlot(commandContext, "sourceSlot"),
																						EntityArgument.getEntities(commandContext, "targets"),
																						SlotArgument.getSlot(commandContext, "slot")
																					)
																				)
																				.then(
																					Commands.argument("modifier", ResourceOrIdArgument.lootModifier(context))
																						.suggests(SUGGEST_MODIFIER)
																						.executes(
																							commandContext -> blockToEntities(
																								(CommandSourceStack)commandContext.getSource(),
																								BlockPosArgument.getLoadedBlockPos(commandContext, "source"),
																								SlotArgument.getSlot(commandContext, "sourceSlot"),
																								EntityArgument.getEntities(commandContext, "targets"),
																								SlotArgument.getSlot(commandContext, "slot"),
																								ResourceOrIdArgument.getLootModifier(commandContext, "modifier")
																							)
																						)
																				)
																		)
																)
														)
														.then(
															Commands.literal("entity")
																.then(
																	Commands.argument("source", EntityArgument.entity())
																		.then(
																			Commands.argument("sourceSlot", SlotArgument.slot())
																				.executes(
																					commandContext -> entityToEntities(
																						commandContext.getSource(),
																						EntityArgument.getEntity(commandContext, "source"),
																						SlotArgument.getSlot(commandContext, "sourceSlot"),
																						EntityArgument.getEntities(commandContext, "targets"),
																						SlotArgument.getSlot(commandContext, "slot")
																					)
																				)
																				.then(
																					Commands.argument("modifier", ResourceOrIdArgument.lootModifier(context))
																						.suggests(SUGGEST_MODIFIER)
																						.executes(
																							commandContext -> entityToEntities(
																								(CommandSourceStack)commandContext.getSource(),
																								EntityArgument.getEntity(commandContext, "source"),
																								SlotArgument.getSlot(commandContext, "sourceSlot"),
																								EntityArgument.getEntities(commandContext, "targets"),
																								SlotArgument.getSlot(commandContext, "slot"),
																								ResourceOrIdArgument.getLootModifier(commandContext, "modifier")
																							)
																						)
																				)
																		)
																)
														)
												)
										)
								)
						)
				)
				.then(
					Commands.literal("modify")
						.then(
							Commands.literal("block")
								.then(
									Commands.argument("pos", BlockPosArgument.blockPos())
										.then(
											Commands.argument("slot", SlotArgument.slot())
												.then(
													Commands.argument("modifier", ResourceOrIdArgument.lootModifier(context))
														.suggests(SUGGEST_MODIFIER)
														.executes(
															commandContext -> modifyBlockItem(
																(CommandSourceStack)commandContext.getSource(),
																BlockPosArgument.getLoadedBlockPos(commandContext, "pos"),
																SlotArgument.getSlot(commandContext, "slot"),
																ResourceOrIdArgument.getLootModifier(commandContext, "modifier")
															)
														)
												)
										)
								)
						)
						.then(
							Commands.literal("entity")
								.then(
									Commands.argument("targets", EntityArgument.entities())
										.then(
											Commands.argument("slot", SlotArgument.slot())
												.then(
													Commands.argument("modifier", ResourceOrIdArgument.lootModifier(context))
														.suggests(SUGGEST_MODIFIER)
														.executes(
															commandContext -> modifyEntityItem(
																(CommandSourceStack)commandContext.getSource(),
																EntityArgument.getEntities(commandContext, "targets"),
																SlotArgument.getSlot(commandContext, "slot"),
																ResourceOrIdArgument.getLootModifier(commandContext, "modifier")
															)
														)
												)
										)
								)
						)
				)
		);
	}

	private static int modifyBlockItem(CommandSourceStack source, BlockPos pos, int slot, Holder<LootItemFunction> modifier) throws CommandSyntaxException {
		Container container = getContainer(source, pos, ERROR_TARGET_NOT_A_CONTAINER);
		if (slot >= 0 && slot < container.getContainerSize()) {
			ItemStack itemStack = applyModifier(source, modifier, container.getItem(slot));
			container.setItem(slot, itemStack);
			source.sendSuccess(() -> Component.translatable("commands.item.block.set.success", pos.getX(), pos.getY(), pos.getZ(), itemStack.getDisplayName()), true);
			return 1;
		} else {
			throw ERROR_TARGET_INAPPLICABLE_SLOT.create(slot);
		}
	}

	private static int modifyEntityItem(CommandSourceStack source, Collection<? extends Entity> targets, int sourceSlot, Holder<LootItemFunction> modifer) throws CommandSyntaxException {
		Map<Entity, ItemStack> map = Maps.<Entity, ItemStack>newHashMapWithExpectedSize(targets.size());

		for (Entity entity : targets) {
			SlotAccess slotAccess = entity.getSlot(sourceSlot);
			if (slotAccess != SlotAccess.NULL) {
				ItemStack itemStack = applyModifier(source, modifer, slotAccess.get().copy());
				if (slotAccess.set(itemStack)) {
					map.put(entity, itemStack);
					if (entity instanceof ServerPlayer) {
						((ServerPlayer)entity).containerMenu.broadcastChanges();
					}
				}
			}
		}

		if (map.isEmpty()) {
			throw ERROR_TARGET_NO_CHANGES.create(sourceSlot);
		} else {
			if (map.size() == 1) {
				Entry<Entity, ItemStack> entry = (Entry<Entity, ItemStack>)map.entrySet().iterator().next();
				source.sendSuccess(
					() -> Component.translatable(
						"commands.item.entity.set.success.single", ((Entity)entry.getKey()).getDisplayName(), ((ItemStack)entry.getValue()).getDisplayName()
					),
					true
				);
			} else {
				source.sendSuccess(() -> Component.translatable("commands.item.entity.set.success.multiple", map.size()), true);
			}

			return map.size();
		}
	}

	private static int setBlockItem(CommandSourceStack source, BlockPos pos, int slot, ItemStack item) throws CommandSyntaxException {
		Container container = getContainer(source, pos, ERROR_TARGET_NOT_A_CONTAINER);
		if (slot >= 0 && slot < container.getContainerSize()) {
			container.setItem(slot, item);
			source.sendSuccess(() -> Component.translatable("commands.item.block.set.success", pos.getX(), pos.getY(), pos.getZ(), item.getDisplayName()), true);
			return 1;
		} else {
			throw ERROR_TARGET_INAPPLICABLE_SLOT.create(slot);
		}
	}

	static Container getContainer(CommandSourceStack source, BlockPos pos, Dynamic3CommandExceptionType exception) throws CommandSyntaxException {
		BlockEntity blockEntity = source.getLevel().getBlockEntity(pos);
		if (!(blockEntity instanceof Container)) {
			throw exception.create(pos.getX(), pos.getY(), pos.getZ());
		} else {
			return (Container)blockEntity;
		}
	}

	private static int setEntityItem(CommandSourceStack source, Collection<? extends Entity> targets, int slot, ItemStack item) throws CommandSyntaxException {
		List<Entity> list = Lists.<Entity>newArrayListWithCapacity(targets.size());

		for (Entity entity : targets) {
			SlotAccess slotAccess = entity.getSlot(slot);
			if (slotAccess != SlotAccess.NULL && slotAccess.set(item.copy())) {
				list.add(entity);
				if (entity instanceof ServerPlayer) {
					((ServerPlayer)entity).containerMenu.broadcastChanges();
				}
			}
		}

		if (list.isEmpty()) {
			throw ERROR_TARGET_NO_CHANGES_KNOWN_ITEM.create(item.getDisplayName(), slot);
		} else {
			if (list.size() == 1) {
				source.sendSuccess(
					() -> Component.translatable("commands.item.entity.set.success.single", ((Entity)list.iterator().next()).getDisplayName(), item.getDisplayName()), true
				);
			} else {
				source.sendSuccess(() -> Component.translatable("commands.item.entity.set.success.multiple", list.size(), item.getDisplayName()), true);
			}

			return list.size();
		}
	}

	private static int blockToEntities(CommandSourceStack source, BlockPos pos, int sourceSlot, Collection<? extends Entity> targets, int slot) throws CommandSyntaxException {
		return setEntityItem(source, targets, slot, getBlockItem(source, pos, sourceSlot));
	}

	private static int blockToEntities(
		CommandSourceStack source, BlockPos pos, int sourceSlot, Collection<? extends Entity> targets, int slot, Holder<LootItemFunction> modifier
	) throws CommandSyntaxException {
		return setEntityItem(source, targets, slot, applyModifier(source, modifier, getBlockItem(source, pos, sourceSlot)));
	}

	private static int blockToBlock(CommandSourceStack source, BlockPos sourcePos, int sourceSlot, BlockPos pos, int slot) throws CommandSyntaxException {
		return setBlockItem(source, pos, slot, getBlockItem(source, sourcePos, sourceSlot));
	}

	private static int blockToBlock(CommandSourceStack source, BlockPos sourcePos, int sourceSlot, BlockPos pos, int slot, Holder<LootItemFunction> modifier) throws CommandSyntaxException {
		return setBlockItem(source, pos, slot, applyModifier(source, modifier, getBlockItem(source, sourcePos, sourceSlot)));
	}

	private static int entityToBlock(CommandSourceStack source, Entity sourceEntity, int sourceSlot, BlockPos pos, int slot) throws CommandSyntaxException {
		return setBlockItem(source, pos, slot, getEntityItem(sourceEntity, sourceSlot));
	}

	private static int entityToBlock(CommandSourceStack source, Entity sourceEntity, int sourceSlot, BlockPos pos, int slot, Holder<LootItemFunction> modifier) throws CommandSyntaxException {
		return setBlockItem(source, pos, slot, applyModifier(source, modifier, getEntityItem(sourceEntity, sourceSlot)));
	}

	private static int entityToEntities(CommandSourceStack source, Entity sourceEntity, int sourceSlot, Collection<? extends Entity> targets, int slot) throws CommandSyntaxException {
		return setEntityItem(source, targets, slot, getEntityItem(sourceEntity, sourceSlot));
	}

	private static int entityToEntities(
		CommandSourceStack source, Entity sourceEntity, int sourceSlot, Collection<? extends Entity> targets, int slot, Holder<LootItemFunction> modifier
	) throws CommandSyntaxException {
		return setEntityItem(source, targets, slot, applyModifier(source, modifier, getEntityItem(sourceEntity, sourceSlot)));
	}

	private static ItemStack applyModifier(CommandSourceStack source, Holder<LootItemFunction> modifier, ItemStack stack) {
		ServerLevel serverLevel = source.getLevel();
		LootParams lootParams = new LootParams.Builder(serverLevel)
			.withParameter(LootContextParams.ORIGIN, source.getPosition())
			.withOptionalParameter(LootContextParams.THIS_ENTITY, source.getEntity())
			.create(LootContextParamSets.COMMAND);
		LootContext lootContext = new LootContext.Builder(lootParams).create(Optional.empty());
		lootContext.pushVisitedElement(LootContext.createVisitedEntry(modifier.value()));
		ItemStack itemStack = (ItemStack)modifier.value().apply(stack, lootContext);
		itemStack.limitSize(itemStack.getMaxStackSize());
		return itemStack;
	}

	private static ItemStack getEntityItem(Entity entity, int slot) throws CommandSyntaxException {
		SlotAccess slotAccess = entity.getSlot(slot);
		if (slotAccess == SlotAccess.NULL) {
			throw ERROR_SOURCE_INAPPLICABLE_SLOT.create(slot);
		} else {
			return slotAccess.get().copy();
		}
	}

	private static ItemStack getBlockItem(CommandSourceStack source, BlockPos pos, int slot) throws CommandSyntaxException {
		Container container = getContainer(source, pos, ERROR_SOURCE_NOT_A_CONTAINER);
		if (slot >= 0 && slot < container.getContainerSize()) {
			return container.getItem(slot).copy();
		} else {
			throw ERROR_SOURCE_INAPPLICABLE_SLOT.create(slot);
		}
	}
}