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
{
    static class PoweredLightsAnimators
    {
        private const string PoweredLightTag = "PoweredLight";

        private delegate void ManualPatch(GameObject animatorContainer);

        private readonly record struct AnimatorPatch(
            string AnimatorContainerPath,
            RuntimeAnimatorController AnimatorController,
            bool AddTagAndAnimator,
            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,
            bool AddTagAndAnimator = false,
            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, AddTagAndAnimator, 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>();

        private static AudioClip AudioClipOn = null!;
        private static AudioClip AudioClipOff = null!;
        private static AudioClip AudioClipFlicker = null!;

        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");

            AudioClipOn = assetBundle.LoadAsset<AudioClip>("Assets/LethalCompany/Mods/MuzikaGromche/AudioClips/LightOn.ogg");
            AudioClipOff = assetBundle.LoadAsset<AudioClip>("Assets/LethalCompany/Mods/MuzikaGromche/AudioClips/LightOff.ogg");
            AudioClipFlicker = assetBundle.LoadAsset<AudioClip>("Assets/LethalCompany/Mods/MuzikaGromche/AudioClips/LightFlicker.ogg");

            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";
            const string MineshaftStartTileSpotlight = $"{BasePath}MineshaftStartTileSpotlight (New).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("NarrowHallwayTile2x2", [
                    new("MineshaftSpotlight (1)", MineshaftSpotlight),
                    new("MineshaftSpotlight (2)", MineshaftSpotlight),
                ]),
                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
                        ManualPatch: RenameGameObjectPatch("IndirectLight (1)", "IndirectLight")),
                ]),
                new("PoolTile", [
                    new("PoolLights/HangingLEDBarLight", LEDHangingLight),
                    new("PoolLights/HangingLEDBarLight (4)", LEDHangingLight),
                    new("PoolLights/HangingLEDBarLight (5)", LEDHangingLight),
                ]),
                new("MineshaftStartTile", [
                    new("Cylinder.001 (1)", MineshaftStartTileSpotlight, AddTagAndAnimator: true),
                    new("Cylinder.001 (2)", MineshaftStartTileSpotlight, AddTagAndAnimator: true),
                ]),
            ];

            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
                    }

                    GameObject animationContainer = animationContainerTransform.gameObject;
                    Animator animator = animationContainer.GetComponent<Animator>();

                    if (patch.AddTagAndAnimator)
                    {
                        animationContainer.tag = PoweredLightTag;
                        if (animator == null)
                        {
                            animator = animationContainer.AddComponent<Animator>();
                        }
                        if (!animationContainer.TryGetComponent<PlayAudioAnimationEvent>(out var audioScript))
                        {
                            audioScript = animationContainer.AddComponent<PlayAudioAnimationEvent>();
                            audioScript.audioClip = AudioClipOn;
                            audioScript.audioClip2 = AudioClipOff;
                            audioScript.audioClip3 = AudioClipFlicker;
                        }
                        if (!animationContainer.TryGetComponent<AudioSource>(out var audioSource))
                        {
                            // Copy from an existing AudioSource of another light animator
                            var otherSource = tile.gameObject.GetComponentInChildren<AudioSource>();
                            if (otherSource != null)
                            {
                                audioSource = animationContainer.AddComponent<AudioSource>();
                                audioSource.spatialBlend = 1;
                                audioSource.playOnAwake = false;
                                audioSource.outputAudioMixerGroup = otherSource.outputAudioMixerGroup;
                                audioSource.spread = otherSource.spread;
                                audioSource.rolloffMode = otherSource.rolloffMode;
                                audioSource.maxDistance = otherSource.maxDistance;
                                audioSource.SetCustomCurve(AudioSourceCurveType.CustomRolloff, otherSource.GetCustomCurve(AudioSourceCurveType.CustomRolloff));
                            }
                        }
                    }

                    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(animationContainer);

                    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))]
    static class PoweredLightsAnimatorsPatch
    {
        [HarmonyPatch(nameof(Tile.AddTriggerVolume))]
        [HarmonyPostfix]
        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]
        static bool OnRefreshLightsList(RoundManager __instance)
        {
            RefreshLightsListPatched(__instance);
            // Skip the original method
            return false;
        }

        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));
            }
        }
    }
}