11

Source/ArachnaeSwarm/Hediffs/ARA_GestaltNode/Pawn_GestaltTracker.cs

@@ -0,0 +1,194 @@
+using ArachnaeSwarm;
+using RimWorld;
+using System.Collections.Generic;
+using System.Linq;
+using UnityEngine;
+using Verse;
+
+namespace ArachnaeSwarm {
+    public class Pawn_GestaltTracker : IExposable
+    {
+        private Pawn pawn;
+        private List<GestaltControlGroup> controlGroups = new List<GestaltControlGroup>();
+        private List<Pawn> controlledPawns = new List<Pawn>();
+
+        public Pawn Pawn => pawn;
+        public List<GestaltControlGroup> ControlGroups => controlGroups;
+        public List<Pawn> ControlledPawns => controlledPawns;
+
+        public int UsedBandwidth => (int)ControlledPawns.Sum(p => p.GetStatValue(StatDefOf.BandwidthCost));
+        public int TotalBandwidth => (int)pawn.GetStatValue(StatDefOf.MechBandwidth);
+        public int TotalAvailableControlGroups => (int)pawn.GetStatValue(StatDefOf.MechControlGroups);
+
+        public AcceptanceReport CanControlPawns
+        {
+            get
+            {
+                if (pawn.Downed)
+                    return "GestaltControllerDowned".Translate(pawn.Named("PAWN"));
+                if (pawn.IsPrisoner)
+                    return "GestaltControllerImprisoned".Translate(pawn.Named("PAWN"));
+                if (!pawn.Spawned)
+                    return false;
+                if (pawn.InMentalState)
+                    return "GestaltControllerMentalState".Translate(pawn.Named("PAWN"), pawn.MentalStateDef.Named("MENTALSTATE"));
+                return true;
+            }
+        }
+
+        public Pawn_GestaltTracker() { }
+
+        public Pawn_GestaltTracker(Pawn pawn)
+        {
+            this.pawn = pawn;
+            Notify_ControlGroupAmountChanged();
+        }
+
+        public bool CanControlPawn(Pawn targetPawn)
+        {
+            if (targetPawn.GetOverlord() != null)
+                return false;
+
+            float bandwidthCost = targetPawn.GetStatValue(StatDefOf.BandwidthCost);
+            return UsedBandwidth + bandwidthCost <= TotalBandwidth;
+        }
+
+        public GestaltControlGroup GetControlGroup(Pawn targetPawn)
+        {
+            foreach (GestaltControlGroup group in controlGroups)
+            {
+                if (group.AssignedPawns.Contains(targetPawn))
+                    return group;
+            }
+            return null;
+        }
+
+        public void AssignPawnToControlGroup(Pawn targetPawn)
+        {
+            if (controlGroups.Count == 0)
+            {
+                Notify_ControlGroupAmountChanged();
+                if (controlGroups.Count == 0)
+                {
+                    Log.Warning("Wants to assign pawn to a control group, but there are none!");
+                    return;
+                }
+            }
+
+            // 分配到第一个可用的控制组
+            controlGroups[0].Assign(targetPawn);
+            Notify_BandwidthChanged();
+        }
+
+        public void UnassignPawnFromAnyControlGroup(Pawn targetPawn)
+        {
+            foreach (GestaltControlGroup group in controlGroups)
+            {
+                group.TryUnassign(targetPawn);
+            }
+        }
+
+        public void Notify_ControlGroupAmountChanged()
+        {
+            int availableGroups = TotalAvailableControlGroups;
+
+            // 添加缺少的控制组
+            while (controlGroups.Count < availableGroups)
+            {
+                controlGroups.Add(new GestaltControlGroup(this));
+            }
+
+            // 移除多余的控制组
+            while (controlGroups.Count > availableGroups)
+            {
+                GestaltControlGroup groupToRemove = controlGroups[controlGroups.Count - 1];
+
+                if (!groupToRemove.AssignedPawns.NullOrEmpty())
+                {
+                    // 将pawn重新分配到其他组
+                    foreach (Pawn pawn in groupToRemove.AssignedPawns)
+                    {
+                        if (controlGroups.Count > 1)
+                        {
+                            controlGroups[0].Assign(pawn);
+                        }
+                    }
+                }
+
+                controlGroups.RemoveAt(controlGroups.Count - 1);
+            }
+        }
+
+        public void Notify_BandwidthChanged()
+        {
+            controlledPawns.Clear();
+
+            // 收集所有被控制的pawn
+            foreach (GestaltControlGroup group in controlGroups)
+            {
+                controlledPawns.AddRange(group.AssignedPawns.Where(p => !controlledPawns.Contains(p)));
+            }
+
+            // 检查带宽限制，移除超出带宽的pawn
+            List<Pawn> toRemove = new List<Pawn>();
+            float currentBandwidth = 0f;
+
+            foreach (Pawn controlledPawn in controlledPawns)
+            {
+                float cost = controlledPawn.GetStatValue(StatDefOf.BandwidthCost);
+                if (currentBandwidth + cost <= TotalBandwidth)
+                {
+                    currentBandwidth += cost;
+                }
+                else
+                {
+                    toRemove.Add(controlledPawn);
+                }
+            }
+
+            foreach (Pawn pawnToRemove in toRemove)
+            {
+                UnassignPawnFromAnyControlGroup(pawnToRemove);
+                pawnToRemove.relations.RemoveDirectRelation(ARA_PawnRelationDefOf.ARA_GestaltOverseer, this.pawn);
+            }
+        }
+
+        public void Notify_ApparelChanged()
+        {
+            Notify_BandwidthChanged();
+            Notify_ControlGroupAmountChanged();
+        }
+
+        public void Notify_HediffStateChange(Hediff hediff)
+        {
+            if (hediff != null)
+            {
+                Notify_BandwidthChanged();
+                Notify_ControlGroupAmountChanged();
+            }
+        }
+
+        public IEnumerable<Gizmo> GetGizmos()
+        {
+            yield return new GestaltBandwidthGizmo(this);
+
+            foreach (GestaltControlGroup group in controlGroups)
+            {
+                yield return new GestaltControlGroupGizmo(group);
+            }
+        }
+
+        public void ExposeData()
+        {
+            Scribe_Collections.Look(ref controlGroups, "controlGroups", LookMode.Deep, this);
+            Scribe_Collections.Look(ref controlledPawns, "controlledPawns", LookMode.Reference);
+
+            if (Scribe.mode == LoadSaveMode.PostLoadInit)
+            {
+                controlledPawns?.RemoveAll(x => x == null);
+                if (controlGroups == null)
+                    controlGroups = new List<GestaltControlGroup>();
+            }
+        }
+    }
+}