package net.minecraft.world.entity;

import com.mojang.serialization.Codec;
import java.util.EnumMap;
import java.util.Map;
import java.util.Objects;
import java.util.Map.Entry;
import net.minecraft.world.item.ItemStack;

public class EntityEquipment {
	public static final Codec<EntityEquipment> CODEC = Codec.unboundedMap(EquipmentSlot.CODEC, ItemStack.CODEC).xmap(map -> {
		EnumMap<EquipmentSlot, ItemStack> enumMap = new EnumMap(EquipmentSlot.class);
		enumMap.putAll(map);
		return new EntityEquipment(enumMap);
	}, entityEquipment -> {
		Map<EquipmentSlot, ItemStack> map = new EnumMap(entityEquipment.items);
		map.values().removeIf(ItemStack::isEmpty);
		return map;
	});
	private final EnumMap<EquipmentSlot, ItemStack> items;

	private EntityEquipment(EnumMap<EquipmentSlot, ItemStack> items) {
		this.items = items;
	}

	public EntityEquipment() {
		this(new EnumMap(EquipmentSlot.class));
	}

	public ItemStack set(EquipmentSlot slot, ItemStack stack) {
		stack.getItem().verifyComponentsAfterLoad(stack);
		return (ItemStack)Objects.requireNonNullElse((ItemStack)this.items.put(slot, stack), ItemStack.EMPTY);
	}

	public ItemStack get(EquipmentSlot slot) {
		return (ItemStack)this.items.getOrDefault(slot, ItemStack.EMPTY);
	}

	public boolean isEmpty() {
		for (ItemStack itemStack : this.items.values()) {
			if (!itemStack.isEmpty()) {
				return false;
			}
		}

		return true;
	}

	public void tick(Entity entity) {
		for (Entry<EquipmentSlot, ItemStack> entry : this.items.entrySet()) {
			ItemStack itemStack = (ItemStack)entry.getValue();
			if (!itemStack.isEmpty()) {
				itemStack.inventoryTick(entity.level(), entity, (EquipmentSlot)entry.getKey());
			}
		}
	}

	public void setAll(EntityEquipment equipment) {
		this.items.clear();
		this.items.putAll(equipment.items);
	}

	public void dropAll(LivingEntity entity) {
		for (ItemStack itemStack : this.items.values()) {
			entity.drop(itemStack, true, false);
		}

		this.clear();
	}

	public void clear() {
		this.items.replaceAll((equipmentSlot, itemStack) -> ItemStack.EMPTY);
	}
}