From ddc938a38c9f5814ab1e63666ec93d3e5c82dd85 Mon Sep 17 00:00:00 2001
From: fullwall <fullwall25@gmail.com>
Date: Sun, 11 Jun 2023 16:00:33 +0800
Subject: [PATCH] Simplify mount trait

---
 .../net/citizensnpcs/trait/MountTrait.java    | 43 ++++++++-----------
 1 file changed, 19 insertions(+), 24 deletions(-)

diff --git a/main/src/main/java/net/citizensnpcs/trait/MountTrait.java b/main/src/main/java/net/citizensnpcs/trait/MountTrait.java
index c17d697ca..78f2bb029 100644
--- a/main/src/main/java/net/citizensnpcs/trait/MountTrait.java
+++ b/main/src/main/java/net/citizensnpcs/trait/MountTrait.java
@@ -17,27 +17,26 @@ import net.citizensnpcs.util.NMS;
  */
 @TraitName("mounttrait")
 public class MountTrait extends Trait {
-    private UUID mountedOn;
-    private boolean triggered = false;
+    private UUID currentMount;
     @Persist("mountedon")
-    private String uuid;
+    private UUID uuid;
 
     public MountTrait() {
         super("mounttrait");
     }
 
-    public void checkMounted(Entity mounted) {
-        if (mountedOn == null || (mounted != null && mounted.getUniqueId().equals(mountedOn)))
+    public void checkMounted() {
+        if (uuid == null || uuid.equals(currentMount))
             return;
-        NPC other = CitizensAPI.getNPCRegistry().getByUniqueId(mountedOn);
+        NPC other = CitizensAPI.getNPCRegistry().getByUniqueId(uuid);
         if (other != null && other.isSpawned()) {
             NMS.mount(other.getEntity(), npc.getEntity());
-            triggered = true;
+            currentMount = uuid;
         }
     }
 
     public UUID getMountedOn() {
-        return mountedOn;
+        return currentMount;
     }
 
     @Override
@@ -47,44 +46,40 @@ public class MountTrait extends Trait {
         }
     }
 
+    @Override
+    public void onRemove() {
+        onDespawn();
+    }
+
     @Override
     public void onSpawn() {
-        checkMounted(null);
+        checkMounted();
     }
 
     @Override
     public void run() {
         if (!npc.isSpawned())
             return;
-        if (!triggered && uuid != null) {
-            try {
-                mountedOn = UUID.fromString(uuid);
-            } catch (IllegalArgumentException e) {
-                mountedOn = null;
-            }
-            checkMounted(null);
-        }
         Entity vehicle = NMS.getVehicle(npc.getEntity());
-        if (vehicle == null && !triggered) {
-            mountedOn = null;
+        if (vehicle == null && currentMount != null) {
+            currentMount = null;
         } else if (vehicle instanceof NPCHolder) {
             setMountedOn(((NPCHolder) vehicle).getNPC().getUniqueId());
         }
-        checkMounted(vehicle);
+        checkMounted();
     }
 
     public void setMountedOn(UUID uuid) {
-        this.mountedOn = uuid;
-        this.uuid = uuid.toString();
+        this.uuid = uuid;
     }
 
     public void unmount() {
-        if (mountedOn == null)
+        if (currentMount == null)
             return;
         if (NMS.getVehicle(npc.getEntity()) != null) {
             npc.getEntity().leaveVehicle();
         }
         uuid = null;
-        mountedOn = null;
+        currentMount = null;
     }
 }
\ No newline at end of file