mycroforge/MycroForge.CLI/CodeGen/EntityLinker.cs

using Humanizer;
using MycroForge.CLI.Commands;
using MycroForge.Core;

namespace MycroForge.CLI.CodeGen;

public partial class EntityLinker
{
    #region AssociationTable

    private static readonly string[] AssociationTable =
    [
        "from sqlalchemy import Column, ForeignKey, Table",
        $"from {Features.Db.FeatureName}.entities.entity_base import EntityBase",
        "",
        "%left_entity%_%right_entity%_mapping = Table(",
        "\t\"%left_entity%_%right_entity%_mapping\",",
        "\tEntityBase.metadata,",
        "\tColumn(\"%left_entity%_id\", ForeignKey(\"%left_table%.id\"), primary_key=True),",
        "\tColumn(\"%right_entity%_id\", ForeignKey(\"%right_table%.id\"), primary_key=True),",
        ")"
    ];

    #endregion

    private readonly ProjectContext _context;
    private readonly string _left;
    private readonly string _right;

    public EntityLinker(ProjectContext context, string left, string right)
    {
        _context = context;
        _left = left;
        _right = right;
    }

    public async Task OneToOne()
    {
        var left = await LoadEntity(_left);
        var right = await LoadEntity(_right);

        left.AppendColumn(
            $"\t{right.FieldName}: Mapped[\"{right.ClassName.Pascalize()}\"] = relationship(back_populates=\"{left.FieldName}\")"
        );

        right.AppendColumns([
            $"\t{left.FieldName}_id: Mapped[int] = mapped_column(ForeignKey(\"{left.TableName}.id\"))",
            $"\t{left.FieldName}: Mapped[\"{left.ClassName.Pascalize()}\"] = relationship(back_populates=\"{right.FieldName}\")"
        ]);

        await _context.WriteFile(left.Path, left.Rewrite());
        await _context.WriteFile(right.Path, right.Rewrite());
    }

    public async Task OneToMany()
    {
        var left = await LoadEntity(_left);
        var right = await LoadEntity(_right);

        left.Import(from: "typing", import: "List");
        left.AppendColumn(
            $"\t{right.FieldName.Pluralize()}: Mapped[List[\"{right.ClassName.Pascalize()}\"]] = relationship()");

        right.AppendColumns([
            $"\t{left.FieldName}_id: Mapped[int] = mapped_column(ForeignKey(\"{left.TableName}.id\"))",
            $"\t{left.FieldName}: Mapped[\"{left.ClassName.Pascalize()}\"] = relationship(back_populates=\"{right.FieldName}\")"
        ]);

        await _context.WriteFile(left.Path, left.Rewrite());
        await _context.WriteFile(right.Path, right.Rewrite());
    }

    public async Task ManyToOne()
    {
        var left = await LoadEntity(_left);
        var right = await LoadEntity(_right);

        left.Import(from: "typing", import: "Optional");
        left.AppendColumns([
            $"\t{right.FieldName}_id: Mapped[Optional[int]] = mapped_column(ForeignKey(\"{right.TableName}.id\"))",
            $"\t{right.FieldName}: Mapped[\"{right.ClassName.Pascalize()}\"] = relationship(back_populates=\"{left.FieldName.Pluralize()}\")"
        ]);
        left.Import(from: "typing", import: "List");

        right.AppendColumns([
            $"\t{left.FieldName}_id: Mapped[int] = mapped_column(ForeignKey(\"{left.TableName}.id\"))",
            $"\t{left.FieldName.Pluralize()}: Mapped[List[\"{left.ClassName.Pascalize()}\"]] = relationship(back_populates=\"{right.FieldName}\")"
        ]);

        await _context.WriteFile(left.Path, left.Rewrite());
        await _context.WriteFile(right.Path, right.Rewrite());
    }

    public async Task ManyToMany()
    {
        var left = await LoadEntity(_left);
        var right = await LoadEntity(_right);

        var associationTable = string.Join('\n', AssociationTable)
            .Replace("%left_entity%", left.ClassName.Underscore().ToLower())
            .Replace("%right_entity%", right.ClassName.Underscore().ToLower())
            .Replace("%left_table%", left.TableName)
            .Replace("%right_table%", right.TableName);
        
        var associationTablePath = Path.Join(
            Features.Db.FeatureName, 
            "entities", 
            "associations",
            $"{left.TableName.Singularize()}_{right.TableName.Singularize()}_mapping.py"
        );

        await _context.CreateFile(associationTablePath, associationTable);

        var associationTableImport = associationTablePath
                .Replace(".py", "")
                .Replace('/', '.')
                .Replace('\\', '.')
                .Split('.')
            ;

        var associationTableImportPath = string.Join('.', associationTableImport);
        var associationTableImportName = associationTableImport[^1];

        left.Import(from: string.Join('.', associationTableImportPath), import: associationTableImportName);
        left.Import(from: "typing", import: "List");
        right.Import(from: string.Join('.', associationTableImportPath), import: associationTableImportName);
        right.Import(from: "typing", import: "List");

        left.AppendColumns([
            $"\t{right.FieldName.Pluralize()}: Mapped[List[\"{right.ClassName.Pascalize()}\"]] = relationship(back_populates=\"{left.FieldName.Pluralize()}\", secondary={associationTableImportName})"
        ]);

        right.AppendColumns([
            $"\t{left.FieldName.Pluralize()}: Mapped[List[\"{left.ClassName.Pascalize()}\"]] = relationship(back_populates=\"{right.FieldName.Pluralize()}\", secondary={associationTableImportName})"
        ]);

        await _context.WriteFile(left.Path, left.Rewrite());
        await _context.WriteFile(right.Path, right.Rewrite());

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

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

    private async Task<EntityModel> LoadEntity(string name)
    {
        var fqn = new FullyQualifiedName(name);
        var path = Path.Join(Features.Db.FeatureName, "entities");

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

        path = Path.Join(path, $"{fqn.SnakeCasedName}.py");
        var entity = new EntityModel(fqn.PascalizedName, path, await _context.ReadFile(path));
        entity.Initialize();
        return entity;
    }
}