EnergonRelics/src/main/java/com/thebrokenrail/energonrelics/energy/core/EnergyProviderBlockEntity.java

package com.thebrokenrail.energonrelics.energy.core;

import com.thebrokenrail.energonrelics.config.HardcodedConfig;
import com.thebrokenrail.energonrelics.EnergonRelics;
import com.thebrokenrail.energonrelics.component.NetworkComponent;
import com.thebrokenrail.energonrelics.energy.core.util.Action;
import com.thebrokenrail.energonrelics.energy.core.util.EnergyTickable;
import com.thebrokenrail.energonrelics.energy.core.util.EnergyTicker;
import com.thebrokenrail.energonrelics.util.BlockPosWithDimension;
import net.fabricmc.fabric.api.block.entity.BlockEntityClientSerializable;
import net.minecraft.block.BlockState;
import net.minecraft.block.Blocks;
import net.minecraft.block.entity.BlockEntity;
import net.minecraft.block.entity.BlockEntityType;
import net.minecraft.item.ItemStack;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.Tickable;
import net.minecraft.util.math.BlockPos;
import net.minecraft.util.math.Vec3d;
import net.minecraft.util.registry.RegistryKey;
import net.minecraft.world.World;

import java.util.Collections;
import java.util.List;
import java.util.Objects;

public class EnergyProviderBlockEntity extends BlockEntity implements BlockEntityClientSerializable, Tickable, EnergyTickable {
    public EnergyProviderBlockEntity(BlockEntityType<?> type) {
        super(type);
    }

    public final void addPropagatedAction(Action.PropagatedAction action) {
        if (isEnergyProvider()) {
            handlePropagatedAction(action);
        } else {
            throw new UnsupportedOperationException();
        }
    }

    public final boolean isNetwork(int network) {
        return isEnergyProvider() && EnergonRelics.NETWORK_CHIP_ITEM.getID(stack) == network;
    }

    public final boolean isWithinDistance(Vec3d pos) {
        return getPos().isWithinDistance(pos, HardcodedConfig.POWER_RANGE);
    }

    public RegistryKey<World> getRegistryKey() {
        return Objects.requireNonNull(getWorld()).getRegistryKey();
    }

    private ItemStack stack = ItemStack.EMPTY;

    @Override
    public CompoundTag toTag(CompoundTag tag) {
        super.toTag(tag);
        if (isEnergyProvider()) {
            tag.put("NetworkChip", stack.toTag(new CompoundTag()));
        }
        return tag;
    }

    @Override
    public void fromTag(BlockState state, CompoundTag tag) {
        super.fromTag(state, tag);
        if (isEnergyProvider()) {
            stack = ItemStack.fromTag(tag.getCompound("NetworkChip"));
        } else {
            stack = ItemStack.EMPTY;
        }
    }

    public boolean isEnergyProvider() {
        return false;
    }

    protected void handlePropagatedAction(Action.PropagatedAction action) {
        if (!isEnergyProvider()) {
            throw new UnsupportedOperationException();
        }
    }

    @Override
    public void logicTick() {
    }

    @Override
    public List<EnergyTickable> startTick() {
        if (isEnergyProvider()) {
            NetworkComponent component = NetworkComponent.getInstance((ServerWorld) Objects.requireNonNull(getWorld()));
            List<BlockPosWithDimension> sources = component.getSourcePos(EnergonRelics.NETWORK_CHIP_ITEM.getID(stack));
            if (!sources.contains(new BlockPosWithDimension(getPos(), getWorld().getRegistryKey()))) {
                takeStack(getWorld());
            }
        }
        return Collections.emptyList();
    }

    @Override
    public final void tick() {
        if (hasWorld() && !Objects.requireNonNull(getWorld()).isClient()) {
            EnergyTicker.schedule(this);
        }
    }

    private static void setEnergyProviderSource(World world, BlockPos pos, ItemStack stack, boolean remove) {
        ServerWorld serverWorld = (ServerWorld) world;
        NetworkComponent component = NetworkComponent.getInstance(serverWorld);
        BlockPosWithDimension newPos = new BlockPosWithDimension(pos, world.getRegistryKey());
        if (remove) {
            component.removeSource(EnergonRelics.NETWORK_CHIP_ITEM.getID(stack), newPos);
        } else {
            component.addSource(EnergonRelics.NETWORK_CHIP_ITEM.getID(stack), newPos);
        }
    }

    public ItemStack takeStack(World world) {
        ItemStack newStack = stack.copy();
        setEnergyProviderSource(world, getPos(), newStack, true);
        stack = ItemStack.EMPTY;
        markDirty();
        return newStack;
    }

    public void placeStack(ItemStack newStack, World world) {
        setEnergyProviderSource(world, getPos(), newStack, false);
        stack = newStack.copy();
        markDirty();
    }

    public boolean hasStack() {
        return !stack.isEmpty();
    }

    @Override
    public void fromClientTag(CompoundTag compoundTag) {
        fromTag(Blocks.AIR.getDefaultState(), compoundTag);
    }

    @Override
    public CompoundTag toClientTag(CompoundTag compoundTag) {
        return toTag(compoundTag);
    }

    @Override
    public void markDirty() {
        super.markDirty();
        if (hasWorld()) {
            sync();
        }
    }
}