Buckle improvements (#2508)

* Buckle improvements

* Update doc

* Remove redundant argument name

* Remove import

Content.Server/GameObjects/EntitySystems/BuckleSystem.cs

@@ -1,8 +1,11 @@
-using Content.Server.GameObjects.Components.Buckle;
+#nullable enable
+using Content.Server.GameObjects.Components.Buckle;
+using Content.Server.GameObjects.Components.Strap;
 using Content.Server.GameObjects.EntitySystems.Click;
 using JetBrains.Annotations;
+using Robust.Server.GameObjects.EntitySystemMessages;
 using Robust.Server.GameObjects.EntitySystems;
-using Robust.Shared.GameObjects;
+using Robust.Shared.Containers;
 using Robust.Shared.GameObjects.Components.Transform;
 using Robust.Shared.GameObjects.Systems;
 
@@ -17,7 +20,10 @@ namespace Content.Server.GameObjects.EntitySystems
 
             UpdatesAfter.Add(typeof(InteractionSystem));
             UpdatesAfter.Add(typeof(InputSystem));
+
             SubscribeLocalEvent<MoveEvent>(MoveEvent);
+            SubscribeLocalEvent<EntInsertedIntoContainerMessage>(ContainerModified);
+            SubscribeLocalEvent<EntRemovedFromContainerMessage>(ContainerModified);
         }
 
         public override void Shutdown()
@@ -29,17 +35,69 @@ namespace Content.Server.GameObjects.EntitySystems
 
         private void MoveEvent(MoveEvent ev)
         {
-            if (ev.Sender.TryGetComponent(out BuckleComponent buckle))
+            if (!ev.Sender.TryGetComponent(out BuckleComponent? buckle))
             {
-                buckle.OnMoveEvent(ev);
+                return;
+            }
+
+            var strap = buckle.BuckledTo;
+
+            if (strap == null)
+            {
+                return;
+            }
+
+            var strapPosition = strap.Owner.Transform.Coordinates.Offset(buckle.BuckleOffset);
+
+            if (ev.NewPosition.InRange(EntityManager, strapPosition, 0.2f))
+            {
+                return;
+            }
+
+            buckle.TryUnbuckle(buckle.Owner, true);
+        }
+
+        private void ContainerModified(ContainerModifiedMessage message)
+        {
+            // Not returning is necessary in case an entity has both a buckle and strap component
+            if (message.Entity.TryGetComponent(out BuckleComponent? buckle))
+            {
+                ContainerModifiedReAttach(buckle, buckle.BuckledTo);
+            }
+
+            if (message.Entity.TryGetComponent(out StrapComponent? strap))
+            {
+                foreach (var buckledEntity in strap.BuckledEntities)
+                {
+                    if (!buckledEntity.TryGetComponent(out BuckleComponent? buckled))
+                    {
+                        continue;
+                    }
+
+                    ContainerModifiedReAttach(buckled, strap);
+                }
             }
         }
 
-        public override void Update(float frameTime)
+        private void ContainerModifiedReAttach(BuckleComponent buckle, StrapComponent? strap)
         {
-            foreach (var buckle in ComponentManager.EntityQuery<BuckleComponent>())
+            if (strap == null)
             {
-                buckle.Update();
+                return;
+            }
+
+            var contained = ContainerHelpers.TryGetContainer(buckle.Owner, out var ownContainer);
+            var strapContained = ContainerHelpers.TryGetContainer(strap.Owner, out var strapContainer);
+
+            if (contained != strapContained || ownContainer != strapContainer)
+            {
+                buckle.TryUnbuckle(buckle.Owner, true);
+                return;
+            }
+
+            if (!contained)
+            {
+                buckle.ReAttach(strap);
             }
         }
     }