using System.Linq.Expressions;
using System.Reflection;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Diagnostics;

namespace EFQueryCustomization;

class QueryInterceptor
{
    public static async Task Demo()
    {
        using var db = new MyContext();

        await db.PrepareDatabase();

        await db.Set<Customer>().OrderBy(x => DateTime.Now).LoadAsync();
    }

    class MyContext : MyDbContextBase
    {
        protected override void OnModelCreating(ModelBuilder modelBuilder)
        {
            base.OnModelCreating(modelBuilder);
            modelBuilder.Entity<Customer>();
        }

        protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
        {
            base.OnConfiguring(optionsBuilder);
            optionsBuilder.AddInterceptors(new KeyOrderingExpressionInterceptor());
        }
    }

    class Customer : IEntity
    {
        public int Id { get; set; }
        public string Name { get; set; }
    }

    class KeyOrderingExpressionInterceptor : IQueryExpressionInterceptor
    {
        public Expression QueryCompilationStarting(Expression queryExpression, QueryExpressionEventData eventData)
            => new KeyOrderingExpressionVisitor().Visit(queryExpression);

        class KeyOrderingExpressionVisitor : ExpressionVisitor
        {
            static readonly MethodInfo ThenByMethod = typeof(Queryable).GetMethods().Single(m => m.Name == nameof(Queryable.ThenBy) && m.GetParameters().Length == 2);

            protected override Expression VisitMethodCall(MethodCallExpression? methodCallExpression)
            {
                var methodInfo = methodCallExpression!.Method;
                if (methodInfo.DeclaringType == typeof(Queryable)
                    && methodInfo.Name == nameof(Queryable.OrderBy)
                    && methodInfo.GetParameters().Length == 2)
                {
                    var sourceType = methodCallExpression.Type.GetGenericArguments()[0];
                    if (sourceType.IsAssignableTo(typeof(IEntity)))
                    {
                        var lambdaExpression = (LambdaExpression)((UnaryExpression)methodCallExpression.Arguments[1]).Operand;
                        var entityParameterExpression = lambdaExpression.Parameters[0];

                        return Expression.Call(
                            ThenByMethod.MakeGenericMethod(
                                sourceType,
                                typeof(int)),
                            base.VisitMethodCall(methodCallExpression),
                            Expression.Lambda(
                                typeof(Func<,>).MakeGenericType(entityParameterExpression.Type, typeof(int)),
                                Expression.Property(entityParameterExpression, nameof(IEntity.Id)),
                                entityParameterExpression));
                    }
                }

                return base.VisitMethodCall(methodCallExpression);
            }
        }
    }

    interface IEntity
    {
        int Id { get; }
    }
}
