Fishing2Server/Fantasy/Fantasy.Net/Fantasy.SourceGenerator/Generators/AssemblyInitializerGenerator.cs

using Fantasy.SourceGenerator.Common;
using Microsoft.CodeAnalysis;

namespace Fantasy.SourceGenerator.Generators
{
    /// <summary>
    /// 程序集初始化器生成器
    /// 为每个 Fantasy 项目生成 ModuleInitializer，在程序集加载时自动注册到框架
    /// 支持 Native AOT，无需运行时反射
    /// </summary>
    [Generator]
    public class AssemblyInitializerGenerator : IIncrementalGenerator
    {
        public void Initialize(IncrementalGeneratorInitializationContext context)
        {
            // 获取编译信息
            var compilationProvider = context.CompilationProvider;
            // 注册源代码输出
            context.RegisterSourceOutput(compilationProvider, (spc, compilation) =>
            {
                GenerateModuleInitializer(spc, compilation);
            });
        }

        private static void GenerateModuleInitializer(
            SourceProductionContext context,
            Compilation compilation)
        {
            // 检查是否定义了 FANTASY_NET 或 FANTASY_UNITY 预编译符号
            if (!CompilationHelper.HasFantasyDefine(compilation))
            {
                // 不是 Fantasy 框架项目，跳过代码生成
                return;
            }

            var assemblyName = compilation.AssemblyName ?? "Unknown";
            // 检测是否是 Unity 环境
            var isUnity = CompilationHelper.IsUnityCompilation(compilation);
            var builder = new SourceCodeBuilder();
            // 添加文件头
            builder.AppendLine(GeneratorConstants.AutoGeneratedHeader);
            // 添加 using
            if (isUnity)
            {
                builder.AddUsings(
                    "System",
                    "UnityEngine"
                );
            }
            else
            {
                builder.AddUsings(
                    "System",
                    "System.Runtime.CompilerServices"
                );
            }
            builder.AppendLine();
            // 开始命名空间
            builder.BeginNamespace(GeneratorConstants.GeneratedNamespace);
            // 添加类注释
            builder.AddXmlComment($"Auto-generated assembly initializer for {assemblyName}");
            builder.AddXmlComment("This class is automatically invoked when the assembly is loaded via ModuleInitializer");
            // 开始类定义
            builder.BeginClass("AssemblyInitializer", "internal static");
            // 添加字段
            builder.AppendLine("private static bool _initialized;");
            builder.AppendLine("private static long _assemblyManifestId;");
            builder.AppendLine();

            // 生成 ModuleInitializer 方法
            GenerateInitializeMethod(builder, assemblyName, isUnity);

            // 生成卸载方法
            GenerateUnloadMethod(builder, isUnity);

            // 结束 AssemblyInitializer 类
            builder.EndClass();

            builder.AppendLine();

            // 生成 AssemblyMarker 类（用于强制加载程序集）
            // 类名包含程序集名称以避免冲突，将特殊字符替换为下划线
            var markerClassName = assemblyName.Replace("-", "_").Replace(".", "_") + "_AssemblyMarker";
            builder.AddXmlComment($"Public marker class for forcing {assemblyName} assembly load");
            builder.AddXmlComment("Access this type to ensure the assembly is loaded and ModuleInitializer executes");
            builder.BeginClass(markerClassName, "public static");
            builder.AppendLine("/// <summary>");
            builder.AppendLine("/// Call this method to ensure the assembly is loaded");
            builder.AppendLine("/// This is useful when loading assemblies dynamically via reflection");
            builder.AppendLine("/// </summary>");
            builder.BeginMethod("public static void EnsureLoaded()");
            builder.AppendLine("// Accessing this method ensures the assembly is loaded");
            builder.AppendLine("// ModuleInitializer will execute automatically when assembly loads");
            builder.EndMethod();
            builder.EndClass();

            // 结束命名空间
            builder.EndNamespace();
            // 输出源代码
            context.AddSource("AssemblyInitializer.g.cs", builder.ToString());
        }

        private static void GenerateInitializeMethod(SourceCodeBuilder builder, string assemblyName, bool isUnity)
        {
            if (isUnity)
            {
                builder.AddXmlComment(
                    "Unity runtime initializer - automatically called when entering play mode or on app start");
                builder.AppendLine("[RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.AfterAssembliesLoaded)]");
                builder.BeginMethod("internal static void Initialize()");
            }
            else
            {
                builder.AddXmlComment("Module initializer - automatically called when assembly is loaded");
                builder.AppendLine("[ModuleInitializer]");
                builder.BeginMethod("internal static void Initialize()");
            }

            // 防止重复初始化
            builder.AppendLine("if (_initialized)");
            builder.OpenBrace();
            builder.AppendLine("return;");
            builder.CloseBrace();
            builder.AppendLine();
            builder.AppendLine("_initialized = true;");
            builder.AppendLine();

            // 获取程序集并计算唯一标识
            builder.AddComment("Get assembly and calculate manifest ID");
            builder.AppendLine("var assembly = typeof(AssemblyInitializer).Assembly;");
            builder.AppendLine(
                $"_assemblyManifestId = Fantasy.Helper.HashCodeHelper.ComputeHash64(assembly.GetName().Name ?? \"{assemblyName}\");");
            builder.AppendLine();

            // 注册卸载事件（用于热更新支持）
            builder.AddComment("Register auto-unload for collectible AssemblyLoadContext (hot-reload support)");
            if (isUnity)
            {
                builder.AppendLine(
                    "#if !UNITY_EDITOR && !UNITY_STANDALONE && !UNITY_ANDROID && !UNITY_IOS && !UNITY_WEBGL");
            }

            builder.AppendLine("var loadContext = System.Runtime.Loader.AssemblyLoadContext.GetLoadContext(assembly);");
            builder.AppendLine(
                "if (loadContext != null && loadContext != System.Runtime.Loader.AssemblyLoadContext.Default)");
            builder.OpenBrace();
            builder.AppendLine("loadContext.Unloading += OnAssemblyUnloading;");
            builder.CloseBrace();
            if (isUnity)
            {
                builder.AppendLine("#endif");
            }

            builder.AppendLine();

            // 声明所有可能的注册器变量
            builder.AddComment("Declare registrar variables");
            builder.AppendLine("Fantasy.Assembly.INetworkProtocolRegistrar? networkProtocolRegistrar = null;");
            builder.AppendLine("Fantasy.Assembly.IEventSystemRegistrar? eventSystemRegistrar = null;");
            builder.AppendLine("Fantasy.Assembly.IEntitySystemRegistrar? entitySystemRegistrar = null;");
            builder.AppendLine("Fantasy.Assembly.IMessageHandlerResolver? messageHandlerResolverRegistrar = null;");
            builder.AppendLine(
                "Fantasy.Assembly.IEntityTypeCollectionRegistrar? entityTypeCollectionRegistrar = null;");
            builder.AppendLine(
                "Fantasy.Assembly.INetworkProtocolOpCodeResolver? networkProtocolOpCodeResolverRegistrar = null;");
            builder.AppendLine(
                "Fantasy.Assembly.INetworkProtocolResponseTypeResolver? networkProtocolResponseTypeResolverRegistrar = null;");
            builder.AppendLine("#if FANTASY_NET", false);
            builder.AppendLine(
                "Fantasy.Assembly.ISeparateTableRegistrar? separateTableRegistrar = null;");
            builder.AppendLine("#endif", false);
            builder.AppendLine();

            // 尝试创建各个注册器（如果存在）
            builder.AddComment("Try to create registrars if they were generated in this assembly");
            GenerateTryCreateRegistrar(builder, "NetworkProtocol", "networkProtocolRegistrar");
            GenerateTryCreateRegistrar(builder, "EventSystem", "eventSystemRegistrar");
            GenerateTryCreateRegistrar(builder, "EntitySystem", "entitySystemRegistrar");
            GenerateTryCreateRegistrar(builder, "MessageHandlerResolver", "messageHandlerResolverRegistrar");
            GenerateTryCreateRegistrar(builder, "EntityTypeCollection", "entityTypeCollectionRegistrar");
            GenerateTryCreateRegistrar(builder, "NetworkProtocolOpCodeResolver", "networkProtocolOpCodeResolverRegistrar");
            GenerateTryCreateRegistrar(builder, "NetworkProtocolResponseTypeResolver", "networkProtocolResponseTypeResolverRegistrar");
            builder.AppendLine("#if FANTASY_NET", false);
            GenerateTryCreateRegistrar(builder, "SeparateTable", "separateTableRegistrar");
            builder.AppendLine("#endif", false);
            
            builder.AppendLine();

            // 注册到框架
            builder.AddComment("Register complete AssemblyManifest to the framework");
            builder.AppendLine("#if FANTASY_NET", false);
            builder.AppendLine("Fantasy.Assembly.AssemblyManifest.Register(");
            builder.Indent();
            builder.AppendLine("_assemblyManifestId,");
            builder.AppendLine("assembly,");
            builder.AppendLine("networkProtocolRegistrar,");
            builder.AppendLine("eventSystemRegistrar,");
            builder.AppendLine("entitySystemRegistrar,");
            builder.AppendLine("messageHandlerResolverRegistrar,");
            builder.AppendLine("entityTypeCollectionRegistrar,");
            builder.AppendLine("separateTableRegistrar,");
            builder.AppendLine("networkProtocolOpCodeResolverRegistrar,");
            builder.AppendLine("networkProtocolResponseTypeResolverRegistrar);");
            builder.Unindent();
            builder.AppendLine("#endif", false);
            builder.AppendLine("#if FANTASY_UNITY", false);
            builder.AppendLine("Fantasy.Assembly.AssemblyManifest.Register(");
            builder.Indent();
            builder.AppendLine("_assemblyManifestId,");
            builder.AppendLine("assembly,");
            builder.AppendLine("networkProtocolRegistrar,");
            builder.AppendLine("eventSystemRegistrar,");
            builder.AppendLine("entitySystemRegistrar,");
            builder.AppendLine("messageHandlerResolverRegistrar,");
            builder.AppendLine("entityTypeCollectionRegistrar,");
            builder.AppendLine("networkProtocolOpCodeResolverRegistrar,");
            builder.AppendLine("networkProtocolResponseTypeResolverRegistrar);");
            builder.Unindent();
            builder.AppendLine("#endif", false);
            builder.EndMethod();
        }

        private static void GenerateUnloadMethod(SourceCodeBuilder builder, bool isUnity)
        {
            builder.AppendLine();
            builder.AddXmlComment("Called when AssemblyLoadContext is unloading (for hot-reload support)");

            // Unity 环境下，AssemblyLoadContext 仅在非编辑器/非独立平台可用
            if (isUnity)
            {
                builder.AppendLine("#if !UNITY_EDITOR && !UNITY_STANDALONE && !UNITY_ANDROID && !UNITY_IOS && !UNITY_WEBGL");
            }

            builder.BeginMethod("private static void OnAssemblyUnloading(System.Runtime.Loader.AssemblyLoadContext context)");

            builder.AddComment("Unregister from framework");
            builder.AppendLine("if (_assemblyManifestId != 0)");
            builder.OpenBrace();
            builder.AppendLine("Fantasy.Assembly.AssemblyManifest.Unregister(_assemblyManifestId);");
            builder.CloseBrace();

            builder.EndMethod();

            if (isUnity)
            {
                builder.AppendLine("#endif");
            }
        }

        private static void GenerateTryCreateRegistrar(
            SourceCodeBuilder builder,
            string registrarName,
            string variableName)
        {
            var typeName = $"Fantasy.Generated.{registrarName}Registrar";
            builder.AppendLine($"{variableName} = new {typeName}();");
        }
    }
}