hytale-gravity-flip/src/main/java/com/mythlane/gravityflip/region/RegionRegistry.java

package com.mythlane.gravityflip.region;

import com.hypixel.hytale.component.ArchetypeChunk;
import com.hypixel.hytale.component.ComponentType;
import com.hypixel.hytale.component.Ref;
import com.hypixel.hytale.component.Store;
import com.hypixel.hytale.math.shape.Box;
import com.hypixel.hytale.server.core.modules.entity.component.TransformComponent;
import com.hypixel.hytale.server.core.universe.world.World;
import com.hypixel.hytale.server.core.universe.world.storage.EntityStore;
import com.hypixel.hytale.server.core.util.Config;
import com.mythlane.gravityflip.config.GravityFlipConfig;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicReference;

/**
 * In-memory index of {@link GravityFlipRegion}s with atomic publication for lock-free reads
 * from the tick loop. Mutations go through CRUD methods which swap an immutable snapshot.
 */
public final class RegionRegistry {

    // Lazy init: TransformComponent.getComponentType() triggers Hytale PluginBase static init,
    // which fails under JUL unless the log manager system property is set — avoided in tests.
    private static volatile ComponentType<EntityStore, TransformComponent> transformType;

    private static ComponentType<EntityStore, TransformComponent> transform() {
        ComponentType<EntityStore, TransformComponent> t = transformType;
        if (t == null) {
            synchronized (RegionRegistry.class) {
                t = transformType;
                if (t == null) {
                    t = TransformComponent.getComponentType();
                    transformType = t;
                }
            }
        }
        return t;
    }

    private final GravityFlipConfig config;
    private final Config<GravityFlipConfig> holder;   // nullable in tests

    private final AtomicReference<List<GravityFlipRegion>> regionsSnapshot;

    private final Object mutationLock = new Object();
    // Keyed by Object (not World) because Mockito cannot mock World under JDK 25 in tests.
    private final ConcurrentHashMap<Object, AtomicReference<RegionSnapshot>> snapshots = new ConcurrentHashMap<>();

    public RegionRegistry(GravityFlipConfig cfg) {
        this(cfg, null);
    }

    public RegionRegistry(GravityFlipConfig cfg, Config<GravityFlipConfig> holder) {
        this.config = cfg;
        this.holder = holder;
        this.regionsSnapshot = new AtomicReference<>(List.copyOf(cfg.getRegions()));
    }

    /** Returns the current immutable region-list snapshot. */
    public Collection<GravityFlipRegion> all() {
        return regionsSnapshot.get();
    }

    /** Read-only view of currently-enabled regions. */
    List<GravityFlipRegion> enabled() {
        List<GravityFlipRegion> out = new ArrayList<>();
        for (GravityFlipRegion r : regionsSnapshot.get()) {
            if (r.isEnabled()) out.add(r);
        }
        return out;
    }

    /** Returns the region with the given name from the current snapshot, or null. */
    public GravityFlipRegion find(String name) {
        if (name == null) return null;
        for (GravityFlipRegion x : regionsSnapshot.get()) {
            if (x.getName().equals(name)) return x;
        }
        return null;
    }

    /** Adds a region; throws IllegalArgumentException if the name already exists. */
    public void add(GravityFlipRegion r) {
        synchronized (mutationLock) {
            if (find(r.getName()) != null) {
                throw new IllegalArgumentException("region name already exists: " + r.getName());
            }
            config.getRegions().add(r);
            regionsSnapshot.set(List.copyOf(config.getRegions()));
        }
    }

    /** Removes the region with the given name; returns true iff it existed. */
    public boolean remove(String name) {
        synchronized (mutationLock) {
            boolean removed = config.getRegions().removeIf(x -> x.getName().equals(name));
            if (removed) regionsSnapshot.set(List.copyOf(config.getRegions()));
            return removed;
        }
    }

    /** Toggles the enabled flag on the named region; returns true iff it existed. */
    public boolean setEnabled(String name, boolean enabled) {
        synchronized (mutationLock) {
            for (GravityFlipRegion x : config.getRegions()) {
                if (x.getName().equals(name)) {
                    x.setEnabled(enabled);
                    regionsSnapshot.set(List.copyOf(config.getRegions()));
                    return true;
                }
            }
            return false;
        }
    }

    /** Re-snapshots the region list after the underlying config was mutated externally. */
    public void refreshFromConfig(GravityFlipConfig cfg) {
        synchronized (mutationLock) {
            regionsSnapshot.set(List.copyOf(cfg.getRegions()));
        }
    }

    /** Persists via the bound {@code Config} holder (completed future if none is bound). */
    public CompletableFuture<Void> save() {
        return holder == null ? CompletableFuture.completedFuture(null) : holder.save();
    }

    /** Iterates the ECS for {@code world} and publishes a fresh per-region occupancy snapshot. */
    public void refreshFor(World world) {
        List<GravityFlipRegion> enabled = enabled();
        Map<GravityFlipRegion, Collection<Ref<EntityStore>>> byRegion = new ConcurrentHashMap<>();
        for (GravityFlipRegion r : enabled) {
            byRegion.put(r, new ConcurrentLinkedQueue<>());
        }

        if (enabled.isEmpty()) {
            publishSnapshot(world, snapshotOf(world, byRegion));
            return;
        }

        // Store.forEachEntityParallel requires the WorldThread (assertThread), so dispatch via
        // world.execute(Runnable). Publication becomes asynchronous (<= 1 tick of lag) — tolerable
        // because RegionTickLoop runs @100ms, so snapshot freshness stays <= 100ms in the worst case.
        world.execute(() -> {
            try {
                Store<EntityStore> store = world.getEntityStore().getStore();
                ComponentType<EntityStore, TransformComponent> TRANSFORM = transform();
                // ComponentType IS-A Query, so TRANSFORM is passed directly.
                store.forEachEntityParallel(TRANSFORM, (index, chunk, cmdBuf) -> {
                    TransformComponent t = chunk.getComponent(index, TRANSFORM);
                    com.hypixel.hytale.math.vector.Vector3d pos = t.getPosition();
                    // Copy to locals — getPosition() returns a backing field; never mutated here.
                    double x = pos.x, y = pos.y, z = pos.z;
                    Ref<EntityStore> ref = null;
                    for (GravityFlipRegion r : enabled) {
                        Box box = r.asBox();
                        if (box.containsPosition(x, y, z)) {
                            if (ref == null) ref = chunk.getReferenceTo(index);
                            byRegion.get(r).add(ref);
                        }
                    }
                });
            } catch (Throwable th) {
                // Swallow — keeps the scheduler alive across transient ECS-state errors
                // (e.g., world being torn down). Publish whatever we collected so far.
            }

            publishSnapshot(world, snapshotOf(world, byRegion));
        });
    }

    /** Off-thread consumer entry point; returns null if no snapshot has been published yet. */
    public RegionSnapshot currentSnapshot(World world) {
        if (world == null) return null;
        AtomicReference<RegionSnapshot> ref = snapshots.get(world);
        return ref == null ? null : ref.get();
    }

    void publishSnapshot(World world, RegionSnapshot snap) {
        publishSnapshotByKey(world, snap);
    }

    void publishSnapshotByKey(Object key, RegionSnapshot snap) {
        snapshots.computeIfAbsent(key, k -> new AtomicReference<>()).set(snap);
    }

    RegionSnapshot currentSnapshotByKey(Object key) {
        AtomicReference<RegionSnapshot> ref = snapshots.get(key);
        return ref == null ? null : ref.get();
    }

    private RegionSnapshot snapshotOf(World w,
                                      Map<GravityFlipRegion, Collection<Ref<EntityStore>>> byRegion) {
        long tick = 0L;
        try { tick = w.getTick(); } catch (Throwable ignored) {}
        final long tickId = Math.max(tick, 1L);
        final Map<GravityFlipRegion, Collection<Ref<EntityStore>>> frozen =
                Collections.unmodifiableMap(byRegion);
        return new RegionSnapshot() {
            @Override public Map<GravityFlipRegion, Collection<Ref<EntityStore>>> byRegion() { return frozen; }
            @Override public long tickId() { return tickId; }
            @Override public World world() { return w; }
        };
    }
}