RelicCraft/src/main/java/com/thebrokenrail/reliccraft/mixin/MixinEntity.java

package com.thebrokenrail.reliccraft.mixin;

import com.thebrokenrail.reliccraft.block.TeleportationRestrictorBlock;
import net.minecraft.entity.Entity;
import net.minecraft.entity.EntityType;
import net.minecraft.entity.data.DataTracker;
import net.minecraft.entity.data.TrackedData;
import net.minecraft.entity.data.TrackedDataHandlerRegistry;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.world.World;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
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.callback.CallbackInfo;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;

@SuppressWarnings("unused")
@Mixin(Entity.class)
public abstract class MixinEntity implements TeleportationRestrictorBlock.TeleportingEntity {
    @Shadow
    protected boolean inNetherPortal;

    @Shadow
    public abstract World getEntityWorld();

    @Shadow
    public abstract DataTracker getDataTracker();

    @Unique
    private static final byte DEFAULT_TELEPORT_COOLDOWN = 80;

    @Unique
    private static final TrackedData<Byte> TELEPORT_COOLDOWN = DataTracker.registerData(Entity.class, TrackedDataHandlerRegistry.BYTE);

    @Inject(at = @At("RETURN"), method = "<init>")
    public void init(EntityType<?> type, World world, CallbackInfo info) {
        getDataTracker().startTracking(TELEPORT_COOLDOWN, (byte) 0);
    }

    @Override
    public void resetTeleportCooldown() {
        getDataTracker().set(TELEPORT_COOLDOWN, (byte) 0);
    }

    @Inject(at = @At("RETURN"), method = "baseTick")
    public void baseTick(CallbackInfo info) {
        updateNetherPortal();
        byte data = getDataTracker().get(TELEPORT_COOLDOWN);
        if (getEntityWorld() instanceof ServerWorld && data < DEFAULT_TELEPORT_COOLDOWN) {
            data++;
            getDataTracker().set(TELEPORT_COOLDOWN, data);
        }
    }

    @Unique
    private void updateNetherPortal() {
        if (cannotTeleport()) {
            inNetherPortal = false;
        }
    }

    @Inject(at = @At("HEAD"), method = "canUsePortals", cancellable = true)
    public void canUsePortals(CallbackInfoReturnable<Boolean> info) {
        if (cannotTeleport()) {
            info.setReturnValue(false);
        }
    }

    @Override
    public boolean cannotTeleport() {
        return getDataTracker().get(TELEPORT_COOLDOWN) < DEFAULT_TELEPORT_COOLDOWN;
    }

    @Inject(at = @At("HEAD"), method = "toTag")
    public void toTag(CompoundTag tag, CallbackInfoReturnable<CompoundTag> info) {
        tag.putByte("RelicCraftTeleportCooldown", getDataTracker().get(TELEPORT_COOLDOWN));
    }

    @Inject(at = @At("HEAD"), method = "fromTag")
    public void fromTag(CompoundTag tag, CallbackInfo info) {
        getDataTracker().set(TELEPORT_COOLDOWN, tag.getByte("RelicCraftTeleportCooldown"));
        updateNetherPortal();
    }
}