mycroforge/MycroForge.CLI/CodeGen/MainModifier.cs

using Antlr4.Runtime.Tree;
using MycroForge.Parsing;

namespace MycroForge.CLI.CodeGen;

public class MainModifier : PythonSourceModifier
{
    private PythonParser.Import_fromContext? _lastEntityImport;
    private PythonParser.Import_fromContext? _lastAssociationImport;
    private PythonParser.Import_fromContext? _lastRouterImport;
    private PythonParser.Import_fromContext? _lastRouterInclude;

    private readonly List<string> _lastEntityImportBuffer;
    private readonly List<string> _lastAssociationImportBuffer;
    private readonly List<string> _lastRouterImportBuffer;
    private readonly List<string> _lastRouterIncludeBuffer;

    public MainModifier(string source) : base(source)
    {
        _lastEntityImportBuffer = new();
        _lastAssociationImportBuffer = new();
        _lastRouterImportBuffer = new();
        _lastRouterIncludeBuffer = new();
    }

    public void Initialize()
    {
        var tree = Parser.file_input();

        Visit(tree);

        if (_lastEntityImport is not null)
            _lastEntityImportBuffer.Add(GetOriginalText(_lastEntityImport));

        if (_lastAssociationImport is not null)
            _lastAssociationImportBuffer.Add(GetOriginalText(_lastAssociationImport));

        if (_lastRouterImport is not null)
            _lastRouterImportBuffer.Add(GetOriginalText(_lastRouterImport));

        if (_lastRouterInclude is not null)
            _lastRouterIncludeBuffer.Add(GetOriginalText(_lastRouterInclude));
    }

    private string ToImportString(string from, string import) => $"from {from} import {import}";

    public void ImportEntity(string from, string import)
    {
        _lastEntityImportBuffer.Add(ToImportString(from, import));
    }

    public void ImportAssociation(string from, string import)
    {
        _lastAssociationImportBuffer.Add(ToImportString(from, import));
    }

    public void ImportRouter(string from, string import)
    {
        _lastRouterImportBuffer.Add(ToImportString(from, import));
    }

    public void IncludeRouter(string prefix, string router)
    {
        _lastRouterImportBuffer.Add($"app.include_router(prefix=\"/{prefix}\", router={router}.router)");
        // _lastRouterImportBuffer.Add($"app.include_router(prefix=\"/{prefix}\", router={router}.router)");
    }

    public override string Rewrite()
    {
        if (_lastEntityImport is not null)
            Rewrite(_lastEntityImport, _lastEntityImportBuffer.ToArray());

        if (_lastAssociationImport is not null)
            Rewrite(_lastAssociationImport, _lastAssociationImportBuffer.ToArray());

        if (_lastRouterImport is not null)
            Rewrite(_lastRouterImport, _lastRouterImportBuffer.ToArray());

        return Rewriter.GetText();
    }

    // public override object? VisitPrimary(PythonParser.PrimaryContext context)
    // {
    //     // Console.WriteLine(GetOriginalText(context));
    //     return base.VisitPrimary(context);
    // }
    //
    // public override object? VisitName_or_attr(PythonParser.Name_or_attrContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitName_or_attr(context);
    // }
    //
    // public override object? VisitStatement(PythonParser.StatementContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitStatement(context);
    // }

    // public override object? VisitDotted_name(PythonParser.Dotted_nameContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitDotted_name(context);
    // }
    //
    // public override object? VisitDotted_name(PythonParser.Dotted_nameContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //
    //     return base.VisitDotted_name(context);
    // }
    //
    // public override object? VisitDotted_as_names(PythonParser.Dotted_as_namesContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //
    //     return base.VisitDotted_as_names(context);
    // }
    //
    public override object? VisitErrorNode(IErrorNode node)
    {
        Console.WriteLine(node.GetText());
        return base.VisitErrorNode(node);
    }

    // public override object? VisitValue_pattern(PythonParser.Value_patternContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitValue_pattern(context);
    // }
    //
    // public override object? VisitStar_atom(PythonParser.Star_atomContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitStar_atom(context);
    // }
    //
    // public override object? VisitExpression(PythonParser.ExpressionContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitExpression(context);
    // }
    //
    // public override object? VisitT_primary(PythonParser.T_primaryContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitT_primary(context);
    // }

    // public override object? VisitAttr(PythonParser.AttrContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitAttr(context);
    // }

    // public override object? VisitT_primary(PythonParser.T_primaryContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitT_primary(context);
    // }

    // public override object? VisitAwait_primary(PythonParser.Await_primaryContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitAwait_primary(context);
    // }

    // public override object? VisitTarget_with_star_atom(PythonParser.Target_with_star_atomContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitTarget_with_star_atom(context);
    // }

    // public override object? VisitSingle_subscript_attribute_target(
    //     PythonParser.Single_subscript_attribute_targetContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitSingle_subscript_attribute_target(context);
    // }
    //
    // public override object? VisitSingle_target(PythonParser.Single_targetContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitSingle_target(context);
    // }

    // public override object? VisitName_or_attr(PythonParser.Name_or_attrContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitName_or_attr(context);
    // }

    // public override object? VisitNamed_expression(PythonParser.Named_expressionContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //
    //     return base.VisitNamed_expression(context);
    // }

    // public override object? VisitPrimary(PythonParser.PrimaryContext context)
    // {
    //     Console.WriteLine(GetOriginalText(context));
    //     return base.VisitPrimary(context);
    // }

    public override object? VisitImport_from(PythonParser.Import_fromContext context)
    {
        var text = GetOriginalText(context);

        if (text.StartsWith($"from {Features.Db.FeatureName}.entities.associations"))
            _lastAssociationImport = context;

        if (text.StartsWith($"from {Features.Db.FeatureName}.entities"))
            _lastEntityImport = context;

        if (text.StartsWith($"from {Features.Api.FeatureName}.routers"))
            _lastRouterImport = context;

        return base.VisitImport_from(context);
    }
}