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

package com.thebrokenrail.twine.mixin;

import com.thebrokenrail.twine.Twine;
import com.thebrokenrail.twine.util.BoatInventory;
import com.thebrokenrail.twine.util.BoatUtil;
import com.thebrokenrail.twine.util.EnderChestInventoryWrapper;
import net.minecraft.block.BlockState;
import net.minecraft.block.EnderChestBlock;
import net.minecraft.enchantment.EnchantmentHelper;
import net.minecraft.entity.Entity;
import net.minecraft.entity.damage.DamageSource;
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.nbt.ListTag;
import net.minecraft.screen.GenericContainerScreenHandler;
import net.minecraft.screen.SimpleNamedScreenHandlerFactory;
import net.minecraft.server.network.ServerPlayerEntity;
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 {
    @Inject(at = @At("RETURN"), method = "initDataTracker")
    public void initDataTracker(CallbackInfo info) {
        ((BoatEntity) (Object) this).getDataTracker().startTracking(BoatUtil.HAS_CHEST, BoatChestMode.NONE.name());
    }

    @Inject(at = @At("HEAD"), method = "canAddPassenger", cancellable = true)
    public void canAddPassenger(Entity passenger, CallbackInfoReturnable<Boolean> info) {
        if (((BoatEntity) (Object) this).getPassengerList().size() > 0 && getChestMode() != BoatChestMode.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() != BoatChestMode.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 hasChest = getChestMode();
        ItemStack stack = player.getStackInHand(hand);
        BoatChestMode newMode = BoatChestMode.valueOF(stack.getItem());
        if (newMode != BoatChestMode.NONE && hasChest == BoatChestMode.NONE) {
            List<Entity> passengers = ((BoatEntity) (Object) this).getPassengerList();
            for (int i = 1; i < passengers.size(); i++) {
                passengers.get(i).stopRiding();
            }
            if (!player.getEntityWorld().isClient()) {
                if (newMode.hasItems()) {
                    Twine.CHEST_BOAT_CRITERION.trigger((ServerPlayerEntity) player);
                } else if (newMode == BoatChestMode.ENDER_CHEST) {
                    Twine.ENDER_CHEST_BOAT_CRITERION.trigger((ServerPlayerEntity) player);
                }
                setChestMode(newMode);
                updateInventory();
            }
            if (!player.isCreative()) {
                stack.decrement(1);
            }
            info.setReturnValue(ActionResult.SUCCESS);
        } else if (player.isSneaking()) {
            if (!player.getEntityWorld().isClient()) {
                openInventory(player);
            }
            info.setReturnValue(ActionResult.SUCCESS);
        }
    }

    @Unique
    private void setChestMode(BoatChestMode mode) {
        ((BoatEntity) (Object) this).getDataTracker().set(BoatUtil.HAS_CHEST, mode.name());
    }

    @Unique
    private void dropItemsOnDestroy(boolean isCreative) {
        if (((BoatEntity) (Object) this).getEntityWorld().getGameRules().getBoolean(GameRules.DO_ENTITY_DROPS)) {
            if (!isCreative) {
                ((BoatEntity) (Object) this).dropStack(new ItemStack(getChestMode().getItem()));
            }
            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 BoatInventory items;

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

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

    @Inject(at = @At("RETURN"), method = "writeCustomDataToTag")
    public void writeCustomDataToTag(CompoundTag tag, CallbackInfo info) {
        BoatChestMode hasChest = getChestMode();
        tag.putString("HasChest", hasChest.name());
        if (hasChest.hasItems()) {
            ListTag listTag = new ListTag();

            for (int i = 2; i < this.items.size(); ++i) {
                ItemStack itemStack = this.items.getStack(i);
                if (!itemStack.isEmpty()) {
                    CompoundTag compoundTag = new CompoundTag();
                    compoundTag.putByte("Slot", (byte)i);
                    itemStack.toTag(compoundTag);
                    listTag.add(compoundTag);
                }
            }

            tag.put("Items", listTag);
        }
    }

    @Inject(at = @At("RETURN"), method = "readCustomDataFromTag")
    public void readCustomDataFromTag(CompoundTag tag, CallbackInfo info) {
        BoatChestMode hasChest;
        try {
            hasChest = BoatChestMode.valueOf(tag.getString("HasChest"));
        } catch (IllegalArgumentException e) {
            hasChest = BoatChestMode.NONE;
        }
        setChestMode(hasChest);
        updateInventory();
        if (hasChest.hasItems()) {
            ListTag listTag = tag.getList("Items", 10);

            for (int i = 0; i < listTag.size(); ++i) {
                CompoundTag compoundTag = listTag.getCompound(i);
                int j = compoundTag.getByte("Slot") & 255;
                if (j >= 2 && j < items.size()) {
                    items.setStack(j, ItemStack.fromTag(compoundTag));
                }
            }
        }
    }

    @Override
    public BoatChestMode getChestMode() {
        try {
            return BoatChestMode.valueOf(((BoatEntity) (Object) this).getDataTracker().get(BoatUtil.HAS_CHEST));
        } catch (IllegalArgumentException e) {
            return BoatChestMode.NONE;
        }
    }

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

    @Override
    public void openInventory(PlayerEntity player) {
        if (getChestMode().hasItems()) {
            player.openHandledScreen(new SimpleNamedScreenHandlerFactory((i, playerInventory, playerEntity) -> GenericContainerScreenHandler.createGeneric9x3(i, playerInventory, getChestInventory()), ((BoatEntity) (Object) this).getDisplayName()));
        } else {
            player.openHandledScreen(new SimpleNamedScreenHandlerFactory((i, playerInventory, playerEntity) -> GenericContainerScreenHandler.createGeneric9x3(i, playerInventory, new EnderChestInventoryWrapper((BoatEntity) (Object) this, player.getEnderChestInventory())), EnderChestBlock.CONTAINER_NAME));
        }
    }
}