muzika-gromche/MuzikaGromche/PoweredLightsAnimators.cs

using DunGen;
using HarmonyLib;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Reflection;
using UnityEngine;

namespace MuzikaGromche
{
    internal class PoweredLightsAnimators
    {
        private delegate void ManualPatch(GameObject animatorContainer);

        private readonly record struct AnimatorPatch(string AnimatorContainerPath, RuntimeAnimatorController AnimatorController, ManualPatch? ManualPatch);

        private readonly record struct TilePatch(string TileName, AnimatorPatch[] Patches)
        {
            // We are specifically looking for cloned tiles, not the original prototypes.
            public readonly string TileCloneName = $"{TileName}(Clone)";
        }

        private readonly record struct AnimatorPatchDescriptor(string AnimatorContainerPath, string AnimatorControllerAssetPath, ManualPatch? ManualPatch = null)
        {
            public AnimatorPatch Load(AssetBundle assetBundle)
            {
                var animationController = assetBundle.LoadAsset<RuntimeAnimatorController>(AnimatorControllerAssetPath)
                    ?? throw new FileNotFoundException($"RuntimeAnimatorController not found: {AnimatorControllerAssetPath}", AnimatorControllerAssetPath);
                return new(AnimatorContainerPath, animationController, ManualPatch);
            }
        }

        private readonly record struct TilePatchDescriptor(string[] TileNames, AnimatorPatchDescriptor[] Descriptors)
        {
            public TilePatchDescriptor(string TileName, AnimatorPatchDescriptor[] Descriptors)
                : this([TileName], Descriptors)
            {
            }

            public TilePatch[] Load(AssetBundle assetBundle)
            {
                var patches = Descriptors.Select(d => d.Load(assetBundle)).ToArray();
                return [.. TileNames.Select(tileName => new TilePatch(tileName, patches))];
            }
        }

        private static IDictionary<string, TilePatch> Patches = new Dictionary<string, TilePatch>();

        public static void Load()
        {
            const string BundleFileName = "muzikagromche_poweredlightsanimators";
            string bundlePath = Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), BundleFileName);
            var assetBundle = AssetBundle.LoadFromFile(bundlePath)
                ?? throw new NullReferenceException("Failed to load bundle");

            const string BasePath = "Assets/LethalCompany/Mods/MuzikaGromche/AnimatorControllers/";

            const string PointLight4 = $"{BasePath}Point Light (4) (Patched).controller";
            const string MineshaftSpotlight = $"{BasePath}MineshaftSpotlight (Patched).controller";
            const string LightbulbsLine = $"{BasePath}lightbulbsLineMesh (Patched).controller";
            const string CeilingFan = $"{BasePath}CeilingFan (originally GameObject) (Patched).controller";
            const string LEDHangingLight = $"{BasePath}LEDHangingLight (Patched).controller";

            TilePatchDescriptor[] descriptors =
            [
                // any version
                new("KitchenTile", [
                    new("PoweredLightTypeB", PointLight4),
                    new("PoweredLightTypeB (1)", PointLight4),
                ]),
                // < v70
                new("ManorStartRoom", [
                    new("ManorStartRoom/Chandelier/PoweredLightTypeB (1)", PointLight4),
                    new("ManorStartRoom/Chandelier2/PoweredLightTypeB", PointLight4),
                ]),
                // v70+
                new("ManorStartRoomSmall", [
                    new("ManorStartRoomMesh/Chandelier/PoweredLightTypeB (1)", PointLight4),
                    new("ManorStartRoomMesh/Chandelier2/PoweredLightTypeB", PointLight4),
                ]),
                new("BirthdayRoomTile", [
                    new("Lights/MineshaftSpotlight", MineshaftSpotlight),
                ]),
                new("BathroomTileContainer", [
                    new("MineshaftSpotlight", MineshaftSpotlight),
                    new("LightbulbLine/lightbulbsLineMesh", LightbulbsLine),
                ]),
                new(["BedroomTile", "BedroomTileB"], [
                    new("CeilingFanAnimContainer", CeilingFan),
                    new("MineshaftSpotlight (1)", MineshaftSpotlight),
                ]),
                new("GarageTile", [
                    new("HangingLEDBarLight (3)", LEDHangingLight),
                    new("HangingLEDBarLight (4)", LEDHangingLight,
                        // This HangingLEDBarLight's IndirectLight is wrongly named, so animator couldn't find it
                        RenameGameObjectPatch("IndirectLight (1)", "IndirectLight")),
                ]),
                new("PoolTile", [
                    new("PoolLights/HangingLEDBarLight", LEDHangingLight),
                    new("PoolLights/HangingLEDBarLight (4)", LEDHangingLight),
                    new("PoolLights/HangingLEDBarLight (5)", LEDHangingLight),
                ]),
            ];

            Patches = descriptors
                .SelectMany(d => d.Load(assetBundle))
                .ToDictionary(d => d.TileCloneName, d => d);
        }

        public static void Patch(Tile tile)
        {
            if (tile == null)
            {
                throw new ArgumentNullException(nameof(tile));
            }

            if (Patches.TryGetValue(tile.gameObject.name, out var tilePatch))
            {
                foreach (var patch in tilePatch.Patches)
                {
                    Transform animationContainerTransform = tile.gameObject.transform.Find(patch.AnimatorContainerPath);
                    if (animationContainerTransform == null)
                    {
#if DEBUG
                        throw new NullReferenceException($"{tilePatch.TileName}/{patch.AnimatorContainerPath} Animation Container not found");
#endif
#pragma warning disable CS0162 // Unreachable code detected
                        continue;
#pragma warning restore CS0162 // Unreachable code detected
                    }

                    var animator = animationContainerTransform.gameObject.GetComponent<Animator>();
                    if (animator == null)
                    {
#if DEBUG
                        throw new NullReferenceException($"{tilePatch.TileName}/{patch.AnimatorContainerPath} Animation Component not found");
#endif
#pragma warning disable CS0162 // Unreachable code detected
                        continue;
#pragma warning restore CS0162 // Unreachable code detected
                    }

                    patch.ManualPatch?.Invoke(animationContainerTransform.gameObject);

                    animator.runtimeAnimatorController = patch.AnimatorController;
                    Debug.Log($"{nameof(MuzikaGromche)} {nameof(PoweredLightsAnimatorsPatch)} {tilePatch.TileName}/{patch.AnimatorContainerPath}: Replaced animator controller");
                }
            }
        }

        private static ManualPatch RenameGameObjectPatch(string relativePath, string newName) => animatorContainer =>
        {
            var targetObject = animatorContainer.transform.Find(relativePath)?.gameObject;
            if (targetObject == null)
            {
#if DEBUG
                throw new NullReferenceException($"{animatorContainer.name}/{relativePath}: GameObject not found!");
#endif
#pragma warning disable CS0162 // Unreachable code detected
                return;
#pragma warning restore CS0162 // Unreachable code detected
            }
            targetObject.name = newName;
            Debug.Log($"{nameof(MuzikaGromche)} {nameof(PoweredLightsAnimatorsPatch)} {animatorContainer.name}/{relativePath}: Renamed GameObject");
        };
    }

    [HarmonyPatch(typeof(Tile))]
    internal class PoweredLightsAnimatorsPatch
    {
        [HarmonyPatch("AddTriggerVolume")]
        [HarmonyPostfix]
        public static void OnAddTriggerVolume(Tile __instance)
        {
            PoweredLightsAnimators.Patch(__instance);
        }
    }

    [HarmonyPatch(typeof(RoundManager))]
    static class AllPoweredLightsPatch
    {
        // Vanilla method assumes that GameObjects with tag "PoweredLight" only contain a single Light component each.
        // This is, however, not true for certains double-light setups, such as:
        // - double PointLight (even though one of them is 'Point' another is 'Spot') inside CeilingFanAnimContainer in BedroomTile/BedroomTileB;
        // - MineshaftSpotlight when it has not only `Point Light` but also `IndirectLight` in BirthdayRoomTile;
        // - (maybe more?)
        // In order to fix that, replace singular GetComponentInChildren<Light> with plural GetComponentsInChildren<Light> version.
        [HarmonyPatch(nameof(RoundManager.RefreshLightsList))]
        [HarmonyPrefix]
        private static bool OnRefreshLightsList(RoundManager __instance)
        {
            RefreshLightsListPatched(__instance);
            // Skip the original method
            return false;
        }

        private static void RefreshLightsListPatched(RoundManager self)
        {
            // Reusable list to reduce allocations
            List<Light> lights = [];

            self.allPoweredLights.Clear();
            self.allPoweredLightsAnimators.Clear();

            GameObject[] gameObjects = GameObject.FindGameObjectsWithTag("PoweredLight");
            if (gameObjects == null)
            {
                return;
            }

            foreach (var gameObject in gameObjects)
            {
                Animator animator = gameObject.GetComponentInChildren<Animator>();
                if (!(animator == null))
                {
                    self.allPoweredLightsAnimators.Add(animator);
                    // Patched section: Use list instead of singular GetComponentInChildren<Light>
                    gameObject.GetComponentsInChildren(includeInactive: true, lights);
                    self.allPoweredLights.AddRange(lights);
                }
            }
            foreach (var animator in self.allPoweredLightsAnimators)
            {
                animator.SetFloat("flickerSpeed", UnityEngine.Random.Range(0.6f, 1.4f));
            }
        }
    }
}