using Humanizer;
using MycroForge.Parsing;

namespace MycroForge.CLI.CodeGen;

public partial class EntityLinker
{
    public class EntityModel : PythonSourceModifier
    {
        private readonly string _className;
        private readonly string _path;

        private readonly List<PythonParser.Import_fromContext> _importCtxs;
        private readonly List<string> _importsBuffer;
        private PythonParser.Import_fromContext LastImport => _importCtxs.Last();

        private readonly List<PythonParser.Class_defContext> _classCtxs;
        private PythonParser.AssignmentContext _tableCtx;

        private readonly List<PythonParser.AssignmentContext> _columnCtxs;
        private readonly List<string> _columnsBuffer;
        private PythonParser.AssignmentContext LastColumn => _columnCtxs.Last();

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

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

        public EntityModel(string className, string path, string source) : base(source)
        {
            _className = className;
            _path = path;
            _importCtxs = new();
            _importsBuffer = new();
            _classCtxs = new();
            _tableCtx = default!;
            _columnCtxs = new();
            _columnsBuffer = new();
        }

        public void Initialize()
        {
            var tree = Parser.file_input();

            Visit(tree);

            if (!_classCtxs.Any(c => GetOriginalText(c).Contains(_className)))
                throw new Exception($"Entity {_className} was not found in {_path}.");

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

            _importsBuffer.Add(GetOriginalText(LastImport));
            _columnsBuffer.Add(GetOriginalText(LastColumn));

            InsertRelationshipImport();
            InsertForeignKeyImport();
        }

        private void InsertRelationshipImport()
        {
            var relationship = _importCtxs.FirstOrDefault(import =>
            {
                var text = GetOriginalText(import);
                return text.Contains("sqlalchemy.orm") && text.Contains("relationship");
            });

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

        private void InsertForeignKeyImport()
        {
            var foreignKey = _importCtxs.FirstOrDefault(import =>
            {
                var text = GetOriginalText(import);
                return text.Contains("sqlalchemy") && text.Contains("ForeignKey");
            });

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

        public override object? VisitImport_from(PythonParser.Import_fromContext context)
        {
            _importCtxs.Add(context);
            return base.VisitImport_from(context);
        }

        public override object? VisitClass_def(PythonParser.Class_defContext context)
        {
            _classCtxs.Add(context);
            return base.VisitClass_def(context);
        }

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

            if (text.StartsWith("__tablename__"))
                _tableCtx = context;

            if (text.Contains("Mapped["))
                _columnCtxs.Add(context);

            return base.VisitAssignment(context);
        }

        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 isExisting = _importCtxs.Select(GetOriginalText).Any(ctx => ctx.Contains(from) && ctx.Contains(import));
            var isBuffered = _importsBuffer.Any(txt => txt.Contains(from) && txt.Contains(import));

            if (!isExisting && !isBuffered)
                _importsBuffer.Add($"from {from} import {import}");
        }

        public override string Rewrite()
        {
            Rewrite(LastImport, _importsBuffer.ToArray());
            Rewrite(LastColumn, _columnsBuffer.ToArray());
            return Rewriter.GetText();
        }
    }
}