FlowTask.java

package me.lukas.minecraft2.core;

import org.bukkit.Bukkit;
import org.bukkit.Material;
import org.bukkit.block.Block;
import org.bukkit.block.BlockFace;
import org.bukkit.block.data.Levelled;
import org.bukkit.plugin.Plugin;
import org.bukkit.scheduler.BukkitRunnable;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public class FlowTask extends BukkitRunnable {

  private final Plugin plugin;
  private final Set<Block> activeBlocks = ConcurrentHashMap.newKeySet();
  private static final BlockFace[] HORIZONTAL_FACES = {BlockFace.NORTH, BlockFace.EAST, BlockFace.SOUTH, BlockFace.WEST};
  private static final int PRESSURE_SEARCH_LIMIT = 96;

  public FlowTask(Plugin plugin) {
    this.plugin = plugin;
  }

  public void addActiveBlock(Block block) {
    if (block.getChunk().isLoaded() && block.getType() == Material.WATER) {
      activeBlocks.add(block);
    }
  }

  @Override
  public void run() {
    if (activeBlocks.isEmpty()) return;
    Set<Block> processingBlocks = new HashSet<>(activeBlocks);
    activeBlocks.clear();
    Map<Block, Integer> pendingChanges = new ConcurrentHashMap<>();
    for (Block sourceBlock : processingBlocks) {
      if (!sourceBlock.getChunk().isLoaded() || sourceBlock.getType() != Material.WATER) continue;
      if (flowDown(sourceBlock, pendingChanges)) continue;
      int sourceLevel = getWaterLevel(sourceBlock, pendingChanges);
      if (sourceLevel == 1 && drainPuddle(sourceBlock, pendingChanges)) continue;
      if (sourceLevel > 1) {
        equalizePressure(sourceBlock, pendingChanges);
      }
    }
    if (!pendingChanges.isEmpty()) {
      Bukkit.getScheduler().runTask(plugin, () -> applyChanges(pendingChanges));
    }
  }

  private void applyChanges(Map<Block, Integer> changes) {
    Set<Block> neighborsToWake = new HashSet<>();
    for (Map.Entry<Block, Integer> entry : changes.entrySet()) {
      Block block = entry.getKey();
      int newLevel = entry.getValue();
      if (block.getChunk().isLoaded() && getWaterLevelFromWorld(block) != newLevel) {
        setWaterLevelInWorld(block, newLevel);
        collectNeighborsToWake(block, neighborsToWake);
      }
    }
    for (Block neighbor : neighborsToWake) {
      if (neighbor.getChunk().isLoaded() && neighbor.getType() == Material.WATER) {
        addActiveBlock(neighbor);
      }
    }
  }

  private boolean equalizePressure(Block sourceBlock, Map<Block, Integer> changes) {
    int sourceLevel = getWaterLevel(sourceBlock, changes);
    Block bestTarget = findBestLocalTarget(sourceBlock, sourceLevel, changes);

    if (bestTarget != null) {
      int lowestLevel = getWaterLevel(bestTarget, changes);
      if (sourceLevel > lowestLevel) {
        applyDynamicAdditiveFlow(sourceBlock, bestTarget, changes);
        return true;
      }
    }
    return wideSearch(sourceBlock, sourceLevel, changes);
  }

  private Block findBestLocalTarget(Block sourceBlock, int sourceLevel, Map<Block, Integer> changes) {
    Block bestTarget = null;
    int lowestLevel = sourceLevel;
    for (BlockFace face : HORIZONTAL_FACES) {
      Block neighbor = sourceBlock.getRelative(face);
      if (!neighbor.getChunk().isLoaded()) continue;
      if (canFlowInto(neighbor, changes) || neighbor.getType() == Material.WATER) {
        int neighborLevel = getWaterLevel(neighbor, changes);
        if (neighborLevel < lowestLevel) {
          lowestLevel = neighborLevel;
          bestTarget = neighbor;
        }
      }
    }
    return bestTarget;
  }

  private boolean wideSearch(Block sourceBlock, int sourceLevel, Map<Block, Integer> changes) {
    Queue<Block> openSet = new LinkedList<>();
    Set<Block> closedSet = new HashSet<>();
    closedSet.add(sourceBlock);
    for (BlockFace face : HORIZONTAL_FACES) openSet.add(sourceBlock.getRelative(face));
    Block bestTarget = null;
    int lowestLevel = sourceLevel;
    int searchedBlocks = 0;
    while (!openSet.isEmpty() && searchedBlocks < PRESSURE_SEARCH_LIMIT) {
      Block currentBlock = openSet.poll();
      if (!closedSet.add(currentBlock) || !currentBlock.getChunk().isLoaded()) continue;
      searchedBlocks++;
      if (currentBlock.getType() != Material.WATER && !canFlowInto(currentBlock, changes)) continue;
      int currentLevel = getWaterLevel(currentBlock, changes);
      if (currentLevel < lowestLevel) {
        lowestLevel = currentLevel;
        bestTarget = currentBlock;
      }
      for (BlockFace face : HORIZONTAL_FACES) {
        Block neighbor = currentBlock.getRelative(face);
        if (!closedSet.contains(neighbor)) openSet.add(neighbor);
      }
    }
    if (bestTarget != null && sourceLevel > lowestLevel) {
      applyDynamicAdditiveFlow(sourceBlock, bestTarget, changes);
      return true;
    }
    return false;
  }

  private void applyDynamicAdditiveFlow(Block source, Block target, Map<Block, Integer> changes) {
    int sourceLevel = getWaterLevel(source, changes);
    int targetLevel = getWaterLevel(target, changes);

    if (changes.containsKey(source) && changes.get(source) != sourceLevel) {
      return;
    }

    int difference = sourceLevel - targetLevel;

    int flowAmount = Math.max(1, difference / 3);

    changes.put(source, sourceLevel - flowAmount);

    int originalTargetLevel = getWaterLevelFromWorld(target);
    int plannedTargetLevel = changes.getOrDefault(target, originalTargetLevel);

    changes.put(target, Math.min(8, plannedTargetLevel + flowAmount));
  }

  private boolean drainPuddle(Block b, Map<Block, Integer> c) {
    if (getWaterLevel(b.getRelative(BlockFace.DOWN), c) > 0) return false;
    Queue<Block> q = new LinkedList<>();
    Set<Block> v = new HashSet<>();
    for (BlockFace f : HORIZONTAL_FACES) {
      Block n = b.getRelative(f);
      if (canFlowInto(n, c)) {
        q.add(n);
        v.add(n);
      }
    }
    int r = 1;
    while (!q.isEmpty() && r <= 2) {
      int s = q.size();
      for (int i = 0; i < s; i++) {
        Block p = q.poll();
        Block l = p.getRelative(BlockFace.DOWN);
        if (canFlowInto(l, c) || getWaterLevel(l, c) > 0) {
          int t = getWaterLevel(l, c);
          if (t < 8) {
            c.put(b, 0);
            c.put(l, t + 1);
            return true;
          }
        }
        if (r < 2) {
          for (BlockFace f : HORIZONTAL_FACES) {
            Block next = p.getRelative(f);
            if (!v.contains(next) && canFlowInto(next, c)) {
              v.add(next);
              q.add(next);
            }
          }
        }
      }
      r++;
    }
    return false;
  }

  private boolean flowDown(Block b, Map<Block, Integer> c) {
    Block d = b.getRelative(BlockFace.DOWN);
    int s = getWaterLevel(b, c);
    if (canFlowInto(d, c)) {
      c.put(b, 0);
      c.put(d, s);
      return true;
    }
    if (d.getType() == Material.WATER) {
      int n = getWaterLevel(d, c);
      if (n < 8) {
        int a = 8 - n;
        int f = Math.min(s, a);
        if (f > 0) {
          c.put(d, n + f);
          c.put(b, s - f);
          return true;
        }
      }
    }
    return false;
  }

  private int getWaterLevel(Block b, Map<Block, Integer> c) {
    return c.getOrDefault(b, getWaterLevelFromWorld(b));
  }

  private int getWaterLevelFromWorld(Block b) {
    if (b.getType() == Material.WATER) {
      Levelled d = (Levelled) b.getBlockData();
      return d.getLevel() == 0 ? 8 : 8 - d.getLevel();
    }
    return 0;
  }

  private void setWaterLevelInWorld(Block b, int l) {
    l = Math.min(l, 8);
    if (l <= 0) {
      if (b.getType() == Material.WATER) b.setType(Material.AIR, false);
    } else {
      b.setType(Material.WATER, false);
      Levelled d = (Levelled) b.getBlockData();
      d.setLevel(l == 8 ? 0 : 8 - l);
      b.setBlockData(d, false);
    }
  }

  private void collectNeighborsToWake(Block b, Set<Block> n) {
    for (int x = -1; x <= 1; x++) {
      for (int y = -1; y <= 1; y++) {
        for (int z = -1; z <= 1; z++) {
          if (x == 0 && y == 0 && z == 0) continue;
          n.add(b.getRelative(x, y, z));
        }
      }
    }
  }

  private boolean canFlowInto(Block b, Map<Block, Integer> c) {
    if (c.containsKey(b) && c.get(b) > 0) return false;
    return b.isPassable() && !b.isLiquid();
  }
}