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

using System.CommandLine;
using Humanizer;
using MycroForge.CLI.CodeGen;
using MycroForge.CLI.Commands.Interfaces;

namespace MycroForge.CLI.Commands;

public partial class MycroForge
{
    public partial class Db
    {
        public partial class Generate
        {
            public class Entity : Command, ISubCommandOf<Generate>
            {
                private record ColumnDefinition(string Name, string NativeType, string OrmType);

                private static readonly string[] Template =
                [
                    "from sqlalchemy import %type_imports%",
                    "from sqlalchemy.orm import Mapped, mapped_column",
                    "from db.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 };

                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)
                {
                    _context.AssertDirectoryExists(Features.Db.FeatureName);

                    var path = string.Empty;
                    if (name.Split(':').Select(s => s.Trim()).ToArray() is { Length: 2 } fullName)
                    {
                        path = fullName[0];
                        name = fullName[1];
                    }

                    var _columns = GetColumnDefinitions(columns.ToArray());
                    var className = name.Underscore().Pascalize();
                    var typeImports = string.Join(", ", _columns.Select(c => c.OrmType.Split('(').First()).Distinct());
                    var columnDefinitions = string.Join("\n\t", _columns.Select(ColumnToString));

                    var code = string.Join('\n', Template);
                    code = code.Replace("%type_imports%", typeImports);
                    code = code.Replace("%class_name%", className);
                    code = code.Replace("%table_name%", name.Underscore().ToLower().Pluralize());
                    code = code.Replace("%column_definitions%", columnDefinitions);
                    
                    var folderPath = Path.Join($"{Features.Db.FeatureName}/entities", path);
                    var fileName = $"{name.ToLower()}.py";
                    var filePath = Path.Join(folderPath, fileName);
                    await _context.CreateFile(filePath, code);

                    var importPathParts = new[] { path, 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 DbEnvUpdater(env, importPath, className).Rewrite();
                    await _context.WriteFile($"{Features.Db.FeatureName}/env.py", env);
                }

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

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

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

                    return definitions;
                }

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