hytale-gravity-flip/src/main/java/com/mythlane/gravityflip/physics/GravityApplier.java

package com.mythlane.gravityflip.physics;

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.server.core.entity.UUIDComponent;
import com.hypixel.hytale.server.core.entity.entities.Player;
import com.hypixel.hytale.server.core.entity.entities.player.movement.MovementManager;
import com.hypixel.hytale.server.core.io.PacketHandler;
import com.hypixel.hytale.server.core.modules.entity.component.TransformComponent;
import com.hypixel.hytale.server.core.modules.physics.component.PhysicsValues;
import com.hypixel.hytale.server.core.universe.PlayerRef;
import com.hypixel.hytale.server.core.universe.world.World;
import com.hypixel.hytale.server.core.universe.world.storage.EntityStore;
import com.hypixel.hytale.server.npc.entities.NPCEntity;
import com.hypixel.hytale.server.npc.movement.controllers.MotionController;
import com.hypixel.hytale.server.npc.role.Role;
import com.mythlane.gravityflip.region.GravityFlipRegion;
import com.mythlane.gravityflip.region.RegionSnapshot;

import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * Tick-driven service that toggles {@code PhysicsValues.invertedGravity} on every entity currently
 * inside an enabled region, wakes up players/NPCs so the new settings take effect, and drives the
 * fall-damage guard on entry/exit transitions.
 */
public final class GravityApplier {
    // Lazy ComponentType holders — avoids Hytale PluginBase static init during tests.
    private static volatile ComponentType<EntityStore, PhysicsValues> physicsType;
    private static volatile ComponentType<EntityStore, UUIDComponent> uuidType;
    private static volatile ComponentType<EntityStore, TransformComponent> transformType;
    private static volatile ComponentType<EntityStore, MovementManager> movementManagerType;
    private static volatile ComponentType<EntityStore, PlayerRef>       playerRefType;
    private static volatile ComponentType<EntityStore, Player>          playerType;
    private static volatile ComponentType<EntityStore, NPCEntity>       npcEntityType;

    private static ComponentType<EntityStore, PhysicsValues> physicsType() {
        ComponentType<EntityStore, PhysicsValues> t = physicsType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = physicsType;
                if (t == null) { t = PhysicsValues.getComponentType(); physicsType = t; }
            }
        }
        return t;
    }
    private static ComponentType<EntityStore, UUIDComponent> uuidType() {
        ComponentType<EntityStore, UUIDComponent> t = uuidType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = uuidType;
                if (t == null) { t = UUIDComponent.getComponentType(); uuidType = t; }
            }
        }
        return t;
    }
    private static ComponentType<EntityStore, TransformComponent> transformType() {
        ComponentType<EntityStore, TransformComponent> t = transformType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = transformType;
                if (t == null) { t = TransformComponent.getComponentType(); transformType = t; }
            }
        }
        return t;
    }
    private static ComponentType<EntityStore, MovementManager> movementManagerType() {
        ComponentType<EntityStore, MovementManager> t = movementManagerType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = movementManagerType;
                if (t == null) { t = MovementManager.getComponentType(); movementManagerType = t; }
            }
        }
        return t;
    }
    private static ComponentType<EntityStore, PlayerRef> playerRefType() {
        ComponentType<EntityStore, PlayerRef> t = playerRefType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = playerRefType;
                if (t == null) { t = PlayerRef.getComponentType(); playerRefType = t; }
            }
        }
        return t;
    }
    private static ComponentType<EntityStore, Player> playerType() {
        ComponentType<EntityStore, Player> t = playerType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = playerType;
                if (t == null) { t = Player.getComponentType(); playerType = t; }
            }
        }
        return t;
    }
    private static ComponentType<EntityStore, NPCEntity> npcEntityType() {
        ComponentType<EntityStore, NPCEntity> t = npcEntityType;
        if (t == null) {
            synchronized (GravityApplier.class) {
                t = npcEntityType;
                if (t == null) { t = NPCEntity.getComponentType(); npcEntityType = t; }
            }
        }
        return t;
    }

    // Written/read from ECS worker threads via forEachEntityParallel — concurrent collection required.
    private final Set<UUID> previouslyInverted = ConcurrentHashMap.newKeySet();
    /** First-matched region per UUID at the previous tick — consulted in Pass 2 for markExit. */
    private final ConcurrentHashMap<UUID, GravityFlipRegion> lastKnownRegion = new ConcurrentHashMap<>();

    private final Consumer<Throwable> errorHandler;
    private final FallDamageGuard guard;

    public GravityApplier(Consumer<Throwable> errorHandler) {
        this(errorHandler, new FallDamageGuard());
    }

    public GravityApplier(Consumer<Throwable> errorHandler, FallDamageGuard guard) {
        this.errorHandler = errorHandler == null ? t -> {} : errorHandler;
        this.guard = guard == null ? new FallDamageGuard() : guard;
    }

    /** Builds a new PhysicsValues copying mass/drag from the source and setting invertedGravity. */
    static PhysicsValues buildPhysicsValuesWithFlag(PhysicsValues source, boolean target) {
        FlaggedDecision d = buildFlaggedDecision(source.getMass(), source.getDragCoefficient(), target);
        return new PhysicsValues(d.mass, d.drag, d.invertedGravity);
    }

    /** Pure-data seam for unit tests — PhysicsValues static init is unavailable outside the server runtime. */
    static FlaggedDecision buildFlaggedDecision(double mass, double drag, boolean target) {
        return new FlaggedDecision(mass, drag, target);
    }

    static final class FlaggedDecision {
        final double mass;
        final double drag;
        final boolean invertedGravity;
        FlaggedDecision(double mass, double drag, boolean invertedGravity) {
            this.mass = mass; this.drag = drag; this.invertedGravity = invertedGravity;
        }
    }

    /** Tick entry point; no-op when world or snapshot is null. */
    public void apply(World world, RegionSnapshot snapshot) {
        if (world == null || snapshot == null) return;
        world.execute(() -> applyOnWorldThread(world, snapshot));
    }

    private void applyOnWorldThread(World world, RegionSnapshot snapshot) {
        Collection<GravityFlipRegion> enabledRegions = snapshot.byRegion().keySet();

        try {
            Store<EntityStore> store = world.getEntityStore().getStore();
            ComponentType<EntityStore, PhysicsValues> PHYST = physicsType();
            ComponentType<EntityStore, UUIDComponent>  UUIDT = uuidType();
            ComponentType<EntityStore, TransformComponent> TT = transformType();
            ComponentType<EntityStore, MovementManager> MMT = movementManagerType();
            ComponentType<EntityStore, PlayerRef>       PRT = playerRefType();
            ComponentType<EntityStore, Player>          PLT = playerType();
            ComponentType<EntityStore, NPCEntity>       NPCT = npcEntityType();

            // PASS 1 — flip ON for every entity with PhysicsValues that is inside an enabled region.
            // Pass 1 and Pass 2 cannot be fused: restore requires the complete currentlyInRegion set
            // to diff against previouslyInverted.
            Set<UUID> currentlyInRegion = ConcurrentHashMap.newKeySet();
            store.forEachEntityParallel(PHYST, (index, chunk, cmdBuf) -> {
                TransformComponent t;
                UUIDComponent uc;
                PhysicsValues v;
                try {
                    t = chunk.getComponent(index, TT);
                    uc = chunk.getComponent(index, UUIDT);
                    v  = chunk.getComponent(index, PHYST);
                } catch (Throwable ignored) { return; }
                if (t == null || uc == null || v == null) return;

                com.hypixel.hytale.math.vector.Vector3d pos = t.getPosition();
                double x = pos.x, y = pos.y, z = pos.z;

                // First-match wins for multi-region precedence.
                GravityFlipRegion matchedRegion = null;
                for (GravityFlipRegion r : enabledRegions) {
                    if (r.asBox().containsPosition(x, y, z)) { matchedRegion = r; break; }
                }
                if (matchedRegion == null) return;

                UUID u = uc.getUuid();

                EntityKind kind = classify(chunk, index, MMT, PRT, PLT, NPCT);
                boolean allowed;
                switch (kind) {
                    case PLAYER: allowed = matchedRegion.isAffectPlayers(); break;
                    case NPC:    allowed = matchedRegion.isAffectNpcs();    break;
                    default:     allowed = matchedRegion.isAffectItems();   break;
                }
                if (!allowed) {
                    // Filtered entity behaves as if outside the region.
                    return;
                }

                currentlyInRegion.add(u);
                lastKnownRegion.put(u, matchedRegion);
                guard.markInRegion(u, matchedRegion);

                if (!v.isInvertedGravity()) {
                    Ref<EntityStore> ref = chunk.getReferenceTo(index);
                    cmdBuf.replaceComponent(ref, PHYST,
                            new PhysicsValues(v.getMass(), v.getDragCoefficient(), true));
                }

                wakePlayerOrNpc(chunk, index, v, true, matchedRegion,
                        MMT, PRT, PLT, NPCT);
            });

            // PASS 2 — restore: for every UUID in previouslyInverted \ currentlyInRegion.
            Set<UUID> toRestore = ConcurrentHashMap.newKeySet();
            toRestore.addAll(previouslyInverted);
            toRestore.removeAll(currentlyInRegion);
            if (!toRestore.isEmpty()) {
                long now = System.currentTimeMillis();
                store.forEachEntityParallel(PHYST, (index, chunk, cmdBuf) -> {
                    UUIDComponent uc;
                    PhysicsValues v;
                    try {
                        uc = chunk.getComponent(index, UUIDT);
                        v  = chunk.getComponent(index, PHYST);
                    } catch (Throwable ignored) { return; }
                    if (uc == null || v == null) return;
                    UUID u = uc.getUuid();
                    if (!toRestore.contains(u)) return;

                    if (v.isInvertedGravity()) {
                        Ref<EntityStore> ref = chunk.getReferenceTo(index);
                        cmdBuf.replaceComponent(ref, PHYST,
                                new PhysicsValues(v.getMass(), v.getDragCoefficient(), false));
                    }

                    wakePlayerOrNpc(chunk, index, v, false, null, MMT, PRT, PLT, NPCT);

                    GravityFlipRegion lastRegion = lastKnownRegion.remove(u);
                    if (lastRegion != null) {
                        guard.markExit(u, lastRegion, now);
                    }
                });
            }

            // Tick-thread update after the parallel pass completes.
            previouslyInverted.clear();
            previouslyInverted.addAll(currentlyInRegion);
        } catch (Throwable th) {
            errorHandler.accept(th);
        }
    }

    private enum EntityKind { PLAYER, NPC, OTHER }

    /** Classifies the entity at {@code index} into player / NPC / other (items fall into other). */
    private EntityKind classify(ArchetypeChunk<EntityStore> chunk, int index,
                                ComponentType<EntityStore, MovementManager> MMT,
                                ComponentType<EntityStore, PlayerRef>       PRT,
                                ComponentType<EntityStore, Player>          PLT,
                                ComponentType<EntityStore, NPCEntity>       NPCT) {
        try {
            MovementManager mm = chunk.getComponent(index, MMT);
            if (mm != null) {
                PlayerRef pr = null; Player pl = null;
                try { pr = chunk.getComponent(index, PRT); } catch (Throwable ignored) {}
                try { pl = chunk.getComponent(index, PLT); } catch (Throwable ignored) {}
                if (pr != null && pl != null) return EntityKind.PLAYER;
            }
        } catch (Throwable ignored) {}
        try {
            NPCEntity npc = chunk.getComponent(index, NPCT);
            if (npc != null) return EntityKind.NPC;
        } catch (Throwable ignored) {}
        return EntityKind.OTHER;
    }

    /** Wakes up the entity so the new PhysicsValues take effect (player movement refresh or NPC controller update). */
    private void wakePlayerOrNpc(
            ArchetypeChunk<EntityStore> chunk, int index,
            PhysicsValues sourceValues, boolean targetFlag,
            GravityFlipRegion matchedRegion,
            ComponentType<EntityStore, MovementManager> MMT,
            ComponentType<EntityStore, PlayerRef>       PRT,
            ComponentType<EntityStore, Player>          PLT,
            ComponentType<EntityStore, NPCEntity>       NPCT) {
        PhysicsValues targetValues = buildPhysicsValuesWithFlag(sourceValues, targetFlag);

        // Player branch.
        MovementManager mm = null;
        try { mm = chunk.getComponent(index, MMT); } catch (Throwable ignored) {}
        if (mm != null) {
            PlayerRef pr = null; Player pl = null;
            try { pr = chunk.getComponent(index, PRT); } catch (Throwable ignored) {}
            try { pl = chunk.getComponent(index, PLT); } catch (Throwable ignored) {}
            if (pr != null && pl != null) {
                try {
                    mm.setDefaultSettings(mm.getDefaultSettings(), targetValues, pl.getGameMode());
                    mm.applyDefaultSettings();
                    PacketHandler ph = pr.getPacketHandler();
                    mm.update(ph);
                } catch (Throwable th) {
                    errorHandler.accept(th);
                }
                return;
            }
        }

        // NPC branch.
        NPCEntity npc = null;
        try { npc = chunk.getComponent(index, NPCT); } catch (Throwable ignored) {}
        if (npc != null) {
            try {
                Role role = npc.getRole();
                if (role != null) {
                    MotionController active = role.getActiveMotionController();
                    if (active != null) {
                        active.updatePhysicsValues(targetValues);

                        // Seed forceVelocity.y only on entry — on exit the native damping zeroes it.
                        if (targetFlag && matchedRegion != null) {
                            double vf = matchedRegion.getVerticalForce();
                            try {
                                active.addForce(
                                    new com.hypixel.hytale.math.vector.Vector3d(0, vf, 0),
                                    null);
                            } catch (Throwable th) {
                                errorHandler.accept(th);
                            }
                        }
                    }
                }
            } catch (Throwable th) {
                errorHandler.accept(th);
            }
            return;
        }

        // Items / other: no-op.
    }

    /** Pure data-diff helper for unit tests. */
    public static DiffResult diff(Set<UUID> previous, Set<UUID> current) {
        Set<UUID> toFlip = new HashSet<>(current);
        toFlip.removeAll(previous);
        Set<UUID> toRestore = new HashSet<>(previous);
        toRestore.removeAll(current);
        return new DiffResult(toFlip, toRestore);
    }

    public static final class DiffResult {
        public final Set<UUID> toFlip;
        public final Set<UUID> toRestore;
        DiffResult(Set<UUID> f, Set<UUID> r) { this.toFlip = f; this.toRestore = r; }
    }

    // Test hooks (package-private).

    /** Immutable view of the internal tracker for unit tests. */
    Set<UUID> previouslyInvertedView() {
        return Collections.unmodifiableSet(previouslyInverted);
    }

    /** Test-only tracker override; never call from production code. */
    void __updateTrackerForTest(Set<UUID> newState) {
        previouslyInverted.clear();
        previouslyInverted.addAll(newState);
    }
}