add functionality and begin work on workstation optimization

src/main/java/me/xginko/villageroptimizer/WrappedVillager.java

@@ -83,7 +83,7 @@ public final class WrappedVillager {
     /**
      * @param type OptimizationType the villager should be set to.
      */
-    public void setOptimization(OptimizationType type) {
+    public void setOptimizationType(OptimizationType type) {
         me.xginko.villageroptimizer.VillagerOptimizer.getFoliaLib().getImpl().runAtEntityTimer(villager, setOptimization -> {
             // Keep repeating task until villager is no longer trading with a player
             if (villager.isTrading()) return;

src/main/java/me/xginko/villageroptimizer/commands/VillagerOptimizerCommand.java

@@ -4,21 +4,17 @@ import me.xginko.villageroptimizer.VillagerOptimizer;
 import me.xginko.villageroptimizer.commands.optimizevillagers.OptVillagersRadius;
 import me.xginko.villageroptimizer.commands.unoptimizevillagers.UnOptVillagersRadius;
 import me.xginko.villageroptimizer.commands.villageroptimizer.VillagerOptimizerCmd;
-import org.bukkit.command.Command;
 import org.bukkit.command.CommandExecutor;
 import org.bukkit.command.CommandMap;
-import org.bukkit.command.CommandSender;
+import org.bukkit.command.TabCompleter;
-import org.jetbrains.annotations.NotNull;
 
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 
-public interface VillagerOptimizerCommand extends CommandExecutor {
+public interface VillagerOptimizerCommand extends CommandExecutor, TabCompleter {
 
     String label();
-    @Override
-    boolean onCommand(@NotNull CommandSender sender, @NotNull Command command, @NotNull String label, String[] args);
 
     List<String> NO_TABCOMPLETES = Collections.emptyList();
     List<String> RADIUS_TABCOMPLETES = List.of("5", "10", "25", "50");

src/main/java/me/xginko/villageroptimizer/commands/optimizevillagers/OptVillagersRadius.java

@@ -5,9 +5,9 @@ import me.xginko.villageroptimizer.VillagerOptimizer;
 import me.xginko.villageroptimizer.WrappedVillager;
 import me.xginko.villageroptimizer.commands.VillagerOptimizerCommand;
 import me.xginko.villageroptimizer.config.Config;
+import me.xginko.villageroptimizer.enums.OptimizationType;
 import me.xginko.villageroptimizer.enums.permissions.Bypass;
 import me.xginko.villageroptimizer.enums.permissions.Commands;
-import me.xginko.villageroptimizer.enums.OptimizationType;
 import me.xginko.villageroptimizer.events.VillagerOptimizeEvent;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.TextReplacementConfig;
@@ -15,7 +15,6 @@ import net.kyori.adventure.text.format.NamedTextColor;
 import net.kyori.adventure.text.format.TextDecoration;
 import org.bukkit.command.Command;
 import org.bukkit.command.CommandSender;
-import org.bukkit.command.TabCompleter;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.EntityType;
 import org.bukkit.entity.Player;
@@ -25,7 +24,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 
-public class OptVillagersRadius implements VillagerOptimizerCommand, TabCompleter {
+public class OptVillagersRadius implements VillagerOptimizerCommand {
 
     private final long cooldown;
     private final int max_radius;
@@ -101,7 +100,7 @@ public class OptVillagersRadius implements VillagerOptimizerCommand, TabComplete
                 if (player_has_cooldown_bypass || wVillager.canOptimize(cooldown)) {
                     VillagerOptimizeEvent optimizeEvent = new VillagerOptimizeEvent(wVillager, OptimizationType.COMMAND, player);
                     if (optimizeEvent.callEvent()) {
-                        wVillager.setOptimization(optimizeEvent.getOptimizationType());
+                        wVillager.setOptimizationType(optimizeEvent.getOptimizationType());
                         wVillager.saveOptimizeTime();
                         successCount++;
                     }

src/main/java/me/xginko/villageroptimizer/commands/unoptimizevillagers/UnOptVillagersRadius.java

@@ -4,8 +4,8 @@ import me.xginko.villageroptimizer.VillagerCache;
 import me.xginko.villageroptimizer.VillagerOptimizer;
 import me.xginko.villageroptimizer.WrappedVillager;
 import me.xginko.villageroptimizer.commands.VillagerOptimizerCommand;
-import me.xginko.villageroptimizer.enums.permissions.Commands;
 import me.xginko.villageroptimizer.enums.OptimizationType;
+import me.xginko.villageroptimizer.enums.permissions.Commands;
 import me.xginko.villageroptimizer.events.VillagerUnoptimizeEvent;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.TextReplacementConfig;
@@ -13,7 +13,6 @@ import net.kyori.adventure.text.format.NamedTextColor;
 import net.kyori.adventure.text.format.TextDecoration;
 import org.bukkit.command.Command;
 import org.bukkit.command.CommandSender;
-import org.bukkit.command.TabCompleter;
 import org.bukkit.entity.Entity;
 import org.bukkit.entity.EntityType;
 import org.bukkit.entity.Player;
@@ -23,7 +22,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 
-public class UnOptVillagersRadius implements VillagerOptimizerCommand, TabCompleter {
+public class UnOptVillagersRadius implements VillagerOptimizerCommand {
 
     private final int max_radius;
 
@@ -92,7 +91,7 @@ public class UnOptVillagersRadius implements VillagerOptimizerCommand, TabComple
                 if (wVillager.isOptimized()) {
                     VillagerUnoptimizeEvent unOptimizeEvent = new VillagerUnoptimizeEvent(wVillager, player, OptimizationType.COMMAND);
                     if (unOptimizeEvent.callEvent()) {
-                        wVillager.setOptimization(OptimizationType.NONE);
+                        wVillager.setOptimizationType(OptimizationType.NONE);
                         successCount++;
                     }
                 }

src/main/java/me/xginko/villageroptimizer/commands/villageroptimizer/VillagerOptimizerCmd.java

@@ -11,12 +11,11 @@ import net.kyori.adventure.text.Component;
 import net.kyori.adventure.text.format.NamedTextColor;
 import org.bukkit.command.Command;
 import org.bukkit.command.CommandSender;
-import org.bukkit.command.TabCompleter;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.List;
 
-public class VillagerOptimizerCmd implements TabCompleter, VillagerOptimizerCommand {
+public class VillagerOptimizerCmd implements VillagerOptimizerCommand {
 
     private final List<SubCommand> subCommands;
     private final List<String> tabCompleter;
@@ -38,19 +37,19 @@ public class VillagerOptimizerCmd implements TabCompleter, VillagerOptimizerComm
 
     @Override
     public boolean onCommand(@NotNull CommandSender sender, @NotNull Command command, @NotNull String label, String[] args) {
-        if (args.length > 0) {
+        if (args.length == 0) {
-            boolean cmdExists = false;
+            sendCommandOverview(sender);
-            for (SubCommand subCommand : subCommands) {
+            return true;
+        }
+
+        for (final SubCommand subCommand : subCommands) {
             if (args[0].equalsIgnoreCase(subCommand.getLabel())) {
                 subCommand.perform(sender, args);
-                    cmdExists = true;
+                return true;
-                    break;
             }
         }
-            if (!cmdExists) sendCommandOverview(sender);
+
-        } else {
         sendCommandOverview(sender);
-        }
         return true;
     }
 

src/main/java/me/xginko/villageroptimizer/config/Config.java

@@ -8,7 +8,6 @@ import java.io.File;
 import java.util.List;
 import java.util.Locale;
 
-@SuppressWarnings({"ALL", "EscapedSpace"})
 public class Config {
 
     private final @NotNull ConfigFile config;

src/main/java/me/xginko/villageroptimizer/enums/Keyring.java

@@ -6,7 +6,6 @@ import org.bukkit.Keyed;
 import org.bukkit.NamespacedKey;
 import org.bukkit.persistence.PersistentDataContainer;
 import org.bukkit.plugin.Plugin;
-import org.intellij.lang.annotations.Subst;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Locale;

src/main/java/me/xginko/villageroptimizer/modules/VillagerOptimizerModule.java

@@ -23,6 +23,8 @@ public interface VillagerOptimizerModule {
         modules.add(new OptimizeByBlock());
         modules.add(new OptimizeByWorkstation());
 
+        modules.add(new EnableLeashingVillagers());
+        modules.add(new FixOptimisationAfterCure());
         modules.add(new RestockOptimizedTrades());
         modules.add(new LevelOptimizedProfession());
         modules.add(new RenameOptimizedVillagers());
@@ -30,7 +32,7 @@ public interface VillagerOptimizerModule {
         modules.add(new PreventUnoptimizedTrading());
         modules.add(new PreventOptimizedTargeting());
         modules.add(new PreventOptimizedDamage());
-        modules.add(new EnableLeashingVillagers());
+        modules.add(new UnoptimizeOnJobLoose());
 
         modules.add(new VillagerChunkLimit());
 

src/main/java/me/xginko/villageroptimizer/modules/gameplay/FixOptimisationAfterCure.java

@@ -0,0 +1,49 @@
+package me.xginko.villageroptimizer.modules.gameplay;
+
+import me.xginko.villageroptimizer.VillagerOptimizer;
+import me.xginko.villageroptimizer.WrappedVillager;
+import me.xginko.villageroptimizer.modules.VillagerOptimizerModule;
+import org.bukkit.entity.EntityType;
+import org.bukkit.entity.Villager;
+import org.bukkit.event.EventHandler;
+import org.bukkit.event.EventPriority;
+import org.bukkit.event.HandlerList;
+import org.bukkit.event.Listener;
+import org.bukkit.event.entity.EntityTransformEvent;
+
+import java.util.concurrent.TimeUnit;
+
+public class FixOptimisationAfterCure implements VillagerOptimizerModule, Listener {
+
+    public FixOptimisationAfterCure() {}
+
+    @Override
+    public void enable() {
+        VillagerOptimizer plugin = VillagerOptimizer.getInstance();
+        plugin.getServer().getPluginManager().registerEvents(this, plugin);
+    }
+
+    @Override
+    public void disable() {
+        HandlerList.unregisterAll(this);
+    }
+
+    @Override
+    public boolean shouldEnable() {
+        return true;
+    }
+
+    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
+    private void onTransform(EntityTransformEvent event) {
+        if (
+                event.getTransformReason().equals(EntityTransformEvent.TransformReason.CURED)
+                && event.getTransformedEntity().getType().equals(EntityType.VILLAGER)
+        ) {
+            Villager villager = (Villager) event.getTransformedEntity();
+            VillagerOptimizer.getFoliaLib().getImpl().runAtEntityLater(villager, () -> {
+                WrappedVillager wVillager = VillagerOptimizer.getCache().getOrAdd(villager);
+                wVillager.setOptimizationType(wVillager.getOptimizationType());
+            }, 2, TimeUnit.SECONDS);
+        }
+    }
+}

src/main/java/me/xginko/villageroptimizer/modules/gameplay/LevelOptimizedProfession.java

@@ -1,5 +1,6 @@
 package me.xginko.villageroptimizer.modules.gameplay;
 
+import com.tcoded.folialib.impl.ServerImplementation;
 import me.xginko.villageroptimizer.VillagerOptimizer;
 import me.xginko.villageroptimizer.VillagerCache;
 import me.xginko.villageroptimizer.config.Config;
@@ -22,12 +23,14 @@ import java.util.concurrent.TimeUnit;
 
 public class LevelOptimizedProfession implements VillagerOptimizerModule, Listener {
 
+    private final ServerImplementation scheduler;
     private final VillagerCache villagerCache;
     private final boolean notify_player;
     private final long cooldown_millis;
 
     public LevelOptimizedProfession() {
         shouldEnable();
+        this.scheduler = VillagerOptimizer.getFoliaLib().getImpl();
         this.villagerCache = VillagerOptimizer.getCache();
         Config config = VillagerOptimizer.getConfiguration();
         config.master().addComment("gameplay.level-optimized-profession", """
@@ -68,14 +71,15 @@ public class LevelOptimizedProfession implements VillagerOptimizerModule, Listen
 
             if (wVillager.canLevelUp(cooldown_millis)) {
                 if (wVillager.calculateLevel() > villager.getVillagerLevel()) {
-                    VillagerOptimizer.getFoliaLib().getImpl().runAtEntity(villager, enableAI -> {
+                    scheduler.runAtEntity(villager, enableAI -> {
                         villager.addPotionEffect(new PotionEffect(PotionEffectType.SLOW, 120, 120, false, false));
                         villager.setAware(true);
-                    });
+
-                    VillagerOptimizer.getFoliaLib().getImpl().runAtEntityLater(villager, disableAI -> {
+                        scheduler.runAtEntityLater(villager, disableAI -> {
                             villager.setAware(false);
                             wVillager.saveLastLevelUp();
                         }, 5, TimeUnit.SECONDS);
+                    });
                 }
             } else {
                 if (notify_player) {

src/main/java/me/xginko/villageroptimizer/modules/gameplay/UnoptimizeOnJobLoose.java

@@ -0,0 +1,49 @@
+package me.xginko.villageroptimizer.modules.gameplay;
+
+import me.xginko.villageroptimizer.VillagerCache;
+import me.xginko.villageroptimizer.VillagerOptimizer;
+import me.xginko.villageroptimizer.WrappedVillager;
+import me.xginko.villageroptimizer.enums.OptimizationType;
+import me.xginko.villageroptimizer.modules.VillagerOptimizerModule;
+import org.bukkit.event.EventHandler;
+import org.bukkit.event.EventPriority;
+import org.bukkit.event.HandlerList;
+import org.bukkit.event.Listener;
+import org.bukkit.event.entity.VillagerCareerChangeEvent;
+
+public class UnoptimizeOnJobLoose implements VillagerOptimizerModule, Listener {
+
+    private final VillagerCache villagerCache;
+
+    public UnoptimizeOnJobLoose() {
+        this.villagerCache = VillagerOptimizer.getCache();
+    }
+
+    @Override
+    public void enable() {
+        VillagerOptimizer plugin = VillagerOptimizer.getInstance();
+        plugin.getServer().getPluginManager().registerEvents(this, plugin);
+    }
+
+    @Override
+    public void disable() {
+        HandlerList.unregisterAll(this);
+    }
+
+    @Override
+    public boolean shouldEnable() {
+        return VillagerOptimizer.getConfiguration().getBoolean("gameplay.unoptimize-on-job-loose.enable", true,
+                "Villagers that get their jobs reset will become unoptimized again. Highly recommended to leave on.");
+    }
+
+    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
+    private void onJobReset(VillagerCareerChangeEvent event) {
+        if (!event.getReason().equals(VillagerCareerChangeEvent.ChangeReason.LOSING_JOB)) return;
+
+        WrappedVillager wrappedVillager = villagerCache.getOrAdd(event.getEntity());
+
+        if (wrappedVillager.isOptimized()) {
+            wrappedVillager.setOptimizationType(OptimizationType.NONE);
+        }
+    }
+}

src/main/java/me/xginko/villageroptimizer/modules/optimization/OptimizeByBlock.java

@@ -98,7 +98,7 @@ public class OptimizeByBlock implements VillagerOptimizerModule, Listener {
         if (!player.hasPermission(Optimize.BLOCK.get())) return;
         if (only_while_sneaking && !player.isSneaking()) return;
 
-        final Location blockLoc = placed.getLocation();
+        final Location blockLoc = placed.getLocation().toCenterLocation();
         WrappedVillager closestOptimizableVillager = null;
         double closestDistance = Double.MAX_VALUE;
 
@@ -109,7 +109,7 @@ public class OptimizeByBlock implements VillagerOptimizerModule, Listener {
             if (profession.equals(Villager.Profession.NONE) || profession.equals(Villager.Profession.NITWIT)) continue;
 
             WrappedVillager wVillager = villagerCache.getOrAdd(villager);
-            final double distance = entity.getLocation().distance(blockLoc);
+            final double distance = entity.getLocation().distanceSquared(blockLoc);
 
             if (distance < closestDistance && wVillager.canOptimize(cooldown_millis)) {
                 closestOptimizableVillager = wVillager;
@@ -123,7 +123,7 @@ public class OptimizeByBlock implements VillagerOptimizerModule, Listener {
             VillagerOptimizeEvent optimizeEvent = new VillagerOptimizeEvent(closestOptimizableVillager, OptimizationType.BLOCK, player, event.isAsynchronous());
             if (!optimizeEvent.callEvent()) return;
 
-            closestOptimizableVillager.setOptimization(optimizeEvent.getOptimizationType());
+            closestOptimizableVillager.setOptimizationType(optimizeEvent.getOptimizationType());
             closestOptimizableVillager.saveOptimizeTime();
 
             if (notify_player) {
@@ -162,7 +162,7 @@ public class OptimizeByBlock implements VillagerOptimizerModule, Listener {
         if (!player.hasPermission(Optimize.BLOCK.get())) return;
         if (only_while_sneaking && !player.isSneaking()) return;
 
-        final Location blockLoc = broken.getLocation();
+        final Location blockLoc = broken.getLocation().toCenterLocation();
         WrappedVillager closestOptimizedVillager = null;
         double closestDistance = Double.MAX_VALUE;
 
@@ -171,7 +171,7 @@ public class OptimizeByBlock implements VillagerOptimizerModule, Listener {
             Villager villager = (Villager) entity;
 
             WrappedVillager wVillager = villagerCache.getOrAdd(villager);
-            final double distance = entity.getLocation().distance(blockLoc);
+            final double distance = entity.getLocation().distanceSquared(blockLoc);
 
             if (distance < closestDistance && wVillager.isOptimized()) {
                 closestOptimizedVillager = wVillager;
@@ -184,7 +184,7 @@ public class OptimizeByBlock implements VillagerOptimizerModule, Listener {
         VillagerUnoptimizeEvent unOptimizeEvent = new VillagerUnoptimizeEvent(closestOptimizedVillager, player, OptimizationType.BLOCK, event.isAsynchronous());
         if (!unOptimizeEvent.callEvent()) return;
 
-        closestOptimizedVillager.setOptimization(OptimizationType.NONE);
+        closestOptimizedVillager.setOptimizationType(OptimizationType.NONE);
 
         if (notify_player) {
             final TextReplacementConfig vilProfession = TextReplacementConfig.builder()

src/main/java/me/xginko/villageroptimizer/modules/optimization/OptimizeByNametag.java

@@ -102,7 +102,7 @@ public class OptimizeByNametag implements VillagerOptimizerModule, Listener {
                     villager.customName(newVillagerName);
                 }
 
-                wVillager.setOptimization(optimizeEvent.getOptimizationType());
+                wVillager.setOptimizationType(optimizeEvent.getOptimizationType());
                 wVillager.saveOptimizeTime();
 
                 if (notify_player)
@@ -125,7 +125,7 @@ public class OptimizeByNametag implements VillagerOptimizerModule, Listener {
                 VillagerUnoptimizeEvent unOptimizeEvent = new VillagerUnoptimizeEvent(wVillager, player, OptimizationType.NAMETAG, event.isAsynchronous());
                 if (!unOptimizeEvent.callEvent()) return;
 
-                wVillager.setOptimization(OptimizationType.NONE);
+                wVillager.setOptimizationType(OptimizationType.NONE);
 
                 if (notify_player)
                     VillagerOptimizer.getLang(player.locale()).nametag_unoptimize_success.forEach(player::sendMessage);

src/main/java/me/xginko/villageroptimizer/modules/optimization/OptimizeByWorkstation.java

@@ -1,5 +1,8 @@
 package me.xginko.villageroptimizer.modules.optimization;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.tcoded.folialib.impl.ServerImplementation;
 import io.papermc.paper.event.entity.EntityMoveEvent;
 import me.xginko.villageroptimizer.VillagerCache;
 import me.xginko.villageroptimizer.VillagerOptimizer;
@@ -19,31 +22,41 @@ import org.bukkit.entity.Entity;
 import org.bukkit.entity.EntityType;
 import org.bukkit.entity.Player;
 import org.bukkit.entity.Villager;
+import org.bukkit.entity.memory.MemoryKey;
 import org.bukkit.event.EventHandler;
 import org.bukkit.event.EventPriority;
 import org.bukkit.event.HandlerList;
 import org.bukkit.event.Listener;
 import org.bukkit.event.block.BlockBreakEvent;
 import org.bukkit.event.block.BlockPlaceEvent;
+import org.bukkit.event.entity.VillagerCareerChangeEvent;
+import org.jetbrains.annotations.Nullable;
 
+import java.time.Duration;
+import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
-import static me.xginko.villageroptimizer.utils.CommonUtil.getWorkstationProfession;
-
 public class OptimizeByWorkstation implements VillagerOptimizerModule, Listener {
 
+    private final ServerImplementation scheduler;
     private final VillagerCache villagerCache;
+    private final Cache<UUID, Location> cachedVillagerJobSites;
     private final long cooldown_millis;
     private final double search_radius;
     private final boolean only_while_sneaking, log_enabled, notify_player;
 
     public OptimizeByWorkstation() {
         shouldEnable();
+        this.scheduler = VillagerOptimizer.getFoliaLib().getImpl();
         this.villagerCache = VillagerOptimizer.getCache();
+        this.cachedVillagerJobSites = Caffeine.newBuilder().expireAfterWrite(Duration.ofSeconds(2)).build();
+
+        // Broken and unfinished, in progress
+
         Config config = VillagerOptimizer.getConfiguration();
         config.master().addComment("optimization-methods.workstation-optimization.enable", """
-                When enabled, the closest villager near a matching workstation being placed will be optimized.\s
+                When enabled, villagers that have a job and have been traded with at least once will become optimized,\s
-                If a nearby matching workstation is broken, the villager will become unoptimized again.""");
+                if near their workstation. If the workstation is broken, the villager will become unoptimized again.""");
         this.search_radius = config.getDouble("optimization-methods.workstation-optimization.search-radius-in-blocks", 2.0, """
                 The radius in blocks a villager can be away from the player when he places a workstation.\s
                 The closest unoptimized villager to the player will be optimized.""") / 2;
@@ -74,55 +87,92 @@ public class OptimizeByWorkstation implements VillagerOptimizerModule, Listener
         return VillagerOptimizer.getConfiguration().getBoolean("optimization-methods.workstation-optimization.enable", false);
     }
 
+    private @Nullable Location getJobSite(Villager villager) {
+        Location jobSite = cachedVillagerJobSites.getIfPresent(villager.getUniqueId());
+        if (jobSite == null) {
+            jobSite = villager.getMemory(MemoryKey.JOB_SITE);
+            cachedVillagerJobSites.put(villager.getUniqueId(), jobSite);
+        }
+        return jobSite;
+    }
+
     @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
     private void onVillagerMove(EntityMoveEvent event) {
-        if (!event.getEntity().getType().equals(EntityType.VILLAGER)) return;
+        if (!event.getEntityType().equals(EntityType.VILLAGER)) return;
-        Villager villager = (Villager) event.getEntity();
+
+        final Villager villager = (Villager) event.getEntity();
+
         if (villager.getProfession().equals(Villager.Profession.NONE)) return;
+        if (CommonUtil.canLooseProfession(villager)) return;
 
+        final Location jobSite = getJobSite(villager);
+        if (jobSite == null) return;
+        // Using distanceSquared is faster. 1*1=1 -> 1 block away from the workstation
+        if (!(villager.getLocation().distanceSquared(jobSite) <= 1)) return;
 
+        WrappedVillager wrappedVillager = villagerCache.getOrAdd(villager);
+
+        if (wrappedVillager.canOptimize(cooldown_millis)) {
+            wrappedVillager.setOptimizationType(OptimizationType.WORKSTATION);
+        }
+    }
+
+    @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
+    private void onCareerChange(VillagerCareerChangeEvent event) {
+        if (!event.getReason().equals(VillagerCareerChangeEvent.ChangeReason.EMPLOYED)) return;
+        if (CommonUtil.canLooseProfession(event.getEntity())) return;
+
+        WrappedVillager wrappedVillager = villagerCache.getOrAdd(event.getEntity());
+
+        if (!wrappedVillager.isOptimized() && wrappedVillager.canOptimize(cooldown_millis)) {
+            wrappedVillager.setOptimizationType(OptimizationType.WORKSTATION);
+        }
     }
 
     @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
     private void onBlockPlace(BlockPlaceEvent event) {
         Block placed = event.getBlock();
-        Villager.Profession workstationProfession = getWorkstationProfession(placed.getType());
+        Villager.Profession workstationProfession = CommonUtil.getWorkstationProfession(placed.getType());
         if (workstationProfession.equals(Villager.Profession.NONE)) return;
+
         Player player = event.getPlayer();
         if (!player.hasPermission(Optimize.WORKSTATION.get())) return;
         if (only_while_sneaking && !player.isSneaking()) return;
 
-        final Location workstationLoc = placed.getLocation();
+        final Location workstationLoc = placed.getLocation().toCenterLocation();
-        WrappedVillager closestOptimizableVillager = null;
+        WrappedVillager villagerThatClaimedWorkstation = null;
-        double closestDistance = Double.MAX_VALUE;
 
         for (Entity entity : workstationLoc.getNearbyEntities(search_radius, search_radius, search_radius)) {
             if (!entity.getType().equals(EntityType.VILLAGER)) continue;
             Villager villager = (Villager) entity;
             if (!villager.getProfession().equals(workstationProfession)) continue;
+            // Ignore villagers that haven't been locked into a profession yet, so we don't disturb trade rollers
+            if (CommonUtil.canLooseProfession(villager)) continue;
+            Location jobSite = getJobSite(villager);
+            if (jobSite == null) continue;
+            if (!workstationLoc.equals(jobSite.toBlockLocation())) continue;
 
             WrappedVillager wVillager = villagerCache.getOrAdd(villager);
-            final double distance = entity.getLocation().distance(workstationLoc);
 
-            if (distance < closestDistance && wVillager.canOptimize(cooldown_millis)) {
+            if (wVillager.canOptimize(cooldown_millis)) {
-                closestOptimizableVillager = wVillager;
+                villagerThatClaimedWorkstation = wVillager;
-                closestDistance = distance;
+                break;
             }
         }
 
-        if (closestOptimizableVillager == null) return;
+        if (villagerThatClaimedWorkstation == null) return;
 
-        if (closestOptimizableVillager.canOptimize(cooldown_millis) || player.hasPermission(Bypass.WORKSTATION_COOLDOWN.get())) {
+        if (villagerThatClaimedWorkstation.canOptimize(cooldown_millis) || player.hasPermission(Bypass.WORKSTATION_COOLDOWN.get())) {
-            VillagerOptimizeEvent optimizeEvent = new VillagerOptimizeEvent(closestOptimizableVillager, OptimizationType.WORKSTATION, player, event.isAsynchronous());
+            VillagerOptimizeEvent optimizeEvent = new VillagerOptimizeEvent(villagerThatClaimedWorkstation, OptimizationType.WORKSTATION, player, event.isAsynchronous());
             if (!optimizeEvent.callEvent()) return;
 
-            closestOptimizableVillager.setOptimization(optimizeEvent.getOptimizationType());
+            villagerThatClaimedWorkstation.setOptimizationType(optimizeEvent.getOptimizationType());
-            closestOptimizableVillager.saveOptimizeTime();
+            villagerThatClaimedWorkstation.saveOptimizeTime();
 
             if (notify_player) {
                 final TextReplacementConfig vilProfession = TextReplacementConfig.builder()
                         .matchLiteral("%vil_profession%")
-                        .replacement(closestOptimizableVillager.villager().getProfession().toString().toLowerCase())
+                        .replacement(villagerThatClaimedWorkstation.villager().getProfession().toString().toLowerCase())
                         .build();
                 final TextReplacementConfig placedWorkstation = TextReplacementConfig.builder()
                         .matchLiteral("%workstation%")
@@ -136,11 +186,11 @@ public class OptimizeByWorkstation implements VillagerOptimizerModule, Listener
             if (log_enabled)
                 VillagerOptimizer.getLog().info(player.getName() + " optimized a villager using workstation: '" + placed.getType().toString().toLowerCase() + "'");
         } else {
-            CommonUtil.shakeHead(closestOptimizableVillager.villager());
+            CommonUtil.shakeHead(villagerThatClaimedWorkstation.villager());
             if (notify_player) {
                 final TextReplacementConfig timeLeft = TextReplacementConfig.builder()
                         .matchLiteral("%time%")
-                        .replacement(CommonUtil.formatTime(closestOptimizableVillager.getOptimizeCooldownMillis(cooldown_millis)))
+                        .replacement(CommonUtil.formatTime(villagerThatClaimedWorkstation.getOptimizeCooldownMillis(cooldown_millis)))
                         .build();
                 VillagerOptimizer.getLang(player.locale()).nametag_on_optimize_cooldown.forEach(line -> player.sendMessage(line
                         .replaceText(timeLeft)
@@ -152,7 +202,7 @@ public class OptimizeByWorkstation implements VillagerOptimizerModule, Listener
     @EventHandler(priority = EventPriority.HIGHEST, ignoreCancelled = true)
     private void onBlockBreak(BlockBreakEvent event) {
         Block broken = event.getBlock();
-        Villager.Profession workstationProfession = getWorkstationProfession(broken.getType());
+        Villager.Profession workstationProfession = CommonUtil.getWorkstationProfession(broken.getType());
         if (workstationProfession.equals(Villager.Profession.NONE)) return;
         Player player = event.getPlayer();
         if (!player.hasPermission(Optimize.WORKSTATION.get())) return;
@@ -181,7 +231,7 @@ public class OptimizeByWorkstation implements VillagerOptimizerModule, Listener
         VillagerUnoptimizeEvent unOptimizeEvent = new VillagerUnoptimizeEvent(closestOptimizedVillager, player, OptimizationType.WORKSTATION, event.isAsynchronous());
         if (!unOptimizeEvent.callEvent()) return;
 
-        closestOptimizedVillager.setOptimization(OptimizationType.NONE);
+        closestOptimizedVillager.setOptimizationType(OptimizationType.NONE);
 
         if (notify_player) {
             final TextReplacementConfig vilProfession = TextReplacementConfig.builder()

src/main/java/me/xginko/villageroptimizer/utils/CommonUtil.java

@@ -60,4 +60,9 @@ public class CommonUtil {
             default -> Villager.Profession.NONE;
         };
     }
+
+    public static boolean canLooseProfession(@NotNull Villager villager) {
+        // A villager with a level of 1 and no trading experience is liable to lose its profession.
+        return villager.getVillagerLevel() <= 1 && villager.getVillagerExperience() <= 0;
+    }
 }