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($"\napp.include_router(prefix=\"/{prefix}\", router={router}.router)");
        return this;
    }

    public string Rewrite()
    {
        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;

        if (_lastRouterInclude is not null)
        {
            _source.InsertMultiLine(
                _lastRouterInclude.EndIndex, _routerIncludeBuffer.ToArray()
            );
        }
        else
        {
            _source.InsertMultiLineAtEnd(_routerIncludeBuffer.ToArray());
        }
    }
}