mycroforge/MycroForge.CLI/Commands/MycroForge.Db.Generate.Entity.cs

using System.CommandLine;
using System.Text.RegularExpressions;
using Humanizer;
using MycroForge.CLI.CodeGen;
using MycroForge.Core.Contract;
using MycroForge.Core;

namespace MycroForge.CLI.Commands;

public partial class MycroForge
{
    public partial class Db
    {
        public partial class Generate
        {
            public class Entity : Command, ISubCommandOf<Generate>
            {
                #region Hidden region

                private static string[] SqlAlchemyTypes =
                [
                    "BigInteger",
                    "Boolean",
                    "Date",
                    "DateTime",
                    "Enum",
                    "Double",
                    "Float",
                    "Integer",
                    "Interval",
                    "LargeBinary",
                    "MatchType",
                    "Numeric",
                    "PickleType",
                    "SchemaType",
                    "SmallInteger",
                    "String",
                    "Text",
                    "Time",
                    "Unicode",
                    "UnicodeText",
                    "Uuid",
                    "ARRAY",
                    "BIGINT",
                    "BINARY",
                    "BLOB",
                    "BOOLEAN",
                    "CHAR",
                    "CLOB",
                    "DATE",
                    "DATETIME",
                    "DECIMAL",
                    "DOUBLE",
                    "DOUBLE_PRECISION",
                    "FLOAT",
                    "INT",
                    "JSON",
                    "INTEGER",
                    "NCHAR",
                    "NVARCHAR",
                    "NUMERIC",
                    "REAL",
                    "SMALLINT",
                    "TEXT",
                    "TIME",
                    "TIMESTAMP",
                    "UUID",
                    "VARBINARY",
                    "VARCHAR"
                ];

                private static readonly Regex SqlAlchemyTypeRegex = new(@".*\(.*\)");

                private static readonly string[] Template =
                [
                    "from sqlalchemy import %sqlalchemy_imports%",
                    "from sqlalchemy.orm import Mapped, mapped_column",
                    $"from {Features.Db.FeatureName}.entities.entity_base import EntityBase",
                    "",
                    "class %class_name%(EntityBase):",
                    "\t__tablename__ = \"%table_name%\"",
                    "\tid: Mapped[int] = mapped_column(primary_key=True)",
                    "\t%column_definitions%",
                    "",
                    "\tdef __repr__(self) -> str:",
                    "\t\treturn f\"%class_name%(id={self.id!r})\""
                ];

                private static readonly Argument<string> NameArgument =
                    new(name: "name", description: string.Join('\n', [
                        "The name of the database entity",
                        "",
                        "Supported formats:",
                        "\tEntity",
                        "\tpath/relative/to/entities:Entity",
                    ]));

                private static readonly Option<IEnumerable<string>> ColumnsOption =
                    new(aliases: ["--column", "-c"], description: string.Join('\n', [
                        "Specify the fields to add.",
                        "",
                        "Format:",
                        "\t<name>:<native_type>:<orm_type>",
                        "\t",
                        "\t<name> = Name of the column",
                        "\t<native_type> = The native Python type",
                        "\t<orm_type> = The SQLAlchemy type",
                        "",
                        "Example:",
                        "\tfirst_name:str:String(255)",
                    ])) { AllowMultipleArgumentsPerToken = true };

                #endregion

                private record ColumnDefinition(string Name, string NativeType, string SqlAlchemyType);

                private readonly ProjectContext _context;

                public Entity(ProjectContext context) : base("entity", "Generate and database entity")
                {
                    _context = context;
                    AddAlias("e");
                    AddArgument(NameArgument);
                    AddOption(ColumnsOption);
                    this.SetHandler(ExecuteAsync, NameArgument, ColumnsOption);
                }

                private async Task ExecuteAsync(string name, IEnumerable<string> columns)
                {
                    var fqn = new FullyQualifiedName(name);
                    var folderPath = Path.Join(Features.Db.FeatureName, "entities");

                    if (fqn.HasNamespace)
                        folderPath = Path.Join(folderPath, fqn.Namespace);

                    var sqlAlchemyColumn = GetColumnDefinitions(columns.ToArray());
                    var distinctSqlAlchemyColumnTypes = sqlAlchemyColumn
                        .Select(c => c.SqlAlchemyType.Split('(').First())
                        .Distinct();
                    
                    var sqlAlchemyImport = string.Join(", ", distinctSqlAlchemyColumnTypes);
                    var columnDefinitions = string.Join("\n  ", sqlAlchemyColumn.Select(ColumnToString));

                    var code = string.Join('\n', Template);
                    code = code.Replace("%sqlalchemy_imports%", sqlAlchemyImport);
                    code = code.Replace("%class_name%", fqn.PascalizedName);
                    code = code.Replace("%table_name%", fqn.SnakeCasedName.Pluralize());
                    code = code.Replace("%column_definitions%", columnDefinitions);

                    var fileName = $"{fqn.SnakeCasedName}.py";
                    var filePath = Path.Join(folderPath, fileName);
                    await _context.CreateFile(filePath, code);

                    var importPathParts = new[] { folderPath, fileName.Replace(".py", "") }
                        .Where(s => !string.IsNullOrEmpty(s));

                    var importPath = string.Join('.', importPathParts)
                        .Replace('/', '.')
                        .Replace('\\', '.')
                        .Underscore()
                        .ToLower();

                    var env = await _context.ReadFile($"{Features.Db.FeatureName}/env.py");
                    env = new DbEnvModifier(env, importPath, fqn.PascalizedName).Rewrite();
                    await _context.WriteFile($"{Features.Db.FeatureName}/env.py", env);

                    var main = await _context.ReadFile("main.py");
                    main = new MainModifier(main).Initialize().Import(importPath, fqn.PascalizedName).Rewrite();
                    await _context.WriteFile("main.py", main);
                }

                private List<ColumnDefinition> GetColumnDefinitions(string[] columns)
                {
                    var definitions = new List<ColumnDefinition>();

                    foreach (var column in columns)
                    {
                        if (column.Split(':') is not { Length: 3 } definition)
                            throw new Exception($"Column definition {column} is invalid.");

                        definitions.Add(new ColumnDefinition(definition[0], definition[1], definition[2]));
                    }

                    ValidateSqlAlchemyColumnTypes(definitions);

                    return definitions;
                }

                private static void ValidateSqlAlchemyColumnTypes(List<ColumnDefinition> definitions)
                {
                    foreach (var column in definitions)
                    {
                        if (!SqlAlchemyTypeRegex.IsMatch(column.SqlAlchemyType))
                        {
                            var message = new[]
                            {
                                $"SQLAlchemy column definition {column.SqlAlchemyType} was not properly defined.",
                                "Add parentheses and specify parameters if required, an example is provided below.",
                                "  String(255)",
                                "",
                                "Available options are:",
                                string.Join(Environment.NewLine, SqlAlchemyTypes.Select(type => $"  - {type}"))
                            };

                            throw new(string.Join(Environment.NewLine, message));
                        }

                        var type = column.SqlAlchemyType.Split('(').First();
                        
                        if (!SqlAlchemyTypes.Contains(type))
                        {
                            var message = string.Join(Environment.NewLine, [
                                $"SQLAlchemy column type '{column.SqlAlchemyType}' is not valid, available options are:",
                                string.Join(Environment.NewLine, SqlAlchemyTypes.Select(type => $"  - {type}"))
                            ]);

                            throw new(message);
                        }
                    }
                }

                private static string ColumnToString(ColumnDefinition definition) =>
                    $"{definition.Name}: Mapped[{definition.NativeType}] = mapped_column({definition.SqlAlchemyType})";
            }
        }
    }
}