From f5a9f91de40921cf8746c7194116135dfb024956 Mon Sep 17 00:00:00 2001
From: DrSmugleaf <DrSmugleaf@users.noreply.github.com>
Date: Wed, 8 Jul 2020 15:30:48 +0200
Subject: [PATCH] Fix not staying buckled to a strap when both the buckled and
 the strap go into the same container

---
 .../Components/Mobs/BuckleComponent.cs        | 79 +++++++++++++------
 .../GameObjects/EntitySystems/BuckleSystem.cs | 34 ++++++++
 2 files changed, 90 insertions(+), 23 deletions(-)
 create mode 100644 Content.Server/GameObjects/EntitySystems/BuckleSystem.cs

diff --git a/Content.Server/GameObjects/Components/Mobs/BuckleComponent.cs b/Content.Server/GameObjects/Components/Mobs/BuckleComponent.cs
index 71a7cf55c0..bc309a7240 100644
--- a/Content.Server/GameObjects/Components/Mobs/BuckleComponent.cs
+++ b/Content.Server/GameObjects/Components/Mobs/BuckleComponent.cs
@@ -72,6 +72,8 @@ namespace Content.Server.GameObjects.Components.Mobs
         [ViewVariables]
         protected override bool Buckled => BuckledTo != null;
 
+        public bool ContainerChanged { get; private set; }
+
         [ViewVariables]
         public int Size => _size;
 
@@ -86,6 +88,30 @@ namespace Content.Server.GameObjects.Components.Mobs
             }
         }
 
+        private void ReAttach(StrapComponent strap)
+        {
+            var ownTransform = Owner.Transform;
+            var strapTransform = strap.Owner.Transform;
+
+            ownTransform.GridPosition = strapTransform.GridPosition;
+            ownTransform.AttachParent(strapTransform);
+
+            switch (strap.Position)
+            {
+                case StrapPosition.None:
+                    ownTransform.WorldRotation = strapTransform.WorldRotation;
+                    break;
+                case StrapPosition.Stand:
+                    StandingStateHelper.Standing(Owner);
+                    ownTransform.WorldRotation = strapTransform.WorldRotation;
+                    break;
+                case StrapPosition.Down:
+                    StandingStateHelper.Down(Owner);
+                    ownTransform.WorldRotation = Angle.South;
+                    break;
+            }
+        }
+
         private bool TryBuckle(IEntity user, IEntity to)
         {
             if (user == null || user == to)
@@ -176,26 +202,7 @@ namespace Content.Server.GameObjects.Components.Mobs
                 appearance.SetData(BuckleVisuals.Buckled, true);
             }
 
-            var ownTransform = Owner.Transform;
-            var strapTransform = strap.Owner.Transform;
-
-            ownTransform.GridPosition = strapTransform.GridPosition;
-            ownTransform.AttachParent(strapTransform);
-
-            switch (strap.Position)
-            {
-                case StrapPosition.None:
-                    ownTransform.WorldRotation = strapTransform.WorldRotation;
-                    break;
-                case StrapPosition.Stand:
-                    StandingStateHelper.Standing(Owner);
-                    ownTransform.WorldRotation = strapTransform.WorldRotation;
-                    break;
-                case StrapPosition.Down:
-                    StandingStateHelper.Down(Owner);
-                    ownTransform.WorldRotation = Angle.South;
-                    break;
-            }
+            ReAttach(strap);
 
             BuckledTo = strap;
             BuckleStatus();
@@ -283,14 +290,38 @@ namespace Content.Server.GameObjects.Components.Mobs
             return TryBuckle(user, to);
         }
 
-        private void InsertIntoContainer(EntInsertedIntoContainerMessage message)
+        private void InsertIntoContainer(ContainerModifiedMessage message)
         {
             if (message.Entity != Owner)
             {
                 return;
             }
 
-            TryUnbuckle(Owner, true);
+            ContainerChanged = true;
+        }
+
+        public void Update()
+        {
+            if (!ContainerChanged || BuckledTo == null)
+            {
+                return;
+            }
+
+            var contained = ContainerHelpers.TryGetContainer(Owner, out var ownContainer);
+            var strapContained = ContainerHelpers.TryGetContainer(BuckledTo.Owner, out var strapContainer);
+
+            if (contained != strapContained || ownContainer != strapContainer)
+            {
+                TryUnbuckle(Owner, true);
+                return;
+            }
+
+            if (!contained && !strapContained)
+            {
+                ReAttach(BuckledTo);
+            }
+
+            ContainerChanged = false;
         }
 
         public override void ExposeData(ObjectSerializer serializer)
@@ -311,6 +342,7 @@ namespace Content.Server.GameObjects.Components.Mobs
             base.Initialize();
 
             _entityManager.EventBus.SubscribeEvent<EntInsertedIntoContainerMessage>(EventSource.Local, this, InsertIntoContainer);
+            _entityManager.EventBus.SubscribeEvent<EntRemovedFromContainerMessage>(EventSource.Local, this, InsertIntoContainer);
         }
 
         protected override void Startup()
@@ -331,7 +363,8 @@ namespace Content.Server.GameObjects.Components.Mobs
                 strap.Remove(this);
             }
 
-            BuckledTo = null;
+            TryUnbuckle(Owner, true);
+
             _buckleTime = default;
             BuckleStatus();
         }
diff --git a/Content.Server/GameObjects/EntitySystems/BuckleSystem.cs b/Content.Server/GameObjects/EntitySystems/BuckleSystem.cs
new file mode 100644
index 0000000000..133e8eb652
--- /dev/null
+++ b/Content.Server/GameObjects/EntitySystems/BuckleSystem.cs
@@ -0,0 +1,34 @@
+using Content.Server.GameObjects.Components.Mobs;
+using Content.Server.GameObjects.EntitySystems.Click;
+using Robust.Server.GameObjects.EntitySystems;
+using Robust.Shared.GameObjects;
+using Robust.Shared.GameObjects.Systems;
+
+namespace Content.Server.GameObjects.EntitySystems
+{
+    public class BuckleSystem : EntitySystem
+    {
+        public override void Initialize()
+        {
+            base.Initialize();
+
+            EntityQuery = new TypeEntityQuery(typeof(BuckleComponent));
+
+            UpdatesAfter.Add(typeof(InteractionSystem));
+            UpdatesAfter.Add(typeof(InputSystem));
+        }
+
+        public override void Update(float frameTime)
+        {
+            foreach (var entity in RelevantEntities)
+            {
+                if (!entity.TryGetComponent(out BuckleComponent buckle))
+                {
+                    continue;
+                }
+
+                buckle.Update();
+            }
+        }
+    }
+}