Workspace/Assets/Modules/ShrinkModFramework/Runtime/Context/ShrinkModContextHost.cs

using System;
using System.Collections.Generic;
using System.Linq;
using Cysharp.Threading.Tasks;
using ShrinkContext;

namespace ShrinkModFramework
{
    public sealed class ShrinkModTransactionException : Exception
    {
        public ShrinkModTransactionException(string message, Exception applyError, Exception restoreError = null)
            : base(message, restoreError == null ? applyError : new AggregateException(applyError, restoreError))
        {
            ApplyError = applyError;
            RestoreError = restoreError;
        }

        public Exception ApplyError { get; }
        public Exception RestoreError { get; }
        public bool PreviousCompositionRestored => RestoreError == null;
    }

    /// <summary>
    /// 模组的 Cordis 组合宿主。配置应用是事务：新组合任一 fiber 失败时，重新协调到旧组件源；
    /// Mono/IL2CPP 下旧程序集仍驻留，但旧/新模组实例及其已追踪效应会被正确卸载或恢复。
    /// </summary>
    public sealed class ShrinkModContextHost
    {
        private readonly Dictionary<string, ShrinkModHandle> _mods = new(StringComparer.Ordinal);
        private readonly Dictionary<ShrinkModComponentSource, string> _catalogNames = new();
        private readonly List<ShrinkModComponentSource> _currentSources = new();
        private readonly ShrinkComponentCatalog _catalog = new();
        private long _generation;
        private int _catalogGeneration;
        private bool _applying;

        public ShrinkModContextHost(bool enableHarmonyPatching = true, bool verboseLogging = false)
        {
            EnableHarmonyPatching = enableHarmonyPatching;
            VerboseLogging = verboseLogging;
            Runtime = new ShrinkContextRuntime();
            Loader = new ShrinkContextLoader(Runtime, _catalog);
            RegistryManager = new ShrinkModRegistryManager();
        }

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

        public bool EnableHarmonyPatching { get; }
        public bool VerboseLogging { get; }
        public ShrinkContextRuntime Runtime { get; }
        public ShrinkContextLoader Loader { get; }
        public IReadOnlyDictionary<string, ShrinkModHandle> Mods => _mods;
        public IReadOnlyList<ShrinkModComponentSource> CurrentSources => _currentSources.ToArray();
        internal ShrinkModRegistryManager RegistryManager { get; }

        public ShrinkModRegistry<T> GetOrCreateRegistry<T>(string name) =>
            RegistryManager.GetOrCreateRegistry<T>(name);

        public static string GetModKey(string modId)
        {
            if (string.IsNullOrWhiteSpace(modId))
                throw new ArgumentException("ModId must not be empty.", nameof(modId));
            return "shrink.mod/" + modId.Trim();
        }

        public async UniTask ApplyAsync(IReadOnlyList<ShrinkModComponentSource> desiredSources)
        {
            if (desiredSources == null)
                throw new ArgumentNullException(nameof(desiredSources));
            if (_applying)
                throw new InvalidOperationException("A mod composition transaction is already running.");

            var desired = ReuseUnchangedSources(ValidateAndOrder(desiredSources));
            var previous = _currentSources.ToArray();
            var previousById = previous.ToDictionary(source => source.Info.ModId, StringComparer.Ordinal);
            _applying = true;
            try
            {
                await Loader.ApplyAsync(BuildEntries(desired));
                EnsureActive(desired);

                _currentSources.Clear();
                _currentSources.AddRange(desired);
                foreach (var source in desired)
                {
                    if (!previousById.TryGetValue(source.Info.ModId, out var oldSource) ||
                        !ReferenceEquals(source, oldSource))
                    {
                        OnModReady?.Invoke(_mods[source.Info.ModId]);
                    }
                }
                OnAllModsReady?.Invoke(Mods);
            }
            catch (Exception applyError)
            {
                Exception restoreError = null;
                try
                {
                    await Loader.ApplyAsync(BuildEntries(previous));
                    EnsureActive(previous);
                }
                catch (Exception ex)
                {
                    restoreError = ex;
                }

                throw new ShrinkModTransactionException(
                    restoreError == null
                        ? "Mod composition failed; the previous composition was restored."
                        : "Mod composition failed and restoring the previous composition also failed.",
                    applyError,
                    restoreError);
            }
            finally
            {
                _applying = false;
            }
        }

        public async UniTask ShutdownAsync()
        {
            if (_applying)
                throw new InvalidOperationException("Cannot shut down while a mod transaction is running.");

            await Loader.ApplyAsync(Array.Empty<ShrinkLoaderEntry>());
            _currentSources.Clear();
        }

        internal long NextGeneration() => ++_generation;

        internal void RegisterHandle(ShrinkModHandle handle)
        {
            if (!_mods.TryAdd(handle.Info.ModId, handle))
                throw new InvalidOperationException($"Mod handle already active: {handle.Info.ModId}");
        }

        internal void UnregisterHandle(ShrinkModHandle handle)
        {
            if (_mods.TryGetValue(handle.Info.ModId, out var current) && ReferenceEquals(current, handle))
                _mods.Remove(handle.Info.ModId);
        }

        private IReadOnlyList<ShrinkLoaderEntry> BuildEntries(IEnumerable<ShrinkModComponentSource> sources)
        {
            var entries = new List<ShrinkLoaderEntry>();
            foreach (var source in sources)
            {
                if (!_catalogNames.TryGetValue(source, out var catalogName))
                {
                    catalogName = $"shrink.mod.source/{source.Info.ModId}/{++_catalogGeneration}";
                    _catalogNames.Add(source, catalogName);
                    _catalog.Register(catalogName, () => new ShrinkModComponent(this, source));
                }

                entries.Add(new ShrinkLoaderEntry(source.Info.ModId, catalogName));
            }
            return entries;
        }

        private void EnsureActive(IEnumerable<ShrinkModComponentSource> sources)
        {
            foreach (var source in sources)
            {
                if (!Loader.TryGetFiber(source.Info.ModId, out var fiber))
                    throw new InvalidOperationException($"Mod fiber was not created: {source.Info.ModId}");
                if (fiber.LastError != null)
                    throw new InvalidOperationException($"Mod {source.Info.ModId} failed during apply.", fiber.LastError);
                if (fiber.State != ShrinkFiberState.Active)
                    throw new InvalidOperationException(
                        $"Mod {source.Info.ModId} did not become active (state={fiber.State}).");
            }
        }

        private static List<ShrinkModComponentSource> ValidateAndOrder(
            IReadOnlyList<ShrinkModComponentSource> sources)
        {
            var map = new Dictionary<string, ShrinkModComponentSource>(StringComparer.Ordinal);
            foreach (var source in sources)
            {
                if (source == null)
                    throw new ArgumentException("Mod sources must not contain null.", nameof(sources));
                if (!map.TryAdd(source.Info.ModId, source))
                    throw new InvalidOperationException($"Duplicate mod source id: {source.Info.ModId}");
            }

            foreach (var source in sources)
            {
                foreach (var dependency in source.Info.Dependencies)
                {
                    if (!map.TryGetValue(dependency.ModId, out var target))
                    {
                        if (!dependency.Optional)
                            throw new InvalidOperationException(
                                $"Mod {source.Info.ModId} is missing required dependency {dependency.ModId}.");
                        continue;
                    }

                    if (!string.IsNullOrEmpty(dependency.MinimumVersion) &&
                        ShrinkVersionUtility.Compare(target.Info.Version, dependency.MinimumVersion) < 0)
                    {
                        throw new InvalidOperationException(
                            $"Mod {source.Info.ModId} requires {dependency.ModId} >= {dependency.MinimumVersion}, " +
                            $"but found {target.Info.Version}.");
                    }
                }
            }

            var result = new List<ShrinkModComponentSource>();
            var visiting = new HashSet<string>(StringComparer.Ordinal);
            var visited = new HashSet<string>(StringComparer.Ordinal);
            foreach (var source in sources.OrderBy(item => item.Info.LoadOrder)
                         .ThenBy(item => item.Info.ModId, StringComparer.Ordinal))
            {
                Visit(source, map, visiting, visited, result);
            }
            return result;
        }

        private List<ShrinkModComponentSource> ReuseUnchangedSources(
            IEnumerable<ShrinkModComponentSource> desired)
        {
            var currentById = _currentSources.ToDictionary(source => source.Info.ModId, StringComparer.Ordinal);
            var normalized = new List<ShrinkModComponentSource>();
            foreach (var source in desired)
            {
                if (currentById.TryGetValue(source.Info.ModId, out var current) &&
                    string.Equals(current.Revision, source.Revision, StringComparison.Ordinal))
                {
                    normalized.Add(current);
                }
                else
                {
                    normalized.Add(source);
                }
            }
            return normalized;
        }

        private static void Visit(ShrinkModComponentSource source,
            IReadOnlyDictionary<string, ShrinkModComponentSource> map,
            ISet<string> visiting,
            ISet<string> visited,
            ICollection<ShrinkModComponentSource> result)
        {
            if (visited.Contains(source.Info.ModId))
                return;
            if (!visiting.Add(source.Info.ModId))
                throw new InvalidOperationException($"Circular mod dependency includes {source.Info.ModId}.");

            foreach (var dependency in source.Info.Dependencies.OrderBy(item => item.ModId, StringComparer.Ordinal))
            {
                if (map.TryGetValue(dependency.ModId, out var target))
                    Visit(target, map, visiting, visited, result);
            }

            visiting.Remove(source.Info.ModId);
            visited.Add(source.Info.ModId);
            result.Add(source);
        }
    }
}