using MycroForge.Parsing;

namespace MycroForge.CLI.CodeGen;

public class OrmEnvInitializer : PythonSourceModifier
{
    public OrmEnvInitializer(string source) : base(source)
    {
    }

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

        if (text != "from alembic import context") return null;

        Rewrite(context,
            text,
            "from orm.settings import OrmSettings",
            "from orm.entities.entity_base import EntityBase"
        );

        return base.VisitImport_from(context);
    }

    public override object? VisitAssignment(PythonParser.AssignmentContext context)
    {
        var text = GetOriginalText(context);

        if (text == "target_metadata = None")
        {
            Rewrite(context, "target_metadata = EntityBase.metadata");
        }
        else if (text == "url = config.get_main_option(\"sqlalchemy.url\")")
        {
            Rewrite(context, "url = OrmSettings.get_connectionstring()");
        }
        else if (text.StartsWith("connectable ="))
        {
            // Important note, the indent here is 4 spaces and not tab(s).
            const string indent = "    ";
            Rewrite(context, [
                "url = OrmSettings.get_connectionstring()",
                $"{indent}context.config.set_main_option('sqlalchemy.url', url)",
                $"{indent}{text}"
            ]);
        }

        return base.VisitAssignment(context);
    }
}