Workspace/Assets/Modules/ShrinkModFramework/Runtime/Loading/ShrinkExternalModAssemblyLoader.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Security.Cryptography;
using UnityEngine;

namespace ShrinkModFramework
{
    internal static class ShrinkExternalModAssemblyLoader
    {
        internal sealed class ExternalAssemblyRevision
        {
            public string Path;
            public string Revision;
            public Assembly Assembly;
            public long LoadedBytes;
        }

        internal sealed class RevisionState
        {
            public Dictionary<string, ExternalAssemblyRevision> Current =
                new(StringComparer.OrdinalIgnoreCase);
            public Dictionary<string, string> Known =
                new(StringComparer.OrdinalIgnoreCase);
        }

        private static readonly Dictionary<string, ExternalAssemblyRevision> CurrentAssemblyRevisions =
            new(StringComparer.OrdinalIgnoreCase);
        private static readonly Dictionary<Assembly, ExternalAssemblyRevision> ExternalAssemblyHistory = new();
        private static readonly Dictionary<string, string> KnownAssemblyFiles = new(StringComparer.OrdinalIgnoreCase);
        private static readonly object ResolveLock = new();
        private static bool _resolveRegistered;
        private static int _lastWarnedResidentCount;

        public static IReadOnlyList<Assembly> LoadExternalAssemblies(ShrinkModFrameworkSettings settings, bool verboseLogging)
        {
            return ScanExternalAssemblyRevisions(settings, verboseLogging)
                .Select(revision => revision.Assembly)
                .ToArray();
        }

        internal static IReadOnlyList<ExternalAssemblyRevision> ScanExternalAssemblyRevisions(
            ShrinkModFrameworkSettings settings, bool verboseLogging)
        {
            if (settings != null && !settings.enableExternalDllMods)
            {
                CurrentAssemblyRevisions.Clear();
                KnownAssemblyFiles.Clear();
                return Array.Empty<ExternalAssemblyRevision>();
            }

#if ENABLE_IL2CPP && !UNITY_EDITOR
            Debug.LogWarning("[ShrinkModFramework] IL2CPP 运行时不支持外部 DLL 热加载，已跳过外部模组扫描。");
            return Array.Empty<ExternalAssemblyRevision>();
#else
            var modsDirectory = GetExternalModsDirectory(settings);
            if (settings == null || settings.autoCreateExternalModsDirectory)
                Directory.CreateDirectory(modsDirectory);

            RegisterAssemblyResolve();

            var dllPaths = Directory.GetFiles(modsDirectory, "*.dll", SearchOption.AllDirectories)
                .Select(Path.GetFullPath)
                .OrderBy(path => path, StringComparer.OrdinalIgnoreCase)
                .ToArray();
            var presentPaths = new HashSet<string>(dllPaths, StringComparer.OrdinalIgnoreCase);

            foreach (var dllPath in dllPaths
                         .OrderBy(path => path, StringComparer.OrdinalIgnoreCase))
            {
                var assemblyName = Path.GetFileNameWithoutExtension(dllPath);
                KnownAssemblyFiles[assemblyName] = dllPath;

                try
                {
                    var bytes = File.ReadAllBytes(dllPath);
                    var revision = ComputeSha256(bytes);
                    if (CurrentAssemblyRevisions.TryGetValue(dllPath, out var current) &&
                        string.Equals(current.Revision, revision, StringComparison.Ordinal))
                    {
                        continue;
                    }

                    var pdbPath = Path.ChangeExtension(dllPath, ".pdb");
                    byte[] pdbBytes = null;
                    if (File.Exists(pdbPath))
                        pdbBytes = File.ReadAllBytes(pdbPath);
                    var assembly = pdbBytes != null
                        ? Assembly.Load(bytes, pdbBytes)
                        : Assembly.Load(bytes);

                    var loadedRevision = new ExternalAssemblyRevision
                    {
                        Path = dllPath,
                        Revision = revision,
                        Assembly = assembly,
                        LoadedBytes = bytes.LongLength + (pdbBytes?.LongLength ?? 0L)
                    };
                    CurrentAssemblyRevisions[dllPath] = loadedRevision;
                    ExternalAssemblyHistory[assembly] = loadedRevision;

                    if (verboseLogging)
                        Debug.Log($"[ShrinkModFramework] 已加载外部模组程序集：{assembly.GetName().Name} ({revision[..12]})");
                }
                catch (Exception ex)
                {
                    Debug.LogError($"[ShrinkModFramework] 加载外部 DLL 失败：{dllPath}\n{ex.Message}");
                }
            }

            WarnIfResidentLimitExceeded(settings);

            foreach (var missingPath in CurrentAssemblyRevisions.Keys
                         .Where(path => !presentPaths.Contains(path))
                         .ToArray())
            {
                CurrentAssemblyRevisions.Remove(missingPath);
                var assemblyName = Path.GetFileNameWithoutExtension(missingPath);
                if (KnownAssemblyFiles.TryGetValue(assemblyName, out var knownPath) &&
                    string.Equals(knownPath, missingPath, StringComparison.OrdinalIgnoreCase))
                {
                    KnownAssemblyFiles.Remove(assemblyName);
                }
            }

            return CurrentAssemblyRevisions.Values
                .OrderBy(revision => revision.Path, StringComparer.OrdinalIgnoreCase)
                .ToArray();
#endif
        }

        internal static bool IsExternalAssembly(Assembly assembly) =>
            assembly != null && ExternalAssemblyHistory.ContainsKey(assembly);

        internal static ShrinkExternalAssemblyDiagnostic CaptureDiagnostic(int softLimit)
        {
            var normalizedLimit = Math.Max(1, softLimit);
            var currentAssemblies = new HashSet<Assembly>(
                CurrentAssemblyRevisions.Values.Select(item => item.Assembly));
            var revisions = ExternalAssemblyHistory.Values
                .OrderBy(item => item.Path, StringComparer.OrdinalIgnoreCase)
                .ThenBy(item => item.Revision, StringComparer.Ordinal)
                .Select(item => new ShrinkExternalAssemblyRevisionDiagnostic(
                    item.Path,
                    item.Revision,
                    item.Assembly.GetName().Name ?? string.Empty,
                    item.LoadedBytes,
                    currentAssemblies.Contains(item.Assembly)))
                .ToArray();

            return new ShrinkExternalAssemblyDiagnostic(
                CurrentAssemblyRevisions.Count,
                revisions,
                revisions.Sum(item => item.LoadedBytes),
                normalizedLimit);
        }

        internal static RevisionState CaptureState()
        {
            return new RevisionState
            {
                Current = CurrentAssemblyRevisions.ToDictionary(
                    pair => pair.Key,
                    pair => pair.Value,
                    StringComparer.OrdinalIgnoreCase),
                Known = KnownAssemblyFiles.ToDictionary(
                    pair => pair.Key,
                    pair => pair.Value,
                    StringComparer.OrdinalIgnoreCase)
            };
        }

        internal static void RestoreState(RevisionState state)
        {
            if (state == null)
                throw new ArgumentNullException(nameof(state));

            CurrentAssemblyRevisions.Clear();
            foreach (var pair in state.Current)
                CurrentAssemblyRevisions[pair.Key] = pair.Value;

            KnownAssemblyFiles.Clear();
            foreach (var pair in state.Known)
                KnownAssemblyFiles[pair.Key] = pair.Value;
        }

        internal static bool TryGetCurrentRevision(Assembly assembly, out string revision)
        {
            foreach (var current in CurrentAssemblyRevisions.Values)
            {
                if (ReferenceEquals(current.Assembly, assembly))
                {
                    revision = current.Revision;
                    return true;
                }
            }

            revision = null;
            return false;
        }

        public static string GetExternalModsDirectory(ShrinkModFrameworkSettings settings)
        {
            var folderName = settings != null && !string.IsNullOrWhiteSpace(settings.externalModsFolderName)
                ? settings.externalModsFolderName.Trim()
                : "Mods";
            return Path.Combine(Application.persistentDataPath, folderName);
        }

        internal static void ResetForTesting()
        {
            CurrentAssemblyRevisions.Clear();
            KnownAssemblyFiles.Clear();
            if (_resolveRegistered)
            {
                AppDomain.CurrentDomain.AssemblyResolve -= OnAssemblyResolve;
                _resolveRegistered = false;
            }
        }

        internal static void ResetForDomainReload()
        {
            ResetForTesting();
            ExternalAssemblyHistory.Clear();
            _lastWarnedResidentCount = 0;
        }

        private static void WarnIfResidentLimitExceeded(ShrinkModFrameworkSettings settings)
        {
            var limit = Math.Max(1, settings != null ? settings.externalAssemblyRevisionSoftLimit : 16);
            var residentCount = ExternalAssemblyHistory.Count;
            if (residentCount < limit || residentCount == _lastWarnedResidentCount)
                return;

            _lastWarnedResidentCount = residentCount;
            var bytes = ExternalAssemblyHistory.Values.Sum(item => item.LoadedBytes);
            Debug.LogWarning(
                $"[ShrinkModFramework] 外部 DLL 常驻 revision 已达到 {residentCount} 个（载入文件约 {bytes} bytes，" +
                $"软阈值 {limit}）。Mono 无法卸载这些 Assembly；建议在维护窗口执行 Domain Reload 或重启进程。");
        }

        private static string ComputeSha256(byte[] bytes)
        {
            using var sha256 = SHA256.Create();
            return BitConverter.ToString(sha256.ComputeHash(bytes)).Replace("-", string.Empty);
        }

        private static void RegisterAssemblyResolve()
        {
            lock (ResolveLock)
            {
                if (_resolveRegistered)
                    return;

                AppDomain.CurrentDomain.AssemblyResolve += OnAssemblyResolve;
                _resolveRegistered = true;
            }
        }

        private static Assembly OnAssemblyResolve(object sender, ResolveEventArgs args)
        {
            var requestedName = new AssemblyName(args.Name).Name;
            if (string.IsNullOrWhiteSpace(requestedName))
                return null;

            if (!KnownAssemblyFiles.TryGetValue(requestedName, out var path) || !File.Exists(path))
                return null;

            try
            {
                return Assembly.Load(File.ReadAllBytes(path));
            }
            catch
            {
                return null;
            }
        }
    }
}