DoAfter Refactor (#13225)

Co-authored-by: DrSmugleaf <drsmugleaf@gmail.com>

Content.Server/Implants/ImplanterSystem.cs

@@ -2,6 +2,7 @@ using System.Threading;
 using Content.Server.DoAfter;
 using Content.Server.Guardian;
 using Content.Server.Popups;
+using Content.Shared.DoAfter;
 using Content.Shared.Hands;
 using Content.Shared.IdentityManagement;
 using Content.Shared.Implants;
@@ -11,7 +12,6 @@ using Content.Shared.Mobs.Components;
 using Content.Shared.Popups;
 using Robust.Shared.Containers;
 using Robust.Shared.GameStates;
-using Robust.Shared.Player;
 
 namespace Content.Server.Implants;
 
@@ -30,9 +30,8 @@ public sealed partial class ImplanterSystem : SharedImplanterSystem
         SubscribeLocalEvent<ImplanterComponent, AfterInteractEvent>(OnImplanterAfterInteract);
         SubscribeLocalEvent<ImplanterComponent, ComponentGetState>(OnImplanterGetState);
 
-        SubscribeLocalEvent<ImplanterComponent, ImplanterImplantCompleteEvent>(OnImplantAttemptSuccess);
-        SubscribeLocalEvent<ImplanterComponent, ImplanterDrawCompleteEvent>(OnDrawAttemptSuccess);
-        SubscribeLocalEvent<ImplanterComponent, ImplanterCancelledEvent>(OnImplantAttemptFail);
+        SubscribeLocalEvent<ImplanterComponent, DoAfterEvent<ImplantEvent>>(OnImplant);
+        SubscribeLocalEvent<ImplanterComponent, DoAfterEvent<DrawEvent>>(OnDraw);
     }
 
     private void OnImplanterAfterInteract(EntityUid uid, ImplanterComponent component, AfterInteractEvent args)
@@ -40,12 +39,6 @@ public sealed partial class ImplanterSystem : SharedImplanterSystem
         if (args.Target == null || !args.CanReach || args.Handled)
             return;
 
-        if (component.CancelToken != null)
-        {
-            args.Handled = true;
-            return;
-        }
-
         //Simplemobs and regular mobs should be injectable, but only regular mobs have mind.
         //So just don't implant/draw anything that isn't living or is a guardian
         //TODO: Rework a bit when surgery is in to work with implant cases
@@ -84,6 +77,9 @@ public sealed partial class ImplanterSystem : SharedImplanterSystem
     /// <param name="implanter">The implanter being used</param>
     public void TryImplant(ImplanterComponent component, EntityUid user, EntityUid target, EntityUid implanter)
     {
+        if (component.CancelToken != null)
+            return;
+
         _popup.PopupEntity(Loc.GetString("injector-component-injecting-user"), target, user);
 
         var userName = Identity.Entity(user, EntityManager);
@@ -92,15 +88,16 @@ public sealed partial class ImplanterSystem : SharedImplanterSystem
         component.CancelToken?.Cancel();
         component.CancelToken = new CancellationTokenSource();
 
-        _doAfter.DoAfter(new DoAfterEventArgs(user, component.ImplantTime, component.CancelToken.Token, target, implanter)
+        var implantEvent = new ImplantEvent();
+
+        _doAfter.DoAfter(new DoAfterEventArgs(user, component.ImplantTime, component.CancelToken.Token,target:target, used:implanter)
         {
             BreakOnUserMove = true,
             BreakOnTargetMove = true,
             BreakOnDamage = true,
             BreakOnStun = true,
-            UsedFinishedEvent = new ImplanterImplantCompleteEvent(implanter, target),
-            UserCancelledEvent = new ImplanterCancelledEvent()
-        });
+            NeedHand = true
+        }, implantEvent);
     }
 
     /// <summary>
@@ -118,15 +115,16 @@ public sealed partial class ImplanterSystem : SharedImplanterSystem
         component.CancelToken?.Cancel();
         component.CancelToken = new CancellationTokenSource();
 
-        _doAfter.DoAfter(new DoAfterEventArgs(user, component.DrawTime, component.CancelToken.Token, target ,implanter)
+        var drawEvent = new DrawEvent();
+
+        _doAfter.DoAfter(new DoAfterEventArgs(user, component.DrawTime, target:target,used:implanter)
         {
             BreakOnUserMove = true,
             BreakOnTargetMove = true,
             BreakOnDamage = true,
             BreakOnStun = true,
-            UsedFinishedEvent = new ImplanterDrawCompleteEvent(implanter, user, target),
-            UsedCancelledEvent = new ImplanterCancelledEvent()
-        });
+            NeedHand = true
+        }, drawEvent);
     }
 
     private void OnImplanterGetState(EntityUid uid, ImplanterComponent component, ref ComponentGetState args)
@@ -134,55 +132,47 @@ public sealed partial class ImplanterSystem : SharedImplanterSystem
         args.State = new ImplanterComponentState(component.CurrentMode, component.ImplantOnly);
     }
 
-    private void OnImplantAttemptSuccess(EntityUid uid, ImplanterComponent component, ImplanterImplantCompleteEvent args)
+    private void OnImplant(EntityUid uid, ImplanterComponent component, DoAfterEvent<ImplantEvent> args)
     {
-        component.CancelToken?.Cancel();
-        component.CancelToken = null;
-        Implant(args.Implanter, args.Target, component);
-    }
-
-    private void OnDrawAttemptSuccess(EntityUid uid, ImplanterComponent component, ImplanterDrawCompleteEvent args)
-    {
-        component.CancelToken?.Cancel();
-        component.CancelToken = null;
-        Draw(args.Implanter, args.User, args.Target, component);
-    }
-
-    private void OnImplantAttemptFail(EntityUid uid, ImplanterComponent component, ImplanterCancelledEvent args)
-    {
-        component.CancelToken?.Cancel();
-        component.CancelToken = null;
-    }
-
-    private sealed class ImplanterImplantCompleteEvent : EntityEventArgs
-    {
-        public EntityUid Implanter;
-        public EntityUid Target;
-
-        public ImplanterImplantCompleteEvent(EntityUid implanter, EntityUid target)
+        if (args.Cancelled)
         {
-            Implanter = implanter;
-            Target = target;
+            component.CancelToken = null;
+            return;
         }
+
+        if (args.Handled || args.Args.Target == null || args.Args.Used == null)
+            return;
+
+        Implant(args.Args.Used.Value, args.Args.Target.Value, component);
+
+        args.Handled = true;
+        component.CancelToken = null;
     }
 
-    private sealed class ImplanterCancelledEvent : EntityEventArgs
+    private void OnDraw(EntityUid uid, ImplanterComponent component, DoAfterEvent<DrawEvent> args)
     {
-
-    }
-
-    private sealed class ImplanterDrawCompleteEvent : EntityEventArgs
-    {
-        public EntityUid Implanter;
-        public EntityUid User;
-        public EntityUid Target;
-
-        public ImplanterDrawCompleteEvent(EntityUid implanter, EntityUid user, EntityUid target)
+        if (args.Cancelled)
         {
-            Implanter = implanter;
-            User = user;
-            Target = target;
+            component.CancelToken = null;
+            return;
         }
+
+        if (args.Handled || args.Args.Used == null || args.Args.Target == null)
+            return;
+
+        Draw(args.Args.Used.Value, args.Args.User, args.Args.Target.Value, component);
+
+        args.Handled = true;
+        component.CancelToken = null;
     }
 
+    private sealed class ImplantEvent : EntityEventArgs
+    {
+
+    }
+
+    private sealed class DrawEvent : EntityEventArgs
+    {
+
+    }
 }