rjw_menstruation/1.4/source/RJW_Menstruation/RJW_Menstruation/HediffComps/HediffComp_PregeneratedBabies.cs

using HarmonyLib;
using RimWorld;
using rjw;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Reflection.Emit;
using Verse;

namespace RJW_Menstruation
{
    public class HediffComp_PregeneratedBabies : HediffComp
    {
        public List<Pawn> babies;
        // Unused, but can't hurt to track
        protected Dictionary<Pawn, Pawn> enzygoticSiblings;

        protected static readonly MethodInfo RandomLastName = typeof(PregnancyUtility).GetMethod("RandomLastName", BindingFlags.Static | BindingFlags.NonPublic, null, new Type[] { typeof(Pawn), typeof(Pawn), typeof(Pawn) }, null);

        public bool HasBaby
        {
            get => !babies.NullOrEmpty();            
        }

        public Pawn PopBaby()
        {
            if (babies.NullOrEmpty()) return null;

            Pawn firstBaby = babies.First();
            babies.Remove(firstBaby);
            return firstBaby;
        }

        public override void CompPostPostRemoved()
        {
            // At this point in the hediff removal process, the new hediff is already on the pawn
            // But it is possible that there is no new hediff (be it a birth, miscarrage, or dev edit)
            base.CompPostPostRemoved();

            // Send the babies from this comp to the new one
            switch (parent)
            {
                case Hediff_Pregnant hediff_Pregnant:
                    Hediff_Labor labor = (Hediff_Labor)Pawn.health.hediffSet.hediffs.Where(hediff => hediff is Hediff_Labor).MaxByWithFallback(hediff => hediff.loadID);
                    HediffComp_PregeneratedBabies laborcomp = labor?.TryGetComp<HediffComp_PregeneratedBabies>();
                    if (laborcomp == null) return;
                    laborcomp.babies = this.babies;
                    laborcomp.enzygoticSiblings = this.enzygoticSiblings;
                    break;
                case Hediff_Labor hediff_Labor:
                    Hediff_LaborPushing pushing = (Hediff_LaborPushing)Pawn.health.hediffSet.hediffs.Where(hediff => hediff is Hediff_LaborPushing).MaxByWithFallback(hediff => hediff.loadID);
                    HediffComp_PregeneratedBabies pushingcomp = pushing?.TryGetComp<HediffComp_PregeneratedBabies>();
                    if (pushingcomp == null) return;
                    pushingcomp.babies = this.babies;
                    pushingcomp.enzygoticSiblings = this.enzygoticSiblings;
                    break;
                case Hediff_LaborPushing hediff_LaborPushing:
                    // Nothing to do, the laborpushing transpiler will pick it up
                    break;
            }
        }

        public override void CompExposeData()
        {
            base.CompExposeData();
            Scribe_Collections.Look(ref babies, "babies", LookMode.Deep);
            Scribe_Collections.Look(ref enzygoticSiblings, "enzygoticSiblings", keyLookMode: LookMode.Reference, valueLookMode: LookMode.Reference);
        }


        public void AddNewBaby(Pawn mother, Pawn father)
        {
            if (babies == null) babies = new List<Pawn>();
            PawnKindDef babyPawnKind = PregnancyCommon.BabyPawnKindDecider(mother, father, true);
            PawnGenerationRequest request = new PawnGenerationRequest
                (
                    kind: babyPawnKind,
                    faction: mother.Faction,
                    allowDowned: true,
                    // fixedLastName seems not to actually do anything, as we eventually end up in PawnBioAndNameGenerator.NameResolvedFrom, which ignores its forcedLastName argument
                    fixedLastName: (string)RandomLastName.Invoke(null, new object[] { mother, mother, xxx.is_human(father) ? father : null }),
                    forceNoIdeo: true,
                    forcedEndogenes: PregnancyUtility.GetInheritedGenes(father, mother),
                    forcedXenotype: XenotypeDefOf.Baseliner,
                    developmentalStages: DevelopmentalStage.Newborn
                );
            int division = 1;
            Pawn firstbaby = null;
            while (Rand.Chance(Configurations.EnzygoticTwinsChance) && division < Configurations.MaxEnzygoticTwins) division++;
            if (division > 1 && enzygoticSiblings == null) enzygoticSiblings = new Dictionary<Pawn, Pawn>();
            for (int i = 0; i < division; i++)
            {
                Pawn baby = PawnGenerator.GeneratePawn(request);
                if (baby == null) break;
                PregnancyCommon.SetupBabyXenotype(mother, father, baby);    // Probably redundant with Biotech post-birth xenotyping
                baby.Drawer.renderer.graphics.ResolveAllGraphics();
                if (division > 1)
                {
                    if (i == 0)
                    {
                        firstbaby = baby;
                        request.FixedGender = baby.gender;
                        request.ForcedEndogenes = baby.genes?.Endogenes.Select(gene => gene.def).ToList();
                    }
                    else
                    {
                        enzygoticSiblings.Add(baby, firstbaby);

                        if (baby.story != null)
                        {
                            baby.story.headType = firstbaby.story.headType;
                            baby.story.hairDef = firstbaby.story.hairDef;
                            baby.story.HairColor = firstbaby.story.HairColor;
                            baby.story.bodyType = firstbaby.story.bodyType;
                            baby.story.furDef = firstbaby.story.furDef;
                            baby.story.skinColorOverride = firstbaby.story.skinColorOverride;
                        }

                        if (baby.genes != null)
                        {
                            baby.genes.SetXenotypeDirect(firstbaby.genes.Xenotype);
                            baby.genes.xenotypeName = firstbaby.genes.xenotypeName;
                            baby.genes.iconDef = firstbaby.genes.iconDef;
                            baby.genes.hybrid = firstbaby.genes.hybrid;
                        }

                        if (baby.IsHAR())
                            HARCompatibility.CopyHARProperties(baby, firstbaby);

                        // MultiplePregnancy calls this post-birth because RJW resets private parts
                        // So xenotype things shouldn't be shared
                        PregnancyCommon.ProcessIdenticalSibling(baby, firstbaby);
                    }
                }
                babies.Add(baby);
                // These get cleared out later, but setting the relations now will help keep track of things.
                baby.SetMother(mother);
                if (mother != father)
                {
                    if (father.gender != Gender.Female) baby.SetFather(father);
                    else baby.relations.AddDirectRelation(PawnRelationDefOf.Parent, father);
                }
            }
        }
    }

    [HarmonyPatch(typeof(PregnancyUtility), nameof(PregnancyUtility.ApplyBirthOutcome))]
    public static class ApplyBirthOutcome_PregeneratedBabies_Patch
    {
        private static Pawn GetPregeneratedBaby(PawnGenerationRequest request, Thing birtherThing)
        {
            // Don't test for the config set here. We can do it at the functions that call ApplyBirthOutcome
            // Easier to work out twins that way

            // From e.g. a vat
            if (!(birtherThing is Pawn mother) || !xxx.is_human(mother))
                return PawnGenerator.GeneratePawn(request);

            // No babies found. Could be an unmodified pregnancy
            HediffComp_PregeneratedBabies comp = mother.health.hediffSet.GetFirstHediff<Hediff_LaborPushing>()?.TryGetComp<HediffComp_PregeneratedBabies>();
            if (comp == null || !comp.HasBaby)
                return PawnGenerator.GeneratePawn(request);

            Pawn baby = comp.PopBaby();
            if (baby == null) return PawnGenerator.GeneratePawn(request);   // Shouldn't happen
            baby.ageTracker.AgeBiologicalTicks = 0;
            baby.ageTracker.AgeChronologicalTicks = 0;
            baby.babyNamingDeadline = Find.TickManager.TicksGame + GenDate.TicksPerDay;
            if (request.ForceDead) baby.Kill(null, null);
            return baby;
        }

        private static readonly MethodInfo ApplyBirthOutcome = typeof(PregnancyUtility).GetMethod(nameof(PregnancyUtility.ApplyBirthOutcome));
        private static readonly int birtherThing = ApplyBirthOutcome.GetParameters().FirstIndexOf(parameter => parameter.Name == "birtherThing" && parameter.ParameterType == typeof(Thing));
        private static readonly MethodInfo GeneratePawn = typeof(PawnGenerator).GetMethod(nameof(PawnGenerator.GeneratePawn), new Type[] {typeof (PawnGenerationRequest)});

        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
        {
            if (birtherThing < 0) throw new InvalidOperationException("Could not locate index of birtherThing");
            if (GeneratePawn?.ReturnType != typeof(Pawn)) throw new InvalidOperationException("GeneratePawn not found");
            foreach (CodeInstruction instruction in instructions)
            {
                if (instruction.Calls(GeneratePawn))
                {
                    yield return new CodeInstruction(OpCodes.Ldarg, birtherThing);
                    yield return CodeInstruction.Call(typeof(ApplyBirthOutcome_PregeneratedBabies_Patch), nameof(GetPregeneratedBaby));
                }
                else yield return instruction;
            }
        }
    }

    [HarmonyPatch(typeof(Hediff_LaborPushing), nameof(Hediff_LaborPushing.PreRemoved))]
    public static class Hediff_LaborPushing_PreRemoved_Patch
    {
        private static Thing ApplyBirthLoop(OutcomeChance outcome, float quality, Precept_Ritual ritual, List<GeneDef> genes, Pawn geneticMother, Thing birtherThing, Pawn father, Pawn doctor, LordJob_Ritual lordJobRitual, RitualRoleAssignments assignments)
        {
            if (birtherThing is Pawn mother)
            {
                HediffComp_PregeneratedBabies comp = mother.health.hediffSet.GetFirstHediff<Hediff_LaborPushing>().TryGetComp<HediffComp_PregeneratedBabies>();
                if (comp?.HasBaby ?? false)
                {
                    OutcomeChance thisOutcome = outcome;
                    Precept_Ritual precept_Ritual = (Precept_Ritual)comp.Pawn.Ideo.GetPrecept(RimWorld.PreceptDefOf.ChildBirth);
                    float birthQuality = PregnancyUtility.GetBirthQualityFor(mother);
                    do
                    {
                        Pawn baby = comp.babies[0];
                        Pawn thisFather = baby.GetFather() ?? father;
                        baby.relations.ClearAllRelations();   // To keep ApplyBirthOutcome from erroring when it tries to set up relations

                        PregnancyUtility.ApplyBirthOutcome(thisOutcome, quality, ritual, genes, geneticMother, birtherThing, thisFather, doctor, lordJobRitual, assignments);
                        // No more babies if mom dies halfway through. Unrealistic maybe, but saves a lot of headache in ApplyBirthOutcome
                        if (mother.Dead) break;
                        if (xxx.is_human(baby))
                            mother.records.Increment(xxx.CountOfBirthHuman);
                        else if (xxx.is_animal(baby))
                            mother.records.Increment(xxx.CountOfBirthAnimal);
                        thisOutcome = ((RitualOutcomeEffectWorker_ChildBirth)precept_Ritual.outcomeEffect).GetOutcome(birthQuality, null);
                    } while (comp.HasBaby);

                    // PreRemoved doesn't use the return value
                    return null;
                }    
            }

            return PregnancyUtility.ApplyBirthOutcome(outcome, quality, ritual, genes, geneticMother, birtherThing, father, doctor, lordJobRitual, assignments);
        }

        private static readonly MethodInfo ApplyBirthOutcome = typeof(PregnancyUtility).GetMethod(nameof(PregnancyUtility.ApplyBirthOutcome),
            new Type[] {typeof(OutcomeChance), typeof(float), typeof(Precept_Ritual), typeof(List<GeneDef>), typeof(Pawn), typeof(Thing), typeof(Pawn), typeof(Pawn), typeof(LordJob_Ritual), typeof(RitualRoleAssignments)});
        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
        {
            if (ApplyBirthOutcome?.ReturnType != typeof(Thing)) throw new InvalidOperationException("ApplyBirthOutcome not found");
            foreach (CodeInstruction instruction in instructions)
            {
                if (instruction.Calls(ApplyBirthOutcome))
                    yield return CodeInstruction.Call(typeof(Hediff_LaborPushing_PreRemoved_Patch), nameof(Hediff_LaborPushing_PreRemoved_Patch.ApplyBirthLoop));
                else yield return instruction;
            }
        }
    }

    [HarmonyPatch(typeof(RitualOutcomeEffectWorker_ChildBirth), nameof (RitualOutcomeEffectWorker_ChildBirth.Apply))]
    public static class Ritual_ChildBirth_Apply_Patch
    {
        private static Thing ApplyBirthLoop(OutcomeChance outcome, float quality, Precept_Ritual ritual, List<GeneDef> genes, Pawn geneticMother, Thing birtherThing, Pawn father, Pawn doctor, LordJob_Ritual lordJobRitual, RitualRoleAssignments assignments)
        {
            if (birtherThing is Pawn mother)
            {
                HediffComp_PregeneratedBabies comp = mother.health.hediffSet.GetFirstHediff<Hediff_LaborPushing>().TryGetComp<HediffComp_PregeneratedBabies>();
                if (comp?.HasBaby ?? false)
                {
                    // Much the same as the other one

                    // Don't reroll the outcome every time, I think
                    // This is all one ritual, so every baby has the same ritual outcome
                    // I don't think this will add the ritual memory every time?
                    // Though even if it does, that's probably okay. More babies more memories after all
                    do
                    {
                        Pawn baby = comp.babies[0];
                        Pawn thisFather = baby.GetFather() ?? father;
                        baby.relations.ClearAllRelations();

                        PregnancyUtility.ApplyBirthOutcome(outcome, quality, ritual, genes, geneticMother, birtherThing, thisFather, doctor, lordJobRitual, assignments);
                        if (mother.Dead) break;
                        if (xxx.is_human(baby))
                            mother.records.Increment(xxx.CountOfBirthHuman);
                        else if (xxx.is_animal(baby))
                            mother.records.Increment(xxx.CountOfBirthAnimal);
                    } while (comp.HasBaby);

                    // The ritual version doesn't use the return value, either
                    return null;
                }
            }
            return PregnancyUtility.ApplyBirthOutcome(outcome, quality, ritual, genes, geneticMother, birtherThing, father, doctor, lordJobRitual, assignments);
        }
        private static readonly MethodInfo ApplyBirthOutcome = typeof(PregnancyUtility).GetMethod(nameof(PregnancyUtility.ApplyBirthOutcome),
            new Type[] { typeof(OutcomeChance), typeof(float), typeof(Precept_Ritual), typeof(List<GeneDef>), typeof(Pawn), typeof(Thing), typeof(Pawn), typeof(Pawn), typeof(LordJob_Ritual), typeof(RitualRoleAssignments) });
        public static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
        {
            if (ApplyBirthOutcome?.ReturnType != typeof(Thing)) throw new InvalidOperationException("ApplyBirthOutcome not found");
            foreach (CodeInstruction instruction in instructions)
            {
                if (instruction.Calls(ApplyBirthOutcome))
                    yield return CodeInstruction.Call(typeof(Ritual_ChildBirth_Apply_Patch), nameof(Ritual_ChildBirth_Apply_Patch.ApplyBirthLoop));
                else yield return instruction;
            }
        }
    }

    // HAR patches ApplyBirthOutcome to produce multiple babies based on the mother's littersize. But the pregenerated babies system already makes multiple babies
    // So make it always consider the mother to have one baby
    public static class HAR_LitterSize_Undo
    {
        public static void Postfix(ref int __result, Pawn mother)
        {
            if (Configurations.PregnancySource != Configurations.PregnancyType.Biotech || !Configurations.EnableBiotechTwins) return;
            __result = 0;
            return;
        }
    }
}