copper-farmers/src/main/java/com/dynamitos/mixin/MoveItemsTaskMixin.java

package com.dynamitos.mixin;

import java.util.Map;
import java.util.Optional;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Predicate;

import org.spongepowered.asm.mixin.Debug;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;

import com.dynamitos.Copperfarmers;
import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;

import net.minecraft.block.BlockState;
import net.minecraft.block.Blocks;
import net.minecraft.entity.EquipmentSlot;
import net.minecraft.entity.ai.brain.MemoryModuleState;
import net.minecraft.entity.ai.brain.MemoryModuleType;
import net.minecraft.entity.ai.brain.task.MoveItemsTask;
import net.minecraft.entity.ai.brain.task.MultiTickTask;
import net.minecraft.entity.mob.PathAwareEntity;
import net.minecraft.item.ItemStack;
import net.minecraft.item.Items;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.math.BlockPos;
import net.minecraft.util.math.Box;
import net.minecraft.world.World;

@Mixin(MoveItemsTask.class)
public abstract class MoveItemsTaskMixin extends MultiTickTask<PathAwareEntity> {
    private MoveItemsTaskAccessor thisAccessor = (MoveItemsTaskAccessor) this;

    public MoveItemsTaskMixin(Map<MemoryModuleType<?>, MemoryModuleState> requiredMemoryState) {
        super(requiredMemoryState);
        // TODO Auto-generated constructor stub
    }

    @Inject(method = "findStorage", at = @At("HEAD"), cancellable = true)
    private void findStorage(ServerWorld world, PathAwareEntity entity,
            CallbackInfoReturnable<Optional<MoveItemsTask.Storage>> callbackInfo) {
        Copperfarmers.LOGGER.debug("Wrapped findStorage called");

        Box box = ((MoveItemsTaskAccessor) this).copperfarmers$getSearchBoundingBox(entity);
        if (!entity.getEquippedStack(EquipmentSlot.MAINHAND).isOf(Items.WHEAT_SEEDS)) {
            return;
        }
        for (int z = (int) box.minZ; z <= box.maxZ; z++) {
            for (int y = (int) box.minY; y <= box.maxY; y++) {
                for (int x = (int) box.minX; x <= box.maxX; x++) {
                    var pos = new BlockPos(x, y, z);
                    var blockState = world.getBlockState(pos);
                    // if block is farmland
                    if (blockState.isOf(Blocks.FARMLAND) && world.isAir(pos.up())) {
                        Copperfarmers.LOGGER.debug("Found farmland at " + pos);
                        if (Math.random() < 0.5) {
                            callbackInfo
                                    .setReturnValue(
                                            Optional.of(new MoveItemsTask.Storage(pos, null, null, blockState)));
                            return;
                        }
                    }
                }
            }
        }
    }

    @Inject(method = "testContainer", at = @At("HEAD"), cancellable = true)
    private void testContainer(PathAwareEntity entity, BlockState state, CallbackInfoReturnable<Boolean> callbackInfo) {
        Copperfarmers.LOGGER.debug("Wrapped testContainer called");
        if (!MoveItemsTaskAccessor.copperfarmers$canPickUpItem(entity) && state.isOf(Blocks.FARMLAND)) {
            Copperfarmers.LOGGER.debug("testContainer returning true for farmland");
            callbackInfo.setReturnValue(true);
        }
    }

    @Inject(method = "tickInteracting", at = @At("HEAD"), cancellable = true)
    private void tickInteracting(MoveItemsTask.Storage storage, World world, PathAwareEntity entity,
            CallbackInfo callbackInfo) {
        Copperfarmers.LOGGER.debug("Wrapped tickInteracting called");
        if (!thisAccessor.copperfarmers$isWithinRange(2.0, storage, world, entity,
                thisAccessor.copperfarmers$atCenterY(entity))) {
            thisAccessor.copperfarmers$transitionToTravelling(entity);
        } else {
            thisAccessor.copperfarmers$setInteractionTicks(thisAccessor.copperfarmers$getInteractionTicks() + 1);
            thisAccessor.copperfarmers$setLookTarget(storage, entity);
            Copperfarmers.LOGGER.debug("Interaction ticks: " + thisAccessor.copperfarmers$getInteractionTicks());
            if (thisAccessor.copperfarmers$getInteractionTicks() >= 60) {
                Copperfarmers.LOGGER.debug("Interacting with storage at " + storage.pos());
                this.selectInteractionState2(
                        entity,
                        storage,
                        this::takeStack2,
                        (entityxx, inventory) -> thisAccessor.copperfarmers$invalidateTargetStorage(entity),
                        this::placeStack2,
                        (entityxx, inventory) -> thisAccessor.copperfarmers$invalidateTargetStorage(entity));
                thisAccessor.copperfarmers$transitionToTravelling(entity);
            }
        }
        callbackInfo.cancel();
    }

    @Inject(method = "transitionToInteracting", at = @At("HEAD"), cancellable = true)
    private void transitionToInteracting(MoveItemsTask.Storage storage, PathAwareEntity entity,
            CallbackInfo callbackInfo) {
        Copperfarmers.LOGGER.debug("Wrapped transitionToInteracting called");
        this.selectInteractionState2(
                entity,
                storage,
                this.createSetInteractionStateCallback2(MoveItemsTask.InteractionState.PICKUP_ITEM),
                this.createSetInteractionStateCallback2(MoveItemsTask.InteractionState.PICKUP_NO_ITEM),
                this.createSetInteractionStateCallback2(MoveItemsTask.InteractionState.PLACE_ITEM),
                this.createSetInteractionStateCallback2(MoveItemsTask.InteractionState.PLACE_NO_ITEM));
        thisAccessor.copperfarmers$setNavigationState(MoveItemsTask.NavigationState.INTERACTING);
        callbackInfo.cancel();
    }

    @Inject(method = "isUnchanged", at = @At("HEAD"), cancellable = true)
    private void isUnchanged(World world, MoveItemsTask.Storage storage, CallbackInfoReturnable<Boolean> callbackInfo) {
        Copperfarmers.LOGGER.debug("Wrapped isUnchanged called");
        if (storage.state().isOf(Blocks.FARMLAND) && world.isAir(storage.pos().up())) {
            Copperfarmers.LOGGER.debug("isUnchanged returning true for farmland at " + storage.pos());
            callbackInfo.setReturnValue(true);
            return;
        }
        Copperfarmers.LOGGER.debug("Storage state: " + storage.state() + " isAir: " + world.isAir(storage.pos().up()));
        if (storage.blockEntity() == null) {
            Copperfarmers.LOGGER.debug("No blockentity " + storage.pos());
            callbackInfo.setReturnValue(false);
            return;
        }
    }

    private void placeStack2(PathAwareEntity entity, MoveItemsTask.Storage storage) {
        Copperfarmers.LOGGER.debug("Wrapped placeStack called");
        BlockPos pos = storage.pos();
        ServerWorld world = (ServerWorld) entity.getEntityWorld();
        if (world.getBlockState(pos).isOf(Blocks.FARMLAND) && world.isAir(pos.up())
                && entity.getEquippedStack(EquipmentSlot.MAINHAND).isOf(Items.WHEAT_SEEDS)) {
            world.setBlockState(pos.up(), Blocks.WHEAT.getDefaultState());
            entity.getEquippedStack(EquipmentSlot.MAINHAND).decrement(1);
            if (entity.getEquippedStack(EquipmentSlot.MAINHAND).isEmpty()) {
                thisAccessor.copperfarmers$resetVisitedPositions(entity);
            } else {
                thisAccessor.copperfarmers$invalidateTargetStorage(entity);
            }
            Copperfarmers.LOGGER.debug("Placed wheat seeds at " + pos.up());
            return;
        }
        if (storage.inventory() == null) {
            Copperfarmers.LOGGER.debug("No inventory to place items into at " + pos);
            return;
        }
        ItemStack itemStack = MoveItemsTaskAccessor.copperfarmers$insertStack(entity, storage.inventory());
        storage.inventory().markDirty();
        entity.equipStack(EquipmentSlot.MAINHAND, itemStack);
        if (itemStack.isEmpty()) {
            thisAccessor.copperfarmers$resetVisitedPositions(entity);
        } else {
            thisAccessor.copperfarmers$invalidateTargetStorage(entity);
        }
    }

    private void takeStack2(PathAwareEntity entity, MoveItemsTask.Storage storage) {
        entity.equipStack(EquipmentSlot.MAINHAND,
                MoveItemsTaskAccessor.copperfarmers$extractStack(storage.inventory()));
        entity.setDropGuaranteed(EquipmentSlot.MAINHAND);
        storage.inventory().markDirty();
        thisAccessor.copperfarmers$resetVisitedPositions(entity);
    }

    private void selectInteractionState2(
            PathAwareEntity entity,
            MoveItemsTask.Storage storage,
            BiConsumer<PathAwareEntity, MoveItemsTask.Storage> pickupItemCallback,
            BiConsumer<PathAwareEntity, MoveItemsTask.Storage> pickupNoItemCallback,
            BiConsumer<PathAwareEntity, MoveItemsTask.Storage> placeItemCallback,
            BiConsumer<PathAwareEntity, MoveItemsTask.Storage> placeNoItemCallback) {
        Copperfarmers.LOGGER.debug("Wrapped selectInteractionState called");
        if (storage == null && entity.getEquippedStack(EquipmentSlot.MAINHAND).isOf(Items.WHEAT_SEEDS)) {
            placeItemCallback.accept(entity, storage);
        }
        if (MoveItemsTaskAccessor.copperfarmers$canPickUpItem(entity)) {
            if (MoveItemsTaskAccessor.copperfarmers$hasItem(storage.inventory())) {
                pickupItemCallback.accept(entity, storage);
            } else {
                pickupNoItemCallback.accept(entity, storage);
            }
        } else if (storage.state().isOf(Blocks.FARMLAND)
                || MoveItemsTaskAccessor.copperfarmers$canInsert(entity, storage.inventory())) {
            placeItemCallback.accept(entity, storage);
        } else {
            placeNoItemCallback.accept(entity, storage);
        }
    }

    private BiConsumer<PathAwareEntity, MoveItemsTask.Storage> createSetInteractionStateCallback2(
            MoveItemsTask.InteractionState state) {
        return (entity, inventory) -> thisAccessor.copperfarmers$setInteractionState(state);
    }
}