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

package com.thebrokenrail.energonrelics.component;

import com.thebrokenrail.energonrelics.EnergonRelics;
import com.thebrokenrail.energonrelics.api.block.entity.core.EnergyProviderBlockEntity;
import com.thebrokenrail.energonrelics.util.BlockPosWithDimension;
import net.minecraft.block.entity.BlockEntity;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.NbtOps;
import net.minecraft.nbt.Tag;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.Identifier;
import net.minecraft.util.math.BlockPos;
import net.minecraft.util.registry.RegistryKey;
import net.minecraft.world.PersistentState;
import net.minecraft.world.World;

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

public class NetworkComponent extends PersistentState {
    private static class Entry {
        private final List<BlockPosWithDimension> 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 networksTag = (ListTag) list;
            for (int i = 0; i < networksTag.size(); i++) {
                networks.add(getEntry(networksTag.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");
                RegistryKey<World> dimension = World.CODEC.parse(NbtOps.INSTANCE, sources.getCompound(i).get("Dimension")).result().orElse(World.OVERWORLD);
                entry.sources.add(new BlockPosWithDimension(new BlockPos(x, y, z), dimension));
            }
        }
        return entry;
    }

    private CompoundTag saveEntry(Entry entry) {
        CompoundTag tag = new CompoundTag();
        tag.putInt("ID", entry.id);
        ListTag sources = new ListTag();
        for (BlockPosWithDimension pos : entry.sources) {
            CompoundTag posTag = new CompoundTag();
            posTag.putInt("X", pos.pos.getX());
            posTag.putInt("Y", pos.pos.getY());
            posTag.putInt("Z", pos.pos.getZ());
            Identifier.CODEC.encodeStart(NbtOps.INSTANCE, pos.dimension.getValue()).result().ifPresent(tags -> posTag.put("Dimension", tags));
            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<BlockPosWithDimension> getSourcesPos(int id) {
        Entry entry = getOrCreate(id);
        return entry.sources;
    }

    private final Map<Integer, List<EnergyProviderBlockEntity>> cache = new HashMap<>();

    private List<EnergyProviderBlockEntity> getSourcesFromCache(World world, int id) {
        if (cache.containsKey(id)) {
            world.getProfiler().push("getSourcesFromCache");

            List<EnergyProviderBlockEntity> list = cache.get(id);

            boolean valid = true;
            for (EnergyProviderBlockEntity entity : list) {
                if (!entity.isNetwork(id) || entity.isRemoved()) {
                    valid = false;
                    break;
                }
            }

            world.getProfiler().pop();

            if (valid) {
                Collections.shuffle(list, world.random);

                return list;
            } else {
                cache.remove(id);

                return null;
            }
        } else {
            return null;
        }
    }

    private List<EnergyProviderBlockEntity> getSourcesFromCacheMiss(World world, int id) {
        world.getProfiler().push("getSourcesFromCacheMiss");

        List<BlockPosWithDimension> sources = getSourcesPos(id);
        List<BlockPosWithDimension> valid = new ArrayList<>();
        Iterator<BlockPosWithDimension> iterator = sources.iterator();
        List<EnergyProviderBlockEntity> providers = new ArrayList<>();
        boolean dirty = false;
        while (iterator.hasNext()) {
            BlockPosWithDimension pos = iterator.next();
            if (!valid.contains(pos)) {
                BlockEntity entity = Objects.requireNonNull(Objects.requireNonNull(world.getServer()).getWorld(pos.dimension)).getBlockEntity(pos.pos);
                if (entity instanceof EnergyProviderBlockEntity && ((EnergyProviderBlockEntity) entity).isNetwork(id)) {
                    providers.add((EnergyProviderBlockEntity) entity);
                } else {
                    iterator.remove();
                    dirty = true;
                }
                valid.add(pos);
            } else {
                iterator.remove();
                dirty = true;
            }
        }
        if (dirty) {
            markDirty();
        }
        Collections.shuffle(providers, world.random);

        cache.put(id, providers);

        world.getProfiler().pop();
        return providers;
    }

    public List<EnergyProviderBlockEntity> getSources(World world, int id) {
        world.getProfiler().push("getSources");

        List<EnergyProviderBlockEntity> result = getSourcesFromCache(world, id);

        if (result == null) {
            result = getSourcesFromCacheMiss(world, id);
        }

        world.getProfiler().pop();
        return result;
    }

    public void clearCache() {
        cache.clear();
    }

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

    public void removeSource(int id, BlockPosWithDimension 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<BlockPosWithDimension> 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();
    }
}