using MycroForge.Parsing;

namespace MycroForge.CLI.CodeGen;

public class OrmEnvUpdater : PythonSourceModifier
{
    private readonly string _moduleName;
    private readonly string _className;

    public OrmEnvUpdater(string source, string moduleName, string className) : base(source)
    {
        _moduleName = moduleName;
        _className = className;
    }

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

        if (text != "from orm.entities.entity_base import EntityBase") return null;

        Rewrite(context, [
            text,
            $"from orm.entities.{_moduleName} import {_className}"
        ]);

        return base.VisitImport_from(context);
    }
}