EnergonRelics/src/main/java/com/thebrokenrail/energonrelics/component/NetworkComponent.java

package com.thebrokenrail.energonrelics.component;

import com.thebrokenrail.energonrelics.EnergonRelics;
import com.thebrokenrail.energonrelics.energy.core.util.EnergyProvider;
import net.minecraft.block.entity.BlockEntity;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.Tag;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.math.BlockPos;
import net.minecraft.world.PersistentState;
import net.minecraft.world.World;

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

public class NetworkComponent extends PersistentState {
    private static class Entry {
        private final List<BlockPos> sources = new ArrayList<>();
        private int id = 0;
    }

    private final List<Entry> networks = new ArrayList<>();

    private NetworkComponent() {
        super(EnergonRelics.NAMESPACE);
    }

    @Override
    public void fromTag(CompoundTag tag) {
        networks.clear();
        Tag list = tag.get("Networks");
        if (list instanceof ListTag) {
            ListTag networks = (ListTag) list;
            for (int i = 0; i < networks.size(); i++) {
                this.networks.add(getEntry(networks.getCompound(i)));
            }
        }
    }

    private Entry getEntry(CompoundTag tag) {
        Entry entry = new Entry();
        entry.id = tag.getInt("ID");
        entry.sources.clear();
        Tag list = tag.get("Sources");
        if (list instanceof ListTag) {
            ListTag sources = (ListTag) list;
            for (int i = 0; i < sources.size(); i++) {
                int x = sources.getCompound(i).getInt("X");
                int y = sources.getCompound(i).getInt("Y");
                int z = sources.getCompound(i).getInt("Z");
                entry.sources.add(new BlockPos(x, y, z));
            }
        }
        return entry;
    }

    private CompoundTag saveEntry(Entry entry) {
        CompoundTag tag = new CompoundTag();
        tag.putInt("ID", entry.id);
        ListTag sources = new ListTag();
        for (BlockPos pos : entry.sources) {
            CompoundTag posTag = new CompoundTag();
            posTag.putInt("X", pos.getX());
            posTag.putInt("Y", pos.getY());
            posTag.putInt("Z", pos.getZ());
            sources.add(posTag);
        }
        tag.put("Sources", sources);
        return tag;
    }

    @Override
    public CompoundTag toTag(CompoundTag tag) {
        ListTag networksTag = new ListTag();
        for (Entry entry : networks) {
            networksTag.add(saveEntry(entry));
        }
        tag.put("Networks", networksTag);
        return tag;
    }

    private Entry getOrCreate(int id) {
        for (Entry entry : networks) {
            if (entry.id == id) {
                return entry;
            }
        }
        Entry entry = new Entry();
        entry.id = id;
        entry.sources.clear();
        networks.add(entry);
        markDirty();
        return entry;
    }

    public static NetworkComponent getInstance(ServerWorld world) {
        return Objects.requireNonNull(world.getServer().getWorld(World.OVERWORLD)).getPersistentStateManager().getOrCreate(NetworkComponent::new, EnergonRelics.NAMESPACE);
    }

    public List<BlockPos> getSourcePos(int id) {
        Entry entry = getOrCreate(id);
        return entry.sources;
    }

    public List<EnergyProvider> getSource(World world, int id) {
        List<BlockPos> sources = getSourcePos(id);
        List<BlockPos> valid = new ArrayList<>();
        Iterator<BlockPos> iterator = sources.iterator();
        List<EnergyProvider> providers = new ArrayList<>();
        boolean dirty = false;
        while (iterator.hasNext()) {
            BlockPos pos = iterator.next();
            if (!valid.contains(pos)) {
                BlockEntity entity = world.getBlockEntity(pos);
                if (entity instanceof EnergyProvider && ((EnergyProvider) entity).isNetwork(id)) {
                    providers.add((EnergyProvider) entity);
                } else {
                    iterator.remove();
                    dirty = true;
                }
                valid.add(pos);
            } else {
                iterator.remove();
                dirty = true;
            }
        }
        if (dirty) {
            markDirty();
        }
        Collections.shuffle(providers);
        return providers;
    }

    public void addSource(int id, BlockPos pos) {
        Entry entry = getOrCreate(id);
        if (!entry.sources.contains(pos)) {
            entry.sources.add(pos);
        }
        markDirty();
    }

    public void removeSource(int id, BlockPos pos) {
        Entry entry = getOrCreate(id);
        entry.sources.remove(pos);
        markDirty();
    }

    public int create() {
        int id = -1;
        for (Entry entry : networks) {
            id = Math.max(id, entry.id);
        }
        return getOrCreate(id + 1).id;
    }

    @Override
    public void markDirty() {
        List<Integer> ids = new ArrayList<>();
        List<BlockPos> positions = new ArrayList<>();
        Iterator<Entry> iterator = networks.iterator();
        while (iterator.hasNext()) {
            Entry entry = iterator.next();
            boolean remove = false;
            if (ids.contains(entry.id)) {
                remove = true;
            } else {
                ids.add(entry.id);
            }
            if (!Collections.disjoint(positions, entry.sources)) {
                remove = true;
            } else {
                positions.addAll(entry.sources);
            }
            if (remove) {
                iterator.remove();
            }
        }
        super.markDirty();
    }
}