muzika-gromche/MuzikaGromche/DiscoBallManager.cs

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

namespace MuzikaGromche
{
    public static class DiscoBallManager
    {
        // A struct holding a disco ball container object and the name of a tile for which it was designed.
        private readonly record struct TilePatch(string TileName, GameObject DiscoBallContainer)
        {
            // We are specifically looking for cloned tiles, not the original prototypes.
            public readonly string TileCloneName = $"{TileName}(Clone)";
        }

        private static TilePatch[] Patches = [];

        private static readonly List<GameObject> CachedDiscoBalls = [];
        private static readonly List<Animator> CachedDiscoBallAnimators = [];

        private static readonly string[] AnimatorContainersNames = [
            "DiscoBallProp/AnimContainer",
            "DiscoBallProp1/AnimContainer",
            "DiscoBallProp2/AnimContainer",
            "DiscoBallProp3/AnimContainer",
            "DiscoBallProp4/AnimContainer",
            "DiscoBallProp5/AnimContainer",
        ];

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

            (string PrefabPath, string TileName)[] patchDescriptors =
            [
                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerManor.prefab", "ManorStartRoomSmall"),
                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerManorOLD.prefab", "ManorStartRoom"),
                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerFactory.prefab", "StartRoom"),
                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerMineShaft.prefab", "MineshaftStartTile"),
                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerLargeForkTileB.prefab", "LargeForkTileB"),
                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerBirthdayRoomTile.prefab", "BirthdayRoomTile"),
            ];

            Patches = [.. patchDescriptors.Select(d =>
                new TilePatch(d.TileName, assetBundle.LoadAsset<GameObject>(d.PrefabPath))
            )];
        }

        internal static void Patch(Tile tile)
        {
            var query = from patch in Patches
                        where tile.gameObject.name == patch.TileCloneName
                        select patch;

            // Should be just one, but FirstOrDefault() isn't usable with structs
            foreach (var patch in query)
            {
                Patch(tile, patch);
            }
        }

        static void Patch(Tile tile, TilePatch patch)
        {
            var discoBall = UnityEngine.Object.Instantiate(patch.DiscoBallContainer, tile.transform);
            if (discoBall == null)
            {
                return;
            }

            foreach (var animator in FindDiscoBallAnimators(discoBall))
            {
                CachedDiscoBallAnimators.Add(animator);
            }
            CachedDiscoBalls.Add(discoBall);
            discoBall.SetActive(false);

            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Patched tile '{tile.gameObject.name}'");
        }

        static IEnumerable<Animator> FindDiscoBallAnimators(GameObject discoBall)
        {
            foreach (var animatorContainerName in AnimatorContainersNames)
            {
                var transform = discoBall.transform.Find(animatorContainerName);
                if (transform == null)
                {
                    // Not all prefabs have all possible animators, and it's OK
                    continue;
                }

                var animator = transform.gameObject?.GetComponent<Animator>();
                if (animator == null)
                {
                    // This would be weird
                    continue;
                }

                yield return animator;
            }
        }

        public static void Toggle(bool on)
        {
            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Toggle {(on ? "ON" : "OFF")} {CachedDiscoBallAnimators.Count} animators");

            foreach (var discoBall in CachedDiscoBalls)
            {
                discoBall.SetActive(true);
            }
            foreach (var animator in CachedDiscoBallAnimators)
            {
                animator?.SetBool("on", on);
            }
        }

        public static void Enable()
        {
            Toggle(true);
        }

        public static void Disable()
        {
            Toggle(false);
        }

        internal static void Clear()
        {
            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Clearing {CachedDiscoBalls.Count} disco balls & {CachedDiscoBallAnimators.Count} animators");
            CachedDiscoBallAnimators.Clear();
            CachedDiscoBalls.Clear();
        }
    }

    [HarmonyPatch(typeof(Tile))]
    static class DiscoBallTilePatch
    {
        [HarmonyPatch(nameof(Tile.AddTriggerVolume))]
        [HarmonyPostfix]
        static void OnAddTriggerVolume(Tile __instance)
        {
            DiscoBallManager.Patch(__instance);
        }
    }

    [HarmonyPatch(typeof(RoundManager))]
    static class DiscoBallDespawnPatch
    {
        [HarmonyPatch(nameof(RoundManager.DespawnPropsAtEndOfRound))]
        [HarmonyPatch(nameof(RoundManager.OnDestroy))]
        [HarmonyPrefix]
        static void OnDestroy(RoundManager __instance)
        {
            var _ = __instance;
            DiscoBallManager.Clear();
        }
    }
}