ShrinkModFramework/Runtime/Loading/ShrinkModLoader.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using ShrinkEventBus;
using UnityEngine;

namespace ShrinkModFramework
{
    public static class ShrinkModLoader
    {
        private sealed class DiscoveredMod
        {
            public ShrinkModInfo Info;
            public Type EntryType;
        }

        private static readonly Dictionary<string, ShrinkModHandle> LoadedMods = new(StringComparer.Ordinal);
        private static readonly List<ShrinkModHandle> LoadSequence = new();
        private static readonly ShrinkModRegistryManager RegistryManager = new();
        private static string[] _authorizedDllPaths = Array.Empty<string>();
        private static bool _hostComposed;

        /// <summary>可选解析策略：请求方程序集（可能为 null）、请求名称、候选授权 DLL 的绝对路径。
        /// 只约束本框架 AssemblyResolve 回调，不是 CLR 已加载程序集的访问沙箱。</summary>
        public static Func<Assembly, AssemblyName, string, bool> AssemblyResolvePolicy { get; set; }

        public static IReadOnlyDictionary<string, ShrinkModHandle> LoadComposed(
            ShrinkModFrameworkSettings settings, IReadOnlyList<ShrinkModComponentSource> sources)
        {
            if (sources == null) throw new ArgumentNullException(nameof(sources));
            if (LoadedMods.Count != 0) throw new InvalidOperationException("Cannot compose over legacy loaded mods.");
            var previous = Mods.ToDictionary(pair => pair.Key, pair => pair.Value.Generation, StringComparer.Ordinal);
            var result = ShrinkModCordisRuntime.ApplyComposedAsync(settings, sources).GetAwaiter().GetResult();
            _hostComposed = true;
            PublishReady(result, previous);
            return result;
        }


        public static bool IsLoaded { get; private set; }
        public static IReadOnlyDictionary<string, ShrinkModHandle> Mods =>
            ShrinkModCordisRuntime.IsInitialized ? ShrinkModCordisRuntime.Mods : LoadedMods;

        public static event Action<ShrinkModHandle> OnModReady;
        public static event Action<IReadOnlyDictionary<string, ShrinkModHandle>> OnAllModsReady;

        public static IReadOnlyDictionary<string, ShrinkModHandle> LoadAll(ShrinkModFrameworkSettings settings = null)
            => LoadInternal(settings, Array.Empty<string>());

        /// <summary>
        /// 加载工程内模组，并且只加载调用方显式授权的外部 DLL 路径。
        /// 路径集合是完整白名单，不会递归扫描模组目录；后续 revision 刷新也只复用该白名单。
        /// </summary>
        public static IReadOnlyDictionary<string, ShrinkModHandle> LoadAuthorized(
            ShrinkModFrameworkSettings settings,
            IEnumerable<string> authorizedDllPaths)
        {
            if (authorizedDllPaths == null)
                throw new ArgumentNullException(nameof(authorizedDllPaths));

            return LoadInternal(settings, authorizedDllPaths);
        }

        private static IReadOnlyDictionary<string, ShrinkModHandle> LoadInternal(
            ShrinkModFrameworkSettings settings,
            IEnumerable<string> authorizedDllPaths)
        {
            if (_hostComposed) throw new InvalidOperationException("Use LoadComposed to update a host composition.");
            settings ??= ShrinkModFrameworkSettings.Instance;
            _authorizedDllPaths = authorizedDllPaths
                .Where(path => !string.IsNullOrWhiteSpace(path))
                .Select(path => path.Trim())
                .Distinct(StringComparer.OrdinalIgnoreCase)
                .ToArray();

            if (settings == null || settings.useContextHost)
                return ApplyContextComposition(settings, _authorizedDllPaths);

            if (IsLoaded)
            {
                Debug.Log("[ShrinkModLoader] 模组已装载，跳过重复装载。");
                return Mods;
            }

            var verboseLogging = settings == null || settings.verboseLogging;

            try
            {
                ShrinkModNetworkManager.Configure(settings == null || settings.enableNetworkSync, verboseLogging);
                ShrinkExternalModAssemblyLoader.LoadExternalAssemblies(
                    settings, verboseLogging, _authorizedDllPaths);

                var discovered = DiscoverMods(settings);
                var ordered = ResolveLoadOrder(discovered, allowExistingLoadedDependencies: true);

                foreach (var discoveredMod in ordered)
                {
                    var handle = InstantiateAndTrack(discoveredMod, settings, verboseLogging);
                    LoadSequence.Add(handle);
                }

                RunLifecycle(LoadSequence, verboseLogging);

                IsLoaded = true;
                OnAllModsReady?.Invoke(LoadedMods);
                Debug.Log($"[ShrinkModLoader] 模组装载完成，共 {LoadedMods.Count} 个。");
                return LoadedMods;
            }
            catch
            {
                ResetForTesting();
                throw;
            }
        }

        public static IReadOnlyDictionary<string, ShrinkModHandle> LoadNewExternalMods(ShrinkModFrameworkSettings settings = null)
        {
            if (_hostComposed) return Mods; // A watcher cannot reconstruct host manifest metadata.
            settings ??= ShrinkModFrameworkSettings.Instance;
            if (settings == null || settings.useContextHost)
                return ApplyContextComposition(settings, _authorizedDllPaths);
            if (!IsLoaded)
                return LoadInternal(settings, _authorizedDllPaths);

            var verboseLogging = settings == null || settings.verboseLogging;
            ShrinkExternalModAssemblyLoader.LoadExternalAssemblies(
                settings, verboseLogging, _authorizedDllPaths);

            var discovered = DiscoverMods(settings)
                .Where(mod => !LoadedMods.ContainsKey(mod.Info.ModId))
                .ToList();

            if (discovered.Count == 0)
                return LoadedMods;

            ValidateNewMods(discovered);
            var ordered = ResolveLoadOrder(discovered, allowExistingLoadedDependencies: true);
            var newHandles = new List<ShrinkModHandle>(ordered.Count);

            try
            {
                foreach (var discoveredMod in ordered)
                {
                    var handle = InstantiateAndTrack(discoveredMod, settings, verboseLogging);
                    LoadSequence.Add(handle);
                    newHandles.Add(handle);
                }

                RunLifecycle(newHandles, verboseLogging);
                return LoadedMods;
            }
            catch
            {
                foreach (var handle in newHandles)
                {
                    handle.EventBinding?.Dispose();
                    handle.EventBinding = null;
                    EventBus.RemoveBus(ShrinkBusKey.Mod(handle.Info.ModId));
                    LoadedMods.Remove(handle.Info.ModId);
                    LoadSequence.Remove(handle);
                }

                throw;
            }
        }

        public static bool TryGetMod(string modId, out ShrinkModHandle handle)
        {
            if (!string.IsNullOrWhiteSpace(modId))
                return Mods.TryGetValue(modId, out handle);

            handle = null;
            return false;
        }

        public static IReadOnlyShrinkModRegistry<T> GetOrCreateRegistry<T>(string name)
            => ShrinkModCordisRuntime.IsInitialized
                ? ShrinkModCordisRuntime.GetOrCreateRegistry<T>(name)
                : RegistryManager.GetOrCreateRegistry<T>(name);

        internal static void ResetForTesting()
        {
            ShrinkModCordisRuntime.ResetForTesting();
            DisposeLegacyEventBindings();
            LoadedMods.Clear();
            LoadSequence.Clear();
            RegistryManager.Clear();
            IsLoaded = false;
            OnModReady = null;
            OnAllModsReady = null;
            ShrinkModNetworkManager.ResetForDomainReload();
            ShrinkExternalModAssemblyLoader.ResetForTesting();
            ShrinkHarmonyPatchService.ResetForTesting();
            _authorizedDllPaths = Array.Empty<string>();
            _hostComposed = false;
            AssemblyResolvePolicy = null;
        }

        internal static void ResetForDomainReload()
        {
            ShrinkModCordisRuntime.ResetForDomainReload();
            DisposeLegacyEventBindings();
            LoadedMods.Clear();
            LoadSequence.Clear();
            RegistryManager.Clear();
            IsLoaded = false;
            OnModReady = null;
            OnAllModsReady = null;
            ShrinkModNetworkManager.ResetForDomainReload();
            ShrinkExternalModAssemblyLoader.ResetForDomainReload();
            ShrinkHarmonyPatchService.ResetForTesting();
            _authorizedDllPaths = Array.Empty<string>();
            _hostComposed = false;
            AssemblyResolvePolicy = null;
        }

        private static IReadOnlyDictionary<string, ShrinkModHandle> ApplyContextComposition(
            ShrinkModFrameworkSettings settings,
            IEnumerable<string> authorizedDllPaths)
        {
            var previousGenerations = Mods.ToDictionary(
                pair => pair.Key,
                pair => pair.Value.Generation,
                StringComparer.Ordinal);

            var result = ShrinkModCordisRuntime
                .ApplyDiscoveredAsync(settings, authorizedDllPaths)
                .GetAwaiter()
                .GetResult();
            PublishReady(result, previousGenerations);
            return result;
        }

        private static void PublishReady(IReadOnlyDictionary<string, ShrinkModHandle> result,
            IReadOnlyDictionary<string, long> previousGenerations)
        {
            IsLoaded = true;
            foreach (var pair in result.OrderBy(pair => pair.Key, StringComparer.Ordinal))
            {
                if (!previousGenerations.TryGetValue(pair.Key, out var generation) ||
                    generation != pair.Value.Generation)
                {
                    OnModReady?.Invoke(pair.Value);
                }
            }
            OnAllModsReady?.Invoke(result);
        }

        private static ShrinkModHandle InstantiateAndTrack(DiscoveredMod discoveredMod,
            ShrinkModFrameworkSettings settings, bool verboseLogging)
        {
            var instance = CreateInstance(discoveredMod);
            var handle = new ShrinkModHandle(discoveredMod.Info, instance);
            LoadedMods.Add(handle.Info.ModId, handle);

            ShrinkHarmonyPatchService.ApplyPatchesIfNeeded(
                handle.Info,
                settings == null || settings.enableHarmonyPatching,
                verboseLogging);

            return handle;
        }

        private static List<DiscoveredMod> DiscoverMods(ShrinkModFrameworkSettings settings)
        {
            var results = new List<DiscoveredMod>();
            var prefixes = settings?.assemblyNamePrefixes;

            foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies())
            {
                if (ShrinkExternalModAssemblyLoader.IsExternalAssembly(assembly) &&
                    !ShrinkExternalModAssemblyLoader.TryGetCurrentRevision(assembly, out _))
                {
                    continue;
                }
                if (!ShouldScanAssembly(assembly, prefixes))
                    continue;

                foreach (var type in GetTypesSafely(assembly))
                {
                    if (type == null || type.IsAbstract || type.IsInterface)
                        continue;

                    var modAttribute = type.GetCustomAttribute<ShrinkModAttribute>(false);
                    if (modAttribute == null)
                        continue;

                    if (!typeof(IShrinkMod).IsAssignableFrom(type))
                    {
                        Debug.LogWarning($"[ShrinkModLoader] 类型 {type.FullName} 标记了 ShrinkMod，但没有实现 IShrinkMod，已跳过。");
                        continue;
                    }

                    if (type.GetConstructor(Type.EmptyTypes) == null)
                        throw new InvalidOperationException($"模组 {type.FullName} 缺少无参构造函数，无法实例化。");

                    var dependencies = type
                        .GetCustomAttributes<ShrinkModDependencyAttribute>(false)
                        .Select(attr => new ShrinkModDependency(attr.ModId, attr.MinimumVersion, attr.Optional, attr.Ordering, attr.VersionRange))
                        .ToArray();

                    results.Add(new DiscoveredMod
                    {
                        EntryType = type,
                        Info = new ShrinkModInfo(
                            modAttribute.ModId,
                            modAttribute.DisplayName,
                            modAttribute.Version,
                            modAttribute.LoadOrder,
                            type,
                            modAttribute.AutoApplyHarmonyPatches,
                            dependencies)
                    });
                }
            }

            ValidateDiscoveredMods(results);
            return results;
        }

        private static void ValidateDiscoveredMods(List<DiscoveredMod> discovered)
        {
            var seen = new HashSet<string>(StringComparer.Ordinal);
            foreach (var mod in discovered)
            {
                if (!seen.Add(mod.Info.ModId))
                    throw new InvalidOperationException($"发现重复模组 ID：{mod.Info.ModId}");
            }

            var map = discovered.ToDictionary(x => x.Info.ModId, x => x, StringComparer.Ordinal);
            foreach (var mod in discovered)
            {
                foreach (var dependency in mod.Info.Dependencies)
                {
                    if (!map.TryGetValue(dependency.ModId, out var target))
                    {
                        if (!dependency.Optional && !LoadedMods.ContainsKey(dependency.ModId))
                            throw new InvalidOperationException($"模组 {mod.Info.ModId} 缺少必选依赖 {dependency.ModId}");
                        continue;
                    }

                    ValidateDependencyVersion(mod.Info, dependency, target.Info.Version);
                }
            }
        }

        private static void ValidateNewMods(List<DiscoveredMod> newMods)
        {
            var seen = new HashSet<string>(LoadedMods.Keys, StringComparer.Ordinal);
            foreach (var mod in newMods)
            {
                if (!seen.Add(mod.Info.ModId))
                    throw new InvalidOperationException($"发现重复模组 ID：{mod.Info.ModId}");
            }

            var map = newMods.ToDictionary(x => x.Info.ModId, x => x, StringComparer.Ordinal);
            foreach (var mod in newMods)
            {
                foreach (var dependency in mod.Info.Dependencies)
                {
                    if (map.TryGetValue(dependency.ModId, out var newTarget))
                    {
                        ValidateDependencyVersion(mod.Info, dependency, newTarget.Info.Version);
                        continue;
                    }

                    if (LoadedMods.TryGetValue(dependency.ModId, out var loadedTarget))
                    {
                        ValidateDependencyVersion(mod.Info, dependency, loadedTarget.Info.Version);
                        continue;
                    }

                    if (!dependency.Optional)
                        throw new InvalidOperationException($"模组 {mod.Info.ModId} 缺少必选依赖 {dependency.ModId}");
                }
            }
        }

        private static void ValidateDependencyVersion(ShrinkModInfo owner, ShrinkModDependency dependency, string actualVersion)
        {
            ShrinkModDependencyGraph.ValidateVersion(owner, dependency, actualVersion);
        }

        private static List<DiscoveredMod> ResolveLoadOrder(List<DiscoveredMod> discovered, bool allowExistingLoadedDependencies)
        {
            var map = discovered.ToDictionary(mod => mod.Info.ModId, StringComparer.Ordinal);
            var infos = discovered.Select(mod => mod.Info).ToList();
            if (allowExistingLoadedDependencies)
                infos.AddRange(LoadedMods.Values.Where(handle => !map.ContainsKey(handle.Info.ModId)).Select(handle => handle.Info));
            var ordered = ShrinkModDependencyGraph.Order(infos);
            foreach (var loaded in LoadedMods.Keys)
                if (ShrinkModDependencyGraph.Predecessors(loaded, infos).Any(map.ContainsKey))
                    throw new InvalidOperationException($"Cannot load a new mod before already loaded mod {loaded}; use ContextHost.");
            return ordered.Where(info => map.ContainsKey(info.ModId)).Select(info => map[info.ModId]).ToList();
        }

        private static IShrinkMod CreateInstance(DiscoveredMod discovered)
        {
            try
            {
                return (IShrinkMod)Activator.CreateInstance(discovered.EntryType);
            }
            catch (Exception ex)
            {
                throw new InvalidOperationException($"实例化模组 {discovered.Info.ModId} 失败：{ex.Message}", ex);
            }
        }

        private static void RunLifecycle(IEnumerable<ShrinkModHandle> handles, bool verboseLogging)
        {
            var stagedHandles = handles.ToList();
            RunPhase(stagedHandles, ShrinkModState.Constructed, verboseLogging, (mod, context) => mod.OnConstruct(context), "构造");
            RunPhase(stagedHandles, ShrinkModState.ContentRegistered, verboseLogging,
                (mod, context) => mod.OnRegisterContent(context), "注册内容");
            RegisterOptionalRuntimeIntegrations(stagedHandles, verboseLogging);
            RunPhase(stagedHandles, ShrinkModState.Initialized, verboseLogging,
                (mod, context) => mod.OnInitialize(context), "初始化");
            RunPhase(stagedHandles, ShrinkModState.Ready, verboseLogging, (mod, context) => mod.OnReady(context), "完成就绪");
        }

        private static void RegisterOptionalRuntimeIntegrations(IEnumerable<ShrinkModHandle> handles, bool verboseLogging)
        {
            foreach (var handle in handles)
                handle.EventBinding = ShrinkModOptionalRuntimeIntegration.RegisterModInstance(
                    handle.Instance, handle.Info, verboseLogging);

            ShrinkModOptionalRuntimeIntegration.RefreshGlobalBindings(verboseLogging);
        }

        private static void DisposeLegacyEventBindings()
        {
            foreach (var handle in LoadedMods.Values)
            {
                handle.EventBinding?.Dispose();
                handle.EventBinding = null;
                EventBus.RemoveBus(ShrinkBusKey.Mod(handle.Info.ModId));
            }
        }

        private static void RunPhase(IEnumerable<ShrinkModHandle> handles, ShrinkModState targetState, bool verboseLogging,
            Action<IShrinkMod, ShrinkModContext> callback, string phaseName)
        {
            foreach (var handle in handles)
            {
                var context = new ShrinkModContext(handle.Info, RegistryManager, LoadedMods, verboseLogging);

                try
                {
                    callback(handle.Instance, context);
                    handle.State = targetState;

                    if (verboseLogging)
                        Debug.Log($"[ShrinkModLoader] {handle.Info.ModId} 已完成阶段：{phaseName}");

                    if (targetState == ShrinkModState.Ready)
                        OnModReady?.Invoke(handle);
                }
                catch (Exception ex)
                {
                    throw new InvalidOperationException(
                        $"模组 {handle.Info.ModId} 在阶段 '{phaseName}' 执行失败：{ex.Message}", ex);
                }
            }
        }

        private static bool ShouldScanAssembly(Assembly assembly, string[] prefixes)
        {
            var name = assembly.GetName().Name;
            if (string.IsNullOrEmpty(name))
                return false;

            if (name.StartsWith("Unity", StringComparison.Ordinal) ||
                name.StartsWith("System", StringComparison.Ordinal) ||
                name.StartsWith("mscorlib", StringComparison.Ordinal) ||
                name.StartsWith("netstandard", StringComparison.Ordinal))
                return false;

            if (prefixes == null || prefixes.Length == 0)
                return true;

            for (var i = 0; i < prefixes.Length; i++)
            {
                var prefix = prefixes[i];
                if (!string.IsNullOrWhiteSpace(prefix) &&
                    name.StartsWith(prefix, StringComparison.OrdinalIgnoreCase))
                    return true;
            }

            return false;
        }

        private static IEnumerable<Type> GetTypesSafely(Assembly assembly)
        {
            try
            {
                return assembly.GetTypes();
            }
            catch (ReflectionTypeLoadException ex)
            {
                return ex.Types.Where(t => t != null);
            }
            catch
            {
                return Array.Empty<Type>();
            }
        }
    }
}