Omega Stacks

OmegaStackEntity

package io.github.daomephsta.omegastacks.common.entity;

import static java.util.stream.Collectors.toCollection;

import java.util.*;

import com.google.common.collect.*;

import io.github.daomephsta.omegastacks.common.OmegaStacks;
import net.fabricmc.fabric.api.util.NbtType;
import net.minecraft.client.network.packet.EntitySpawnS2CPacket;
import net.minecraft.entity.*;
import net.minecraft.entity.damage.DamageSource;
import net.minecraft.entity.data.*;
import net.minecraft.item.ItemStack;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.network.Packet;
import net.minecraft.world.World;

public class OmegaStackEntity extends Entity
{
	private static final TrackedData<ItemStack> DISPLAY_STACK_0 = 
		DataTracker.registerData(OmegaStackEntity.class, TrackedDataHandlerRegistry.ITEM_STACK),
		
												DISPLAY_STACK_1 = 
		DataTracker.registerData(OmegaStackEntity.class, TrackedDataHandlerRegistry.ITEM_STACK),
		
												DISPLAY_STACK_2 = 
		DataTracker.registerData(OmegaStackEntity.class, TrackedDataHandlerRegistry.ITEM_STACK),
		
												DISPLAY_STACK_3 = 
		DataTracker.registerData(OmegaStackEntity.class, TrackedDataHandlerRegistry.ITEM_STACK),
		
												DISPLAY_STACK_4 = 
		DataTracker.registerData(OmegaStackEntity.class, TrackedDataHandlerRegistry.ITEM_STACK);
	
	private static final List<TrackedData<ItemStack>> DISPLAY_STACKS = 
		ImmutableList.of(DISPLAY_STACK_0, DISPLAY_STACK_1, DISPLAY_STACK_2, DISPLAY_STACK_3, DISPLAY_STACK_4);
	
	private final Iterable<ItemStack> displayStacks = Iterables.transform(DISPLAY_STACKS, getDataTracker()::get);
	
	private int age;
	private float health = 5;
	private SubStackList subStacks;

	public OmegaStackEntity(World world, List<ItemEntity> subStackEntities)
	{
		this(OmegaStacks.OMEGA_STACK_ENTITY_TYPE, world, subStackEntities);
	}
	
	public OmegaStackEntity(EntityType<OmegaStackEntity> entityType, World world)
	{
		this(entityType, world, Collections.emptyList());
	}

	private OmegaStackEntity(EntityType<OmegaStackEntity> entityType, World world, List<ItemEntity> subStackEntities)
	{
		super(entityType, world);
		this.subStacks = new SubStackList(this, subStackEntities);
	}

	@Override
	protected void initDataTracker()
	{
		for (TrackedData<ItemStack> displayStack : DISPLAY_STACKS)
		{
			getDataTracker().startTracking(displayStack, ItemStack.EMPTY);
		}
	}
	
	@Override
	public boolean damage(DamageSource source, float amount)
	{
		if (this.isInvulnerableTo(source))
			return false;
		
		this.scheduleVelocityUpdate();
		health -= amount;
		if (health <= 0.0F)
			this.remove();
		return true;
	}
	
	@Override
	public void tick()
	{
		super.tick();
		System.out.println(this.getPos());
		age++;
	}
	
	public Iterable<ItemStack> getDisplayStacks()
	{
		return displayStacks;
	}

	@Override
	protected void readCustomDataFromTag(CompoundTag tag)
	{
		this.age = tag.getInt("age");
		this.health = tag.getFloat("health");
		this.subStacks = tag.getList("sub_stacks", NbtType.COMPOUND).stream()
			.map(subTag -> ItemStack.fromTag((CompoundTag) subTag))
			.collect(toCollection(() -> new SubStackList(this)));
	}

	@Override
	protected void writeCustomDataToTag(CompoundTag tag)
	{
		tag.putInt("age", age);
		tag.putFloat("health", health);
		tag.put("sub_stacks", subStacks.stream()
				.map(stack -> stack.toTag(new CompoundTag()))
				.collect(toCollection(ListTag::new)));
	}

	@Override
	public Packet<?> createSpawnPacket()
	{
		return new EntitySpawnS2CPacket(this);
	}
	
	private static class SubStackList extends ForwardingList<ItemStack>
	{
		private final OmegaStackEntity omegaStack;
		private final List<ItemStack> subStacks;

		private SubStackList(OmegaStackEntity omegaStack)
		{
			this(omegaStack, Collections.emptyList());
		}
		
		private SubStackList(OmegaStackEntity omegaStack, List<ItemEntity> subStackEntities)
		{
			this.omegaStack = omegaStack;
			this.subStacks = new ArrayList<>();
			for (ItemEntity itemEntity : subStackEntities)
			{
				ItemStack stack = itemEntity.getStack();
				ItemStack toAdd = stack;
				for (ItemStack subStack : subStacks)
				{
					boolean canMerge = subStack.getItem() == stack.getItem() && ItemStack.areTagsEqual(subStack, stack);
					if (canMerge)
					{
						subStack.addAmount(stack.getAmount());
						toAdd = ItemStack.EMPTY;
					}
				}
				if (!toAdd.isEmpty())
					subStacks.add(toAdd);
				itemEntity.remove();
			}
			System.out.println(subStacks);
		}

		@Override
		protected List<ItemStack> delegate()
		{
			return subStacks;
		}

		@Override
		public void add(int index, ItemStack stack)
		{
			subStacks.add(index, stack);
			if (index < DISPLAY_STACKS.size())
				omegaStack.getDataTracker().set(DISPLAY_STACKS.get(index), stack);
		}
		
		@Override
		public ItemStack get(int index)
		{
			return subStacks.get(index);
		}

		@Override
		public int size()
		{
			return subStacks.size();
		}
	}
}