package me.earthme.luminol.entitylimiter;

import me.earthme.luminol.config.modules.optimizations.EntityLimiterConfig;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.Mob;
import net.minecraft.world.entity.item.ItemEntity;
import net.minecraft.world.entity.projectile.Projectile;
import net.minecraft.world.entity.vehicle.AbstractMinecart;
import net.minecraft.world.entity.vehicle.Boat;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.level.chunk.LevelChunk;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

/**
 * Neko Core built-in entity limiter (ported from LagFixer EntityLimiterModule, GPL-3.0).
 *
 * <p>Two independent enforcement modes:
 * <ul>
 *   <li><b>Spawn-time gate</b> ({@link #shouldBlock}) – called from
 *       {@code ServerLevel.addEntity}; cancels the spawn when the target chunk
 *       already holds {@code >= limit} entities of the same category.</li>
 *   <li><b>Overflow sweep</b> ({@link #tickLevel}) – driven from the regionised
 *       {@code ServerLevel.tick}; every {@code overflowIntervalSeconds} it scans the
 *       current region's entities, buckets them per chunk, and discards anything
 *       above {@code limit × overflowMultiplier}.</li>
 * </ul>
 *
 * <p>Folia-safe: the sweep only touches entities owned by the ticking region
 * (via {@code getLocalEntitiesCopy()}), so removals always run on the correct thread.
 * No background executor is used, which avoids the upstream chunk-map iterator drift.
 */
public final class LuminolEntityLimiter {

    private static final Logger LOGGER = LoggerFactory.getLogger("NekoCore/EntityLimiter");
    private static volatile boolean started = false;
    /** Per-level tick accumulator, keyed by level dimension id. Sweep fires when it reaches the interval. */
    private static final Map<String, Integer> TICK_ACCUM = new HashMap<>();

    private LuminolEntityLimiter() {}

    /** Logs the active configuration once at startup. Idempotent. */
    public static synchronized void start() {
        if (started || !EntityLimiterConfig.enabled) return;
        started = true;
        LOGGER.info("EntityLimiter started (creatures={}, items={}, vehicles={}, projectiles={}, overflow={}/{}s)",
            EntityLimiterConfig.creatures, EntityLimiterConfig.items, EntityLimiterConfig.vehicles,
            EntityLimiterConfig.projectiles, EntityLimiterConfig.overflowEnabled, EntityLimiterConfig.overflowIntervalSeconds);
    }

    public static synchronized void stop() {
        started = false;
        TICK_ACCUM.clear();
    }

    // ── Spawn-time gate ─────────────────────────────────────────────────────────
    //
    // Per-chunk, per-region-tick classification count cache.
    //
    // The naive implementation scanned the whole region's entity list (an O(region) copy) on every
    // single spawn, making a wartime burst — TNT dropping thousands of items, a horde spawning — an
    // O(N²) (worse, given the per-call array copy) catastrophe: the limiter meant to *prevent* lag
    // became the lag source. This cache rebuilds a chunk's category counts at most once per region
    // tick (scanning only that one chunk's entity slice, O(chunk)), then increments locally as more
    // entities of the same chunk spawn in the same tick. A burst of M spawns into one chunk costs
    // O(K + M) instead of O(M · region).
    //
    // Folia-safe: the spawn gate always runs on the region thread that owns the chunk, so a
    // ThreadLocal cache is region-isolated. The generation stamp (bumped each region tick in
    // {@link #tickLevel}) invalidates stale entries lazily.

    /** counts[0..3] = creatures/items/vehicles/projectiles for a chunk this tick. */
    private static final ThreadLocal<it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap<int[]>> CHUNK_COUNTS =
            ThreadLocal.withInitial(it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap::new);
    /** Per-chunk generation stamp the cached counts were built at (parallel to CHUNK_COUNTS). */
    private static final ThreadLocal<it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap> CHUNK_GEN =
            ThreadLocal.withInitial(() -> {
                it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap m = new it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap();
                m.defaultReturnValue(Long.MIN_VALUE);
                return m;
            });
    /** Monotonic generation, bumped once per region tick so cached chunk counts expire. */
    private static final ThreadLocal<long[]> GENERATION = ThreadLocal.withInitial(() -> new long[]{0L});

    /** Returns true if {@code entity} should be blocked because its chunk is at the per-category cap. */
    public static boolean shouldBlock(Entity entity, LevelChunk chunk) {
        if (!EntityLimiterConfig.enabled || isWhitelisted(entity)) return false;
        int limit = getLimit(entity);
        if (limit < 1) return false;
        int cat = categoryIndex(entity);
        if (cat < 0) return false;

        final int[] counts = chunkCounts(chunk);
        if (counts == null) return false; // off-region; do not block

        if (counts[cat] >= limit) {
            return true; // at cap — block, and do NOT count the rejected entity
        }
        counts[cat]++; // accepted: keep the running total current for the next same-tick spawn
        return false;
    }

    /**
     * Returns this chunk's per-category counts for the current region tick, rebuilding from the
     * chunk's own entity slice on a cache miss / stale generation. {@code null} if off-region.
     */
    private static int[] chunkCounts(LevelChunk chunk) {
        io.papermc.paper.threadedregions.RegionizedWorldData worldData = chunk.getLevel().getCurrentWorldData();
        if (worldData == null) return null; // off-region (e.g. worldgen thread)

        final long gen = GENERATION.get()[0];
        final long key = chunk.getPos().toLong();
        final it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap<int[]> map = CHUNK_COUNTS.get();
        final it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap genMap = CHUNK_GEN.get();
        int[] counts = map.get(key);
        if (counts != null && genMap.get(key) == gen) {
            return counts; // fresh this tick
        }
        if (counts == null) {
            counts = new int[4];
            map.put(key, counts);
        } else {
            java.util.Arrays.fill(counts, 0);
        }
        genMap.put(key, gen);
        // One scan of this chunk's own entities (O(chunk), not O(region)).
        for (Entity e : localChunkEntities(chunk)) {
            if (e == null || e.isRemoved() || isWhitelisted(e)) continue;
            int c = categoryIndex(e);
            if (c >= 0) counts[c]++;
        }
        return counts;
    }

    private static java.util.List<Entity> localChunkEntities(LevelChunk chunk) {
        try {
            ca.spottedleaf.moonrise.patches.chunk_system.scheduling.NewChunkHolder holder =
                    ((ca.spottedleaf.moonrise.patches.chunk_system.level.ChunkSystemServerLevel) chunk.getLevel())
                            .moonrise$getChunkTaskScheduler().chunkHolderManager
                            .getChunkHolder(chunk.locX, chunk.locZ);
            if (holder == null) return java.util.Collections.emptyList();
            ca.spottedleaf.moonrise.patches.chunk_system.level.entity.ChunkEntitySlices slices = holder.getEntityChunk();
            if (slices == null) return java.util.Collections.emptyList();
            return slices.getAllEntities();
        } catch (Throwable t) {
            return java.util.Collections.emptyList();
        }
    }

    // ── Overflow sweep (driven from ServerLevel.tick on the region thread) ───────
    /**
     * Called once per region tick. Accumulates ticks and, once the configured
     * interval elapses, sweeps this region's entities for overflow.
     */
    public static void tickLevel(ServerLevel level) {
        if (!EntityLimiterConfig.enabled) return;
        // Bump the spawn-gate count-cache generation so this tick's first spawn into each chunk
        // rebuilds from a fresh scan. Cheap (one long increment per region tick).
        long g = ++GENERATION.get()[0];
        // The per-chunk cache maps grow as the region visits chunks and never self-shrink. Clear
        // them periodically (they rebuild lazily, so clearing only costs a re-scan on next spawn)
        // to bound memory on long-lived region threads.
        if ((g & 0x3FFL) == 0L) { // ~every 1024 region ticks (~51s)
            CHUNK_COUNTS.get().clear();
            CHUNK_GEN.get().clear();
        }
        if (!EntityLimiterConfig.overflowEnabled || EntityLimiterConfig.overflowIntervalSeconds <= 0) {
            return;
        }
        final int intervalTicks = EntityLimiterConfig.overflowIntervalSeconds * 20;
        final String key = level.dimension().location().toString();
        int accum = TICK_ACCUM.getOrDefault(key, 0) + 1;
        if (accum < intervalTicks) {
            TICK_ACCUM.put(key, accum);
            return;
        }
        TICK_ACCUM.put(key, 0);
        runOverflowSweep(level);
    }

    private static void runOverflowSweep(ServerLevel level) {
        io.papermc.paper.threadedregions.RegionizedWorldData worldData = level.getCurrentWorldData();
        if (worldData == null) return;
        final int limitC = (int)(EntityLimiterConfig.creatures   * EntityLimiterConfig.overflowMultiplier);
        final int limitI = (int)(EntityLimiterConfig.items       * EntityLimiterConfig.overflowMultiplier);
        final int limitV = (int)(EntityLimiterConfig.vehicles    * EntityLimiterConfig.overflowMultiplier);
        final int limitP = (int)(EntityLimiterConfig.projectiles * EntityLimiterConfig.overflowMultiplier);

        // Per-chunk running counts: index 0=creatures, 1=items, 2=vehicles, 3=projectiles.
        final Map<Long, int[]> perChunk = new HashMap<>();
        int removed = 0;

        // Cap removals per sweep so a sudden entity explosion (e.g. wartime TNT dropping thousands
        // of items at once) doesn't turn a single sweep into a multi-thousand-removal MSPT spike.
        // The remainder is cleaned up on the next sweep. 0 = unlimited.
        final int maxPerSweep = EntityLimiterConfig.overflowMaxPerSweep > 0
                ? EntityLimiterConfig.overflowMaxPerSweep : Integer.MAX_VALUE;
        boolean capped = false;

        // Snapshot so removals during the loop don't disturb iteration. Region-owned entities only.
        final Entity[] snapshot = worldData.getLocalEntitiesCopy();
        for (Entity entity : snapshot) {
            if (entity == null || entity.isRemoved()) continue;
            if (isWhitelisted(entity)) continue;
            if (EntityLimiterConfig.spareNamed && entity.hasCustomName()) continue;

            final ChunkPos cp = entity.chunkPosition();
            final int[] counts = perChunk.computeIfAbsent(cp.toLong(), k -> new int[4]);

            // Still count categories accurately (so per-chunk caps stay correct across sweeps),
            // but stop issuing removals once this sweep's budget is spent.
            if (entity instanceof Mob) {
                if (EntityLimiterConfig.overflowCreatures && ++counts[0] > limitC && removed < maxPerSweep) { entity.remove(Entity.RemovalReason.DISCARDED); removed++; }
            } else if (entity instanceof ItemEntity) {
                if (EntityLimiterConfig.overflowItems && ++counts[1] > limitI && removed < maxPerSweep) { entity.remove(Entity.RemovalReason.DISCARDED); removed++; }
            } else if (entity instanceof AbstractMinecart || entity instanceof Boat) {
                if (EntityLimiterConfig.overflowVehicles && ++counts[2] > limitV && removed < maxPerSweep) { entity.remove(Entity.RemovalReason.DISCARDED); removed++; }
            } else if (entity instanceof Projectile) {
                if (EntityLimiterConfig.overflowProjectiles && ++counts[3] > limitP && removed < maxPerSweep) { entity.remove(Entity.RemovalReason.DISCARDED); removed++; }
            }

            if (removed >= maxPerSweep) { capped = true; break; }
        }

        if (removed > 0) {
            LOGGER.debug("EntityLimiter overflow sweep removed {} entities in {}{}", removed,
                level.dimension().location(), capped ? " (capped, will continue next sweep)" : "");
        }
    }

    // ── Helpers ─────────────────────────────────────────────────────────────────
    private static int getLimit(Entity e) {
        if (e instanceof Mob) return EntityLimiterConfig.creatures;
        if (e instanceof ItemEntity) return EntityLimiterConfig.items;
        if (e instanceof AbstractMinecart || e instanceof Boat) return EntityLimiterConfig.vehicles;
        if (e instanceof Projectile) return EntityLimiterConfig.projectiles;
        return 0;
    }

    private static boolean isWhitelisted(Entity e) {
        String typeName = e.getType().toString().toUpperCase();
        return EntityLimiterConfig.whitelist.stream().anyMatch(typeName::contains);
    }

    /** Maps an entity to its limiter category index (0=creatures,1=items,2=vehicles,3=projectiles), or -1 if uncapped. */
    private static int categoryIndex(Entity e) {
        if (e instanceof Mob) return 0;
        if (e instanceof ItemEntity) return 1;
        if (e instanceof AbstractMinecart || e instanceof Boat) return 2;
        if (e instanceof Projectile) return 3;
        return -1;
    }
}
