public class DirectLightMode implements LightingMode {

    @Override
    public void applyLighting(Block block, int power) {
        applyLighting(Set.of(block), power);
    }

    @Override
    public void removeLighting(Block block) {
        removeLighting(Set.of(block));
    }

    @Override
    public void removeLighting(Set<Block> blocks) {
        Map<World, Set<ChunkPos>> worldToChunksMap = new HashMap<>();

        for (Block block : blocks) {
            World world = block.getWorld();
            ChunkPos chunkPos = new ChunkPos(block.getX() >> 4, block.getZ() >> 4);
            worldToChunksMap.computeIfAbsent(world, w -> new HashSet<>()).add(chunkPos);
        }

        for (Map.Entry<World, Set<ChunkPos>> entry : worldToChunksMap.entrySet()) {
            ServerLevel nmsLevel = ((CraftWorld) entry.getKey()).getHandle();

            // Refresh chunks with light sources
            for (ChunkPos chunkPos : entry.getValue()) {
                removeLighting(chunkPos, entry.getKey());
            }

            // Also refresh neighboring chunks to fix border artifacts
            Set<ChunkPos> neighborChunks = new HashSet<>();
            for (ChunkPos chunkPos : entry.getValue()) {
                for (int dx = -1; dx <= 1; dx++) {
                    for (int dz = -1; dz <= 1; dz++) {
                        if (dx == 0 && dz == 0) continue;
                        neighborChunks.add(new ChunkPos(chunkPos.x + dx, chunkPos.z + dz));
                    }
                }
            }

            for (ChunkPos neighborPos : neighborChunks) {
                refreshChunkLighting(neighborPos, nmsLevel);
            }
        }
    }

    @Override
    public void applyLighting(Set<Block> blocks, int power) {
        Map<ChunkPos, Set<Block>> chunkToBlocksMap = new HashMap<>();

        for (Block block : blocks) {
            ChunkPos chunkPos = new ChunkPos(block.getX() >> 4, block.getZ() >> 4);
            chunkToBlocksMap.computeIfAbsent(chunkPos, c -> new HashSet<>()).add(block);
        }

        // Apply lighting to chunks with light sources
        chunkToBlocksMap.forEach((chunkPos, chunkBlocks) -> applyChunkLighting(chunkPos, chunkBlocks, power));

        // Update neighboring chunks to fix border artifacts
        Set<ChunkPos> neighborChunks = new HashSet<>();
        for (ChunkPos chunkPos : chunkToBlocksMap.keySet()) {
            // Add all 8 neighboring chunks
            for (int dx = -1; dx <= 1; dx++) {
                for (int dz = -1; dz <= 1; dz++) {
                    if (dx == 0 && dz == 0) continue; // Skip the center chunk (already updated)
                    neighborChunks.add(new ChunkPos(chunkPos.x + dx, chunkPos.z + dz));
                }
            }
        }

        neighborChunks.removeIf(chunkToBlocksMap::containsKey);

        // Send light updates to neighboring chunks (without modifying their data)
        World world = blocks.iterator().next().getWorld();
        ServerLevel nmsLevel = ((CraftWorld) world).getHandle();
        for (ChunkPos neighborPos : neighborChunks) {
            refreshChunkLighting(neighborPos, nmsLevel);
        }
    }

    @Override
    public void updateLighting(Set<Block> toUpdate, int power, Set<Block> toRemove) {
        Map<ChunkPos, Set<Block>> updateMap = new HashMap<>();
        Set<ChunkPos> removeChunks = new HashSet<>();

        Set<Block> allBlocks = new HashSet<>(toUpdate);
        allBlocks.addAll(toRemove);

        if (allBlocks.isEmpty()) {
            return;
        }

        World world = allBlocks.iterator().next().getWorld();
        ServerLevel nmsLevel = ((CraftWorld) world).getHandle();

        for (Block block : toRemove) {
            ChunkPos chunkPos = new ChunkPos(block.getX() >> 4, block.getZ() >> 4);
            removeChunks.add(chunkPos);
        }

        for (Block block : toUpdate) {
            ChunkPos chunkPos = new ChunkPos(block.getX() >> 4, block.getZ() >> 4);
            updateMap.computeIfAbsent(chunkPos, c -> new HashSet<>()).add(block);
        }

        Set<ChunkPos> removedOrNeighbouringChunks = new HashSet<>(removeChunks);

        for (ChunkPos chunkPos : updateMap.keySet()) {
            removedOrNeighbouringChunks.addAll(getNeighbouringChunks(chunkPos));
        }

        for (ChunkPos chunkPos : updateMap.keySet()) {
            removedOrNeighbouringChunks.addAll(getNeighbouringChunks(chunkPos));
        }

        removedOrNeighbouringChunks.removeIf(updateMap::containsKey);

        for (ChunkPos chunkPos : removedOrNeighbouringChunks) {
            refreshChunkLighting(chunkPos, nmsLevel);
        }

        for (Map.Entry<ChunkPos, Set<Block>> entry : updateMap.entrySet()) {
            applyChunkLighting(entry.getKey(), entry.getValue(), power);
        }
    }

    private void applyChunkLighting(ChunkPos chunkPos, Set<Block> blocks, int power) {
        ServerLevel nmsLevel = ((CraftWorld) blocks.iterator().next().getWorld()).getHandle();
        LevelChunk chunk = nmsLevel.getChunk(chunkPos.x, chunkPos.z);

        LevelLightEngine lightEngine = nmsLevel.getLightEngine();
        ClientboundLightUpdatePacket packet = new ClientboundLightUpdatePacket(chunk.getPos(), lightEngine, null, new BitSet());

        ClientboundLightUpdatePacketData data = packet.getLightData();

        List<byte[]> blockUpdates = new ArrayList<>();
        List<byte[]> skyUpdates = new ArrayList<>();

        data.getEmptyBlockYMask().clear();
        data.getBlockYMask().clear();
        data.getEmptySkyYMask().clear();
        data.getSkyYMask().clear();

        for (int i = 0; i < lightEngine.getLightSectionCount(); ++i) {
            this.prepareSectionData(chunk.getPos(), lightEngine, LightLayer.BLOCK, i, data.getBlockYMask(), data.getEmptyBlockYMask(), blockUpdates, layer -> {
                for (Block block : blocks) {
                    layer.set(block.getX() & 15, block.getY() & 15, block.getZ() & 15, power);
                }
            });

            this.prepareSectionData(chunk.getPos(), lightEngine, LightLayer.SKY, i, data.getSkyYMask(), data.getEmptySkyYMask(), skyUpdates, layer -> {
                for (Block block : blocks) {
                    layer.set(block.getX() & 15, block.getY() & 15, block.getZ() & 15, power);
                }
            });
        }

        data.getBlockUpdates().clear();
        data.getBlockUpdates().addAll(blockUpdates);

        data.getSkyUpdates().clear();
        data.getSkyUpdates().addAll(skyUpdates);

        for (Player player : Bukkit.getOnlinePlayers()) {
            ((CraftPlayer) player).getHandle().connection.send(packet);
        }
    }

    private void removeLighting(ChunkPos pos, World world) {
        ServerLevel nmsLevel = ((CraftWorld) world).getHandle();
        refreshChunkLighting(pos, nmsLevel);
    }

    private void refreshChunkLighting(ChunkPos pos, ServerLevel nmsLevel) {
        LevelChunk chunk = nmsLevel.getChunk(pos.x, pos.z);

        LevelLightEngine lightEngine = nmsLevel.getLightEngine();
        ClientboundLightUpdatePacket packet = new ClientboundLightUpdatePacket(chunk.getPos(), lightEngine, null, null);

        for (Player player : Bukkit.getOnlinePlayers()) {
            ((CraftPlayer) player).getHandle().connection.send(packet);
        }
    }

    private void prepareSectionData(ChunkPos chunkPos, LevelLightEngine levelLightEngine, LightLayer lightLayer, int index, BitSet skyLight, BitSet blockLight, List<byte[]> updates, Consumer<DataLayer> dataLayerConsumer) {
        DataLayer dataLayerData = levelLightEngine.getLayerListener(lightLayer).getDataLayerData(SectionPos.of(chunkPos, levelLightEngine.getMinLightSection() + index));

        if (dataLayerData == null) {
            return;
        }

        dataLayerData = dataLayerData.copy();
        dataLayerConsumer.accept(dataLayerData);

        if (dataLayerData.isEmpty()) {
            blockLight.set(index);
        } else {
            skyLight.set(index);
            updates.add(dataLayerData.copy().getData());
        }
    }

    private Set<ChunkPos> getNeighbouringChunks(ChunkPos chunkPos) {
        Set<ChunkPos> neighbors = new HashSet<>();
        for (int dx = -1; dx <= 1; dx++) {
            for (int dz = -1; dz <= 1; dz++) {
                if (dx == 0 && dz == 0) continue;
                neighbors.add(new ChunkPos(chunkPos.x + dx, chunkPos.z + dz));
            }
        }
        return neighbors;
    }
}