ArachnaeSwarm/Source/ArachnaeSwarm/Hediff_HiveMindMaster.cs

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

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

        public override string LabelInBrackets => drones.Count.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 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);
                }
            }
        }
    }
}