mycroforge/MycroForge.CLI/CodeGen/EntityLinker.EntityModel.cs

using Humanizer;

namespace MycroForge.CLI.CodeGen;

public partial class EntityLinker
{
    public class EntityModel
    {
        private readonly Source _source;
        private readonly string _className;
        private readonly string _path;

        private readonly List<SourceMatch> _importContexts;
        private readonly List<string> _importsBuffer;
        private readonly List<SourceMatch> _classContexts;
        private SourceMatch _tableContext;
        private readonly List<SourceMatch> _columnContexts;
        private readonly List<string> _columnsBuffer;

        private SourceMatch LastColumn => _columnContexts.Last();
        private SourceMatch LastImport => _importContexts.Last();

        public string ClassName => _className;
        public string Path => _path;
        public string FieldName => _className.Underscore().ToLower();

        public string TableName => _tableContext.Text
            .Replace("__tablename__", string.Empty)
            .Replace("=", string.Empty)
            .Replace("\"", string.Empty)
            .Trim();

        public EntityModel(string className, string path, string source)
        {
            _source = new Source(source);
            _className = className;
            _path = path;
            _importContexts = new();
            _importsBuffer = new();
            _classContexts = new();
            _tableContext = default!;
            _columnContexts = new();
            _columnsBuffer = new();

            var classContexts = _source.FindAll(@"class\s+.*\s*\(EntityBase\)\s*:\s");
            if (!classContexts.Any(c => c.Text.Contains(_className)))
                throw new Exception($"Entity {_className} was not found in {_path}.");
        }

        public void Initialize()
        {
            ReadImports();

            AssertClassExists();

            AssertTableNameExists();

            ReadColumns();

            InsertRelationshipImport();
            
            InsertForeignKeyImport();
        }

        private void ReadImports()
        {
            _importContexts.AddRange(_source.FindAll(@"from\s+.+\s+import\s+.*\s"));
        }

        private void AssertClassExists()
        {
            _classContexts.AddRange(_source.FindAll(@"class\s+.+\s*\(EntityBase\)\s*:\s"));
            if (!_classContexts.Any(c => c.Text.Contains(_className)))
                throw new Exception($"Entity {_className} was not found in {_path}.");
        }

        private void AssertTableNameExists()
        {
            _tableContext = _source.Find(@"__tablename__\s*=\s*.+\s");
            if (string.IsNullOrEmpty(_tableContext.Text))
                throw new Exception($"__tablename__ definition for Entity {_className} was not found in {_path}.");
        }

        private void ReadColumns()
        {
            _columnContexts.AddRange(
                _source.FindAll(@".+\s*:\s*Mapped\[.+\]\s*=\s*(relationship|mapped_column)\(.+\)\s")
            );

            if (_columnContexts.Count == 0)
                throw new Exception($"Entity {_className} has no columns.");
        }

        private void InsertRelationshipImport()
        {
            var relationship = _importContexts.FirstOrDefault(
                import => import.Text.Contains("sqlalchemy.orm") && import.Text.Contains("relationship")
            );

            if (relationship is null)
                _importsBuffer.Add("from sqlalchemy.orm import relationship");
        }

        private void InsertForeignKeyImport()
        {
            var foreignKey = _importContexts.FirstOrDefault(
                import => import.Text.Contains("sqlalchemy") && import.Text.Contains("ForeignKey")
            );

            if (foreignKey is null)
                _importsBuffer.Add("from sqlalchemy import ForeignKey");
        }

        public void AppendColumn(string text)
        {
            _columnsBuffer.Add(text);
        }

        public void AppendColumns(params string[] text)
        {
            _columnsBuffer.Add(string.Join('\n', text));
        }

        public void Import(string from, string import)
        {
            var exists = _importContexts.Any(ctx => ctx.Text.Contains(from) && ctx.Text.Contains(import));
            var buffered = _importsBuffer.Any(txt => txt.Contains(from) && txt.Contains(import));

            if (!exists && !buffered)
                _importsBuffer.Add($"from {from} import {import}");
        }

        public string Rewrite()
        {
            // The order matters! We must rewrite the columns first,
            // so that their indexes are not modified when inserting imports.
            _source.InsertMultiLine(LastColumn.EndIndex, _columnsBuffer.Append("\n").ToArray());
            _source.InsertMultiLine(LastImport.EndIndex, _importsBuffer.Append("\n").ToArray());
            return _source.Text;
        }
    }
}