RevitGen/RevitGen.Generator/SourceGenerationHelper.cs

using Microsoft.CodeAnalysis;

using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;

namespace RevitGen.Generator
{
    /// <summary>
    /// 一个静态辅助类，包含了所有用于生成C#源代码的逻辑。
    /// </summary>
    internal static class SourceGenerationHelper
    {
        private const string RevitCommandAttributeFullName = "RevitGen.Attributes.RevitCommandAttribute";
        private const string CommandHandlerAttributeFullName = "RevitGen.Attributes.CommandHandlerAttribute";

        /// <summary>
        /// 为一个具体的命令生成其 partial class 的【另一半】。
        /// 这一半会实现 IExternalCommand 并提供所有上下文属性。
        /// </summary>
        public static string GenerateCommandPartialClass(INamedTypeSymbol classSymbol)
        {
            var ns = classSymbol.ContainingNamespace.ToDisplayString();
            var className = classSymbol.Name;

            var commandHandlerMethod = classSymbol.GetMembers()
                .OfType<IMethodSymbol>()
                .FirstOrDefault(m => m.GetAttributes().Any(a => a.AttributeClass?.ToDisplayString() == CommandHandlerAttributeFullName));

            if (commandHandlerMethod == null || commandHandlerMethod.Parameters.Any() || !commandHandlerMethod.ReturnsVoid)
            {
                return $"// 错误({className}): 必须有一个被[CommandHandler]标记的、无参数且返回void的方法。";
            }

            var attributeData = classSymbol.GetAttributes().First(ad => ad.AttributeClass?.ToDisplayString() == RevitCommandAttributeFullName);
            var transactionMode = GetAttributeProperty(attributeData, "TransactionMode", 1);

            var source = new StringBuilder();
            source.AppendLine("// <auto-generated/>");
            source.AppendLine($"namespace {ns}");
            source.AppendLine("{");
            source.AppendLine("    using Autodesk.Revit.DB;");
            source.AppendLine("    using Autodesk.Revit.UI;");
            source.AppendLine("    using Autodesk.Revit.Attributes;");
            source.AppendLine("    using System.ComponentModel;");
            source.AppendLine();

            source.AppendLine($"    [Transaction((TransactionMode){transactionMode})]");
            source.AppendLine($"    public partial class {className} : IExternalCommand");
            source.AppendLine("    {");
            source.AppendLine("        private ExternalCommandData _commandData;");
            source.AppendLine("        public UIApplication UIApplication => _commandData.Application;");
            source.AppendLine("        public UIDocument UIDocument => UIApplication.ActiveUIDocument;");
            source.AppendLine("        public Document Document => UIDocument.Document;");
            source.AppendLine("        public View ActiveView => UIDocument.ActiveView;");
            source.AppendLine("        public Result Result { get; set; } = Result.Succeeded;");
            source.AppendLine("        public string ErrorMessage { get; set; } = string.Empty;");
            source.AppendLine("        public ElementSet ElementSet { get; private set; }");
            source.AppendLine();

            source.AppendLine("        public Result Execute(ExternalCommandData commandData, ref string message, ElementSet elements)");
            source.AppendLine("        {");
            source.AppendLine("            this._commandData = commandData;");
            source.AppendLine("            this.ElementSet = elements;");
            source.AppendLine();
            source.AppendLine("            try");
            source.AppendLine("            {");
            if (transactionMode == 1)
            {
                source.AppendLine($"                using (var trans = new Transaction(this.Document, \"{className}\"))");
                source.AppendLine("                {");
                source.AppendLine("                    trans.Start();");
                source.AppendLine($"                    this.{commandHandlerMethod.Name}();");
                source.AppendLine("                    if (this.Result == Result.Succeeded) trans.Commit(); else trans.RollBack();");
                source.AppendLine("                }");
            }
            else
            {
                source.AppendLine($"                this.{commandHandlerMethod.Name}();");
            }
            source.AppendLine("            }");
            source.AppendLine("            catch (System.Exception ex)");
            source.AppendLine("            {");
            source.AppendLine("                this.ErrorMessage = ex.ToString();");
            source.AppendLine("                this.Result = Result.Failed;");
            source.AppendLine("            }");
            source.AppendLine();
            source.AppendLine("            message = this.ErrorMessage;");
            source.AppendLine("            return this.Result;");
            source.AppendLine("        }");
            source.AppendLine("    }");
            source.AppendLine("}");
            return source.ToString();
        }

        /// <summary>
        /// 生成唯一的 IExternalApplication 类，用于创建Revit界面 (选项卡, 面板, 按钮)。
        /// </summary>
        public static string GenerateApplicationClass(IEnumerable<INamedTypeSymbol> commandClasses)
        {
            var source = new StringBuilder();
            source.AppendLine("// <auto-generated/>");
            source.AppendLine("using Autodesk.Revit.UI;");
            source.AppendLine("using System;");
            source.AppendLine("using System.Reflection;");
            source.AppendLine("using System.Windows.Media;");
            source.AppendLine("using System.Windows.Media.Imaging;");
            source.AppendLine("using System.IO;");
            source.AppendLine("using System.Drawing;");
            source.AppendLine("using System.Resources;");

            source.AppendLine();
            source.AppendLine("namespace RevitGen.Runtime");
            source.AppendLine("{");
            source.AppendLine("    public class RevitGenApplication : IExternalApplication");
            source.AppendLine("    {");
            source.AppendLine("        private static ResourceManager _resourceManager;");

            source.AppendLine("        private BitmapSource BitmapToImageSource(Bitmap bitmap)");
            source.AppendLine("        {");
            source.AppendLine("            if (bitmap == null) return null;");
            source.AppendLine("            using (MemoryStream stream = new MemoryStream())");
            source.AppendLine("            {");
            source.AppendLine("                bitmap.Save(stream, System.Drawing.Imaging.ImageFormat.Png);");
            source.AppendLine("                stream.Position = 0;");
            source.AppendLine("                BitmapImage result = new BitmapImage();");
            source.AppendLine("                result.BeginInit();");
            source.AppendLine("                result.CacheOption = BitmapCacheOption.OnLoad;");
            source.AppendLine("                result.StreamSource = stream;");
            source.AppendLine("                result.EndInit();");
            source.AppendLine("                result.Freeze();");
            source.AppendLine("                return result;");
            source.AppendLine("            }");
            source.AppendLine("        }");
            source.AppendLine();

            source.AppendLine("        private BitmapSource LoadImageFromEmbeddedResource(Assembly assembly, string resourcePath)");
            source.AppendLine("        {");
            source.AppendLine("            try");
            source.AppendLine("            {");
            source.AppendLine("                string assemblyName = assembly.GetName().Name;");
            source.AppendLine("                string resourceName = $\"{assemblyName}.{resourcePath.Replace('/', '.').Replace('\\\\', '.')}\";");
            source.AppendLine("                using (Stream stream = assembly.GetManifestResourceStream(resourceName))");
            source.AppendLine("                {");
            source.AppendLine("                    if (stream == null) return null;");
            source.AppendLine("                    var image = new BitmapImage();");
            source.AppendLine("                    image.BeginInit();");
            source.AppendLine("                    image.StreamSource = stream;");
            source.AppendLine("                    image.CacheOption = BitmapCacheOption.OnLoad;");
            source.AppendLine("                    image.EndInit();");
            source.AppendLine("                    image.Freeze();");
            source.AppendLine("                    return image;");
            source.AppendLine("                }");
            source.AppendLine("            }");
            source.AppendLine("            catch { return null; }");
            source.AppendLine("        }");
            source.AppendLine();

            source.AppendLine("        public Result OnStartup(UIControlledApplication app)");
            source.AppendLine("        {");
            source.AppendLine("            var assembly = Assembly.GetExecutingAssembly();");
            source.AppendLine("            string assemblyPath = assembly.Location;");

            source.AppendLine("            if (_resourceManager == null)");
            source.AppendLine("            {");
            source.AppendLine("                string resourceName = assembly.GetName().Name + \".Properties.Resources\";");
            source.AppendLine("                try { _resourceManager = new ResourceManager(resourceName, assembly); } catch { /* ResX 不存在时忽略 */ }");
            source.AppendLine("            }");

            source.AppendLine();

            var commandsWithData = commandClasses.Select(c => new {
                Symbol = c,
                Attribute = c.GetAttributes().First(ad => ad.AttributeClass?.ToDisplayString() == RevitCommandAttributeFullName)
            }).ToList();

            var groupedByTab = commandsWithData.GroupBy(data => GetAttributeProperty<string>(data.Attribute, "TabName", null));

            foreach (var tabGroup in groupedByTab)
            {
                var tabName = tabGroup.Key;

                if (!string.IsNullOrEmpty(tabName))
                {
                    source.AppendLine($"            try {{ app.CreateRibbonTab(\"{tabName}\"); }} catch {{ /* 选项卡已存在 */ }}");
                }

                var groupedByPanel = tabGroup.GroupBy(data => GetAttributeProperty(data.Attribute, "PanelName", "Commands"));
                foreach (var panelGroup in groupedByPanel)
                {
                    var panelName = panelGroup.Key;
                    var panelVar = $"panel_{SanitizeIdentifier(tabName ?? "AddIns")}_{SanitizeIdentifier(panelName)}";

                    if (!string.IsNullOrEmpty(tabName))
                    {
                        source.AppendLine($"            RibbonPanel {panelVar} = app.CreateRibbonPanel(\"{tabName}\", \"{panelName}\");");
                    }
                    else
                    {
                        source.AppendLine($"            RibbonPanel {panelVar} = app.CreateRibbonPanel(\"{panelName}\");");
                    }

                    foreach (var commandData in panelGroup)
                    {
                        var commandSymbol = commandData.Symbol;
                        var attr = commandData.Attribute;
                        var buttonText = (string)attr.ConstructorArguments.First().Value;
                        var fullClassName = commandSymbol.ToDisplayString();
                        var iconName = GetAttributeProperty(attr, "Icon", string.Empty);
                        var tooltip = GetAttributeProperty(attr, "ToolTip", string.Empty);

                        source.AppendLine($"            var pbd_{commandSymbol.Name} = new PushButtonData(\"cmd_{commandSymbol.Name}\", \"{buttonText}\", assemblyPath, \"{fullClassName}\");");

                        if (!string.IsNullOrEmpty(tooltip))
                        {
                            source.AppendLine($"            pbd_{commandSymbol.Name}.ToolTip = \"{tooltip}\";");
                        }

                        if (!string.IsNullOrEmpty(iconName))
                        {
                            // ★★ 核心修正：为每个按钮的图标处理逻辑创建一个独立的局部作用域 ★★
                            source.AppendLine("            {");
                            source.AppendLine($"                string iconExtension = System.IO.Path.GetExtension(\"{iconName}\");");
                            source.AppendLine("                if (!string.IsNullOrEmpty(iconExtension))");
                            source.AppendLine("                {");
                            source.AppendLine($"                    pbd_{commandSymbol.Name}.LargeImage = LoadImageFromEmbeddedResource(assembly, \"{iconName}\");");
                            source.AppendLine("                }");
                            source.AppendLine("                else");
                            source.AppendLine("                {");
                            source.AppendLine($"                    if (_resourceManager != null)");
                            source.AppendLine($"                    {{");
                            source.AppendLine($"                        try");
                            source.AppendLine($"                        {{");
                            source.AppendLine($"                            var iconBitmap = (Bitmap)_resourceManager.GetObject(\"{iconName}\");");
                            source.AppendLine($"                            pbd_{commandSymbol.Name}.LargeImage = BitmapToImageSource(iconBitmap);");
                            source.AppendLine($"                        }}");
                            source.AppendLine($"                        catch {{ /* 资源未找到或加载失败 */ }}");
                            source.AppendLine($"                    }}");
                            source.AppendLine("                }");
                            source.AppendLine("            }"); // ★★ 结束局部作用域 ★★
                        }

                        source.AppendLine($"            {panelVar}.AddItem(pbd_{commandSymbol.Name});");
                    }
                }
            }

            source.AppendLine("            return Result.Succeeded;");
            source.AppendLine("        }");
            source.AppendLine("        public Result OnShutdown(UIControlledApplication app) => Result.Succeeded;");
            source.AppendLine("    }");
            source.AppendLine("}");
            return source.ToString();
        }

        /// <summary>
        /// 一个辅助方法，用于从特性的数据中安全地读取一个具名属性。
        /// </summary>
        private static T GetAttributeProperty<T>(AttributeData attributeData, string propertyName, T defaultValue)
        {
            var namedArgument = attributeData.NamedArguments.FirstOrDefault(kvp => kvp.Key == propertyName);
            if (namedArgument.Value.Value == null)
            {
                return defaultValue;
            }
            return (T)namedArgument.Value.Value;
        }

        /// <summary>
        /// 一个辅助方法，用于清理字符串，使其可以用作C#变量名。
        /// </summary>
        private static string SanitizeIdentifier(string name)
        {
            return Regex.Replace(name, @"[^\w]", "_");
        }
    }
}