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

package com.thebrokenrail.energonrelics.energy.core;

import com.thebrokenrail.energonrelics.component.NetworkComponent;
import com.thebrokenrail.energonrelics.energy.core.util.Action;
import com.thebrokenrail.energonrelics.energy.core.util.EnergyTickable;
import net.minecraft.block.BlockState;
import net.minecraft.block.entity.BlockEntityType;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.IntArrayTag;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.math.Vec3d;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public abstract class EnergyReceiverBlockEntity extends EnergyProviderBlockEntity {
    public EnergyReceiverBlockEntity(BlockEntityType<?> type) {
        super(type);
    }

    private final List<EnergyProviderBlockEntity> providers = new ArrayList<>();
    private final List<Integer> networks = new ArrayList<>();

    private long totalCost = 0;
    private long previousTotalCost = 0;

    private final List<Action.PropagatedAction> sent = new ArrayList<>();

    protected void propagateAction(Action.PropagatedAction action) {
        sent.add(action);

        totalCost = totalCost + action.amountOwed();
        if (providers.size() > 0) {
            action.expandPayments(providers.size());
            for (EnergyProviderBlockEntity provider : providers) {
                provider.addPropagatedAction(action);
            }
        } else {
            action.pay(0);
        }
    }

    protected void addAction(Action action) {
        propagateAction(new Action.PropagatedActionImpl(action, getWorld(), getPos(), getCachedState()));
    }

    @Override
    public List<EnergyTickable> startTick() {
        sent.clear();

        List<EnergyTickable> list = new ArrayList<>(super.startTick());

        ServerWorld world = (ServerWorld) getWorld();
        assert world != null;
        NetworkComponent component = NetworkComponent.getInstance(world);

        providers.clear();
        for (int network : networks) {
            List<EnergyProviderBlockEntity> potentialProviders = component.getSources(world, network);
            for (EnergyProviderBlockEntity provider : potentialProviders) {
                if (provider.getRegistryKey().equals(getWorld().getRegistryKey())) {
                    Vec3d pos = new Vec3d(getPos().getX() + 0.5d, getPos().getY() + 0.5d, getPos().getZ() + 0.5d);
                    if (provider.isWithinDistance(pos)) {
                        providers.add(provider);
                    }
                }
            }
        }

        previousTotalCost = totalCost;
        totalCost = 0;

        list.addAll(providers);

        return list;
    }

    @Override
    public void logicTick() {
        super.logicTick();
        assert getWorld() != null;
        getWorld().getProfiler().push("energyTick");
        energyTick();
        getWorld().getProfiler().pop();
    }

    public long getTotalCost() {
        return previousTotalCost;
    }

    @Override
    protected void handlePropagatedAction(Action.PropagatedAction action) {
        super.handlePropagatedAction(action);
        // Propagate Action To Energy Providers
        if (isEnergyProviderActive() && !sent.contains(action)) {
            propagateAction(action);
        } else {
            action.pay(0);
        }
    }

    @Override
    public CompoundTag toTag(CompoundTag tag) {
        super.toTag(tag);
        tag.put("Networks", new IntArrayTag(networks));
        return tag;
    }

    @Override
    public void fromTag(BlockState state, CompoundTag tag) {
        super.fromTag(state, tag);
        int[] networksArray = tag.getIntArray("Networks");
        networks.clear();
        for (int network : networksArray) {
            networks.add(network);
        }
    }

    protected boolean isEnergyProviderActive() {
        throw new UnsupportedOperationException();
    }

    protected abstract void energyTick();

    public void toggle(int network) {
        if (contains(network)) {
            networks.removeAll(Collections.singletonList(network));
        } else {
            networks.add(network);
        }
        markDirty();
    }

    public boolean contains(int network) {
        return networks.contains(network);
    }
}