Twine/src/main/java/com/thebrokenrail/twine/mixin/MixinBoatEntity.java

package com.thebrokenrail.twine.mixin;

import com.thebrokenrail.twine.util.boat.BoatChestMode;
import com.thebrokenrail.twine.util.boat.BoatChestModes;
import com.thebrokenrail.twine.util.boat.BoatUtil;
import com.thebrokenrail.twine.util.boat.inventory.BoatInventory;
import com.thebrokenrail.twine.util.boat.inventory.BoatInventoryWrapper;
import net.minecraft.block.Block;
import net.minecraft.block.BlockState;
import net.minecraft.enchantment.EnchantmentHelper;
import net.minecraft.entity.Entity;
import net.minecraft.entity.damage.DamageSource;
import net.minecraft.entity.data.DataTracker;
import net.minecraft.entity.data.TrackedData;
import net.minecraft.entity.data.TrackedDataHandlerRegistry;
import net.minecraft.entity.player.PlayerEntity;
import net.minecraft.entity.vehicle.BoatEntity;
import net.minecraft.inventory.Inventory;
import net.minecraft.item.ItemStack;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.screen.SimpleNamedScreenHandlerFactory;
import net.minecraft.server.network.ServerPlayerEntity;
import net.minecraft.sound.SoundEvents;
import net.minecraft.util.ActionResult;
import net.minecraft.util.Hand;
import net.minecraft.util.math.BlockPos;
import net.minecraft.world.GameRules;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Unique;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.Redirect;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;

import java.util.List;

@Mixin(BoatEntity.class)
public class MixinBoatEntity implements BoatUtil {
    private static final TrackedData<Integer> CHEST_MODE = DataTracker.registerData(BoatEntity.class, TrackedDataHandlerRegistry.INTEGER);

    @Unique
    private BoatInventory items = null;
    @Unique
    private ItemStack stack = ItemStack.EMPTY;

    @Inject(at = @At("RETURN"), method = "initDataTracker")
    public void initDataTracker(CallbackInfo info) {
        ((BoatEntity) (Object) this).getDataTracker().startTracking(CHEST_MODE, BoatChestModes.NONE.getID());
    }

    @Inject(at = @At("HEAD"), method = "canAddPassenger", cancellable = true)
    public void canAddPassenger(Entity passenger, CallbackInfoReturnable<Boolean> info) {
        if (((BoatEntity) (Object) this).getPassengerList().size() > 0 && getChestMode() != BoatChestModes.NONE) {
            info.setReturnValue(false);
        }
    }

    @Redirect(at = @At(value = "INVOKE", target = "Ljava/util/List;size()I"), method = "updatePassengerPosition", allow = 2, require = 2)
    public int updatePassengerPosition(List<Entity> list) {
        if (getChestMode() != BoatChestModes.NONE) {
            return list.size() + 1;
        } else {
            return list.size();
        }
    }

    @Inject(at = @At("HEAD"), method = "interact", cancellable = true)
    public void interact(PlayerEntity player, Hand hand, CallbackInfoReturnable<ActionResult> info) {
        BoatChestMode mode = getChestMode();
        ItemStack itemStack = player.getStackInHand(hand);

        BoatChestMode newMode = BoatChestMode.valueOf(Block.getBlockFromItem(itemStack.getItem()));

        if (newMode != BoatChestModes.NONE && mode == BoatChestModes.NONE) {
            List<Entity> passengers = ((BoatEntity) (Object) this).getPassengerList();
            for (int i = 1; i < passengers.size(); i++) {
                passengers.get(i).stopRiding();
            }
            if (!player.getEntityWorld().isClient()) {
                ItemStack newStack;
                if (!player.isCreative()) {
                    newStack = itemStack.split(1);
                } else {
                    newStack = itemStack.copy();
                    newStack.setCount(1);
                }

                setChestItem(newStack);

                newMode.triggerAdvancement((ServerPlayerEntity) player);

                player.getEntityWorld().playSound(null, ((BoatEntity) (Object) this).getBlockPos(), SoundEvents.ITEM_ARMOR_EQUIP_GENERIC, ((BoatEntity) (Object) this).getSoundCategory(), 1f, 1f);
            }
            info.setReturnValue(ActionResult.SUCCESS);
        } else if (player.shouldCancelInteraction()) {
            if (!player.getEntityWorld().isClient()) {
                openInventory(player);
            }
            info.setReturnValue(ActionResult.SUCCESS);
        }
    }

    @Unique
    private void updateChestMode() {
        BoatChestMode mode = BoatChestMode.valueOf(Block.getBlockFromItem(stack.getItem()));

        ((BoatEntity) (Object) this).getDataTracker().set(CHEST_MODE, mode.getID());
    }

    @Unique
    private void dropItemsOnDestroy(boolean isCreative) {
        if (((BoatEntity) (Object) this).getEntityWorld().getGameRules().getBoolean(GameRules.DO_ENTITY_DROPS)) {
            BoatChestMode mode = getChestMode();
            if (mode.containsItems()) {
                ((BoatEntity) (Object) this).dropStack(stack.copy());
            } else {
                if (!isCreative) {
                    ((BoatEntity) (Object) this).dropStack(getStackWithoutItems());
                }
                dropItems();
            }
        }
    }

    @Inject(at = @At(value = "INVOKE", target = "Lnet/minecraft/entity/vehicle/BoatEntity;remove()V"), method = "damage")
    public void damage(DamageSource source, float amount, CallbackInfoReturnable<Boolean> info) {
        dropItemsOnDestroy(source.getAttacker() instanceof PlayerEntity && ((PlayerEntity) source.getAttacker()).isCreative());
    }

    @Inject(at = @At(value = "INVOKE", target = "Lnet/minecraft/entity/vehicle/BoatEntity;remove()V"), method = "fall")
    public void fall(double heightDifference, boolean onGround, BlockState landedState, BlockPos landedPosition, CallbackInfo info) {
        dropItemsOnDestroy(false);
    }

    @Unique
    private void updateInventory() {
        BoatChestMode mode = getChestMode();
        items = mode.hasItems() ? new BoatInventory((BoatEntity) (Object) this, mode.getSize(), stack, mode) : null;
    }

    @Unique
    private void dropItems() {
        if (items != null) {
            for (int i = 0; i < items.size(); ++i) {
                ItemStack itemStack = items.getStack(i).copy();
                if (!itemStack.isEmpty() && !EnchantmentHelper.hasVanishingCurse(itemStack)) {
                    ((BoatEntity) (Object) this).dropStack(itemStack);
                }
            }
        }
    }

    @Unique
    private void setChestItem(ItemStack newStack) {
        stack = newStack;

        updateChestMode();
        updateInventory();
    }

    @Unique
    private ItemStack getStackWithoutItems() {
        ItemStack newStack = stack.copy();
        CompoundTag tag = newStack.getOrCreateTag();

        CompoundTag entityTag = tag.getCompound("BlockEntityTag");
        entityTag.remove("Items");
        if (entityTag.getSize() < 1) {
            tag.remove("BlockEntityTag");
        }

        return newStack;
    }

    @Inject(at = @At("RETURN"), method = "writeCustomDataToTag")
    public void writeCustomDataToTag(CompoundTag tag, CallbackInfo info) {
        tag.put("TwineChestItem", stack.toTag(new CompoundTag()));
    }

    @Inject(at = @At("RETURN"), method = "readCustomDataFromTag")
    public void readCustomDataFromTag(CompoundTag tag, CallbackInfo info) {
        CompoundTag itemTag = tag.getCompound("TwineChestItem");
        ItemStack newStack = ItemStack.fromTag(itemTag);

        setChestItem(newStack);
    }

    @Override
    public BoatChestMode getChestMode() {
        return BoatChestMode.valueOf(((BoatEntity) (Object) this).getDataTracker().get(CHEST_MODE));
    }

    @Override
    public Inventory getChestInventory() {
        return items;
    }

    @Override
    public void openInventory(PlayerEntity player) {
        BoatChestMode mode = getChestMode();
        player.openHandledScreen(new SimpleNamedScreenHandlerFactory((i, playerInventory, playerEntity) -> mode.getScreenHandlerFactory().createMenu(i, playerInventory, playerEntity, MixinBoatEntity.this, inventory -> new BoatInventoryWrapper((Entity) (Object) this, inventory, mode)), mode.getScreenHandlerName(stack)));
    }
}