ArachnaeSwarm/Source/ArachnaeSwarm/ARA_HiveMind/Hediff_HiveMindMaster.cs

using System.Collections.Generic;
using System.Linq;
using RimWorld;
using Verse;

namespace ArachnaeSwarm
{
    public class Hediff_HiveMindMaster : HediffWithComps
    {
        private List<Pawn> drones = new List<Pawn>();

        public override string LabelInBrackets => drones.Count.ToString();

        public override bool ShouldRemove => false;

        public override string TipStringExtra
        {
            get
            {
                if (drones.NullOrEmpty())
                {
                    return "\n" + "ARA_TipString_NoDronesBound".Translate().Colorize(ColoredText.SubtleGrayColor);
                }

                System.Text.StringBuilder stringBuilder = new System.Text.StringBuilder();
                stringBuilder.AppendLine("\n" + "ARA_TipString_ControlledDrones".Translate().Colorize(ColoredText.TipSectionTitleColor));
                foreach (var drone in drones)
                {
                    string backstoryInfo = "";
                    if (drone.story?.TitleShort is string adulthoodTitle && !adulthoodTitle.NullOrEmpty())
                    {
                        backstoryInfo = $" ({adulthoodTitle.Colorize(ColoredText.SubtleGrayColor)})";
                    }
                    stringBuilder.AppendLine($"  - {drone.LabelShortCap}{backstoryInfo}");
                }
                return stringBuilder.ToString();
            }
        }

        public override void ExposeData()
        {
            base.ExposeData();
            Scribe_Collections.Look(ref drones, "drones", LookMode.Reference);
            if (drones == null)
            {
                drones = new List<Pawn>();
            }
        }

        public bool TryBindDrone(Pawn drone)
        {
            if (drone == null || drone.Dead || !drone.Spawned || drone.Map != this.pawn.Map)
            {
                Log.Message($"[ArachnaeSwarm] Cannot bind drone {drone?.LabelShort ?? "null"}: Invalid pawn state.");
                return false;
            }

            Hediff_HiveMindDrone droneHediff = drone.health.hediffSet.GetFirstHediffOfDef(HediffDef.Named("ARA_HiveMindDrone")) as Hediff_HiveMindDrone;
            if (droneHediff == null)
            {
                Log.Message($"[ArachnaeSwarm] Cannot bind drone {drone.LabelShort}: Does not have ARA_HiveMindDrone hediff.");
                return false;
            }

            if (droneHediff.target != null && droneHediff.target != this.pawn)
            {
                Log.Message($"[ArachnaeSwarm] Cannot bind drone {drone.LabelShort}: Already bound to another master ({droneHediff.target.LabelShort}).");
                return false;
            }
            
            if (drones.Contains(drone))
            {
                Log.Message($"[ArachnaeSwarm] Drone {drone.LabelShort} is already bound to this master.");
                return false;
            }

            droneHediff.target = this.pawn; // Set the drone's target to this master
            drones.Add(drone);
            UpdateSeverity();
            Log.Message($"[ArachnaeSwarm] Master {this.pawn.LabelShort} successfully bound drone {drone.LabelShort}.");
            return true;
        }

        public void TryBindAllAvailableDrones()
        {
            if (this.pawn?.Map == null) return;

            // First, clean up any dead, despawned, or destroyed drones from the list.
            int removedCount = drones.RemoveAll(drone => drone == null || drone.Dead || drone.Destroyed);

            // Find all pawns on the map that could be potential drones
            List<Pawn> potentialDrones = this.pawn.Map.mapPawns.AllPawnsSpawned
                .Where(p => p.health.hediffSet.HasHediff(HediffDef.Named("ARA_HiveMindDrone")))
                .ToList();

            int boundCount = 0;

            foreach (var drone in potentialDrones)
            {
                Hediff_HiveMindDrone droneHediff = drone.health.hediffSet.GetFirstHediffOfDef(HediffDef.Named("ARA_HiveMindDrone")) as Hediff_HiveMindDrone;

                // Check if the drone is unlinked (target is null) and not already in our list
                if (droneHediff != null && droneHediff.target == null && !drones.Contains(drone))
                {
                    droneHediff.target = this.pawn; // Set the drone's target to this master
                    drones.Add(drone);
                    Log.Message($"[ArachnaeSwarm] Master {this.pawn.LabelShort} automatically bound drone {drone.LabelShort}.");
                    boundCount++;
                }
            }
            
            // Update severity if anything changed (drones were added or removed)
            if (boundCount > 0 || removedCount > 0)
            {
                UpdateSeverity();
            }
        }
        
        public void DeregisterDrone(Pawn drone)
        {
            if (drones.Contains(drone))
            {
                drones.Remove(drone);
                UpdateSeverity();
            }
        }

        private void UpdateSeverity()
        {
            this.Severity = drones.Count;
        }

        public override void PostRemoved()
        {
            base.PostRemoved();
            // Kill all drones when the master hediff is removed (e.g., master dies)
            foreach (var drone in drones.ToList()) // ToList() to avoid collection modification issues
            {
                if (drone != null && !drone.Dead)
                {
                    Log.Message($"[ArachnaeSwarm] Master {pawn.LabelShort} died, killing drone {drone.LabelShort}.");
                    drone.Kill(null, this);
                }
            }
        }
    }
}