mycroforge/MycroForge.CLI/CodeGen/MainModifier.cs

using MycroForge.Core.CodeGen;

namespace MycroForge.CLI.CodeGen;

public class MainModifier
{
    private readonly Source _source;
    private readonly List<string> _importsBuffer;
    private readonly List<string> _routerIncludeBuffer;
    private SourceMatch? _lastImport;
    private SourceMatch? _lastRouterInclude;

    public MainModifier(string source)
    {
        _source = new Source(source);
        _importsBuffer = new();
        _routerIncludeBuffer = new();
    }

    public MainModifier Initialize()
    {
        _lastImport = _source
            .FindAll($@"from\s+.+\s+import\s+.+\s")
            .LastOrDefault();

        _lastRouterInclude = _source.Find(@"app\s*\.\s*include_router\((.|\s)+\)\s?");

        return this;
    }

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

    public MainModifier Import(string from, string import)
    {
        _importsBuffer.Add(ToImportString(from, import));
        return this;
    }

    public MainModifier IncludeRouter(string prefix, string router)
    {
        _routerIncludeBuffer.Add($"app.include_router(prefix=\"/{prefix}\", router={router}.router)");
        return this;
    }

    public string Rewrite()
    {
        // Make sure to insert the includes before the imports, if done the other way around,
        // the insertions of the includes will change the indexes of the imports.
        InsertIncludes();

        InsertImports();

        return _source.Text;
    }

    private void InsertImports()
    {
        if (_importsBuffer.Count == 0) return;

        if (_lastImport is not null)
        {
            _source.InsertMultiLine(_lastImport.EndIndex, _importsBuffer.ToArray());
        }
        else
        {
            _source.InsertMultiLineAtStart(_importsBuffer.ToArray());
        }
    }

    private void InsertIncludes()
    {
        if (_routerIncludeBuffer.Count == 0) return;

        // Prepend an empty string to the router include buffer,
        // this will ensure that the new entries are all on separate lines.
        var content = _routerIncludeBuffer.Prepend(string.Empty).ToArray();
        
        if (_lastRouterInclude is not null)
        {
            _source.InsertMultiLine(_lastRouterInclude.EndIndex, content);
        }
        else
        {
            _source.InsertMultiLineAtEnd(content);
        }
    }
}