sqrtspace-dotnet/src/SqrtSpace.SpaceTime.EntityFramework/SpaceTimeQueryExtensions.cs

using System.Linq.Expressions;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Infrastructure;
using Microsoft.EntityFrameworkCore.Query;
using SqrtSpace.SpaceTime.Core;

namespace SqrtSpace.SpaceTime.EntityFramework;

/// <summary>
/// Extended query extensions for SpaceTime optimizations
/// </summary>
public static class SpaceTimeQueryExtensions
{
    /// <summary>
    /// Configures the query to use external sorting for large datasets
    /// </summary>
    public static IQueryable<T> UseExternalSorting<T>(this IQueryable<T> query) where T : class
    {
        // Mark the query for external sorting
        return query.TagWith("SpaceTime:UseExternalSorting");
    }

    /// <summary>
    /// Streams query results asynchronously for memory efficiency
    /// </summary>
    public static async IAsyncEnumerable<T> StreamQueryResultsAsync<T>(
        this IQueryable<T> query,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken = default) where T : class
    {
        var context = GetDbContext(query);
        if (context == null)
        {
            // Fallback to regular async enumeration
            await foreach (var item in query.AsAsyncEnumerable().WithCancellation(cancellationToken))
            {
                yield return item;
            }
            yield break;
        }

        // Get total count for batch size calculation
        var totalCount = await query.CountAsync(cancellationToken);
        var batchSize = SpaceTimeCalculator.CalculateSqrtInterval(totalCount);

        // Stream in batches
        for (int offset = 0; offset < totalCount; offset += batchSize)
        {
            var batch = await query
                .Skip(offset)
                .Take(batchSize)
                .ToListAsync(cancellationToken);

            foreach (var item in batch)
            {
                yield return item;
            }

            // Clear change tracker to prevent memory buildup
            context.ChangeTracker.Clear();
        }
    }

    /// <summary>
    /// Processes query results in √n-sized batches with checkpoint support
    /// </summary>
    public static async IAsyncEnumerable<IReadOnlyList<T>> BatchBySqrtNAsync<T>(
        this IQueryable<T> query,
        string? checkpointId = null,
        bool resumeFromCheckpoint = false,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken = default) where T : class
    {
        var context = GetDbContext(query);
        var options = context?.GetService<SpaceTimeOptions>();
        CheckpointManager? checkpointManager = null;

        if (!string.IsNullOrEmpty(checkpointId) && options?.EnableQueryCheckpointing == true)
        {
            checkpointManager = new CheckpointManager(options.CheckpointDirectory);
        }

        try
        {
            var totalCount = await query.CountAsync(cancellationToken);
            var batchSize = SpaceTimeCalculator.CalculateSqrtInterval(totalCount);
            var startOffset = 0;

            // Resume from checkpoint if requested
            if (resumeFromCheckpoint && checkpointManager != null)
            {
                var checkpoint = await checkpointManager.RestoreCheckpointAsync<QueryCheckpoint>(checkpointId!);
                if (checkpoint != null)
                {
                    startOffset = checkpoint.ProcessedCount;
                }
            }

            for (int offset = startOffset; offset < totalCount; offset += batchSize)
            {
                var batch = await query
                    .Skip(offset)
                    .Take(batchSize)
                    .ToListAsync(cancellationToken);

                if (batch.Count == 0)
                    yield break;

                yield return batch;

                // Save checkpoint
                if (checkpointManager != null && !string.IsNullOrEmpty(checkpointId))
                {
                    await checkpointManager.CreateCheckpointAsync(new QueryCheckpoint
                    {
                        ProcessedCount = offset + batch.Count,
                        TotalCount = totalCount
                    }, checkpointId);
                }

                // Clear change tracker if enabled
                if (context != null && options?.EnableSqrtNChangeTracking == true)
                {
                    var trackedCount = context.ChangeTracker.Entries().Count();
                    if (trackedCount > (options.MaxTrackedEntities ?? batchSize))
                    {
                        context.ChangeTracker.Clear();
                    }
                }
            }
        }
        finally
        {
            checkpointManager?.Dispose();
        }
    }

    private static DbContext? GetDbContext<T>(IQueryable<T> query)
    {
        if (query.Provider is IInfrastructure<IServiceProvider> infrastructure)
        {
            var context = infrastructure.Instance.GetService(typeof(DbContext)) as DbContext;
            return context;
        }

        // Fallback: try reflection
        var provider = query.Provider;
        var contextProperty = provider.GetType().GetProperty("Context");
        return contextProperty?.GetValue(provider) as DbContext;
    }

    private class QueryCheckpoint
    {
        public int ProcessedCount { get; set; }
        public int TotalCount { get; set; }
    }
}

/// <summary>
/// Extension methods for DbContext bulk operations
/// </summary>
public static class SpaceTimeDbContextExtensions
{
    /// <summary>
    /// Performs bulk insert with √n buffering for memory efficiency
    /// </summary>
    public static async Task BulkInsertWithSqrtNBufferingAsync<T>(
        this DbContext context,
        IEnumerable<T> entities,
        CancellationToken cancellationToken = default) where T : class
    {
        ArgumentNullException.ThrowIfNull(context);
        ArgumentNullException.ThrowIfNull(entities);

        var options = context.GetService<SpaceTimeOptions>();
        var entityList = entities as IList<T> ?? entities.ToList();
        var totalCount = entityList.Count;
        var batchSize = SpaceTimeCalculator.CalculateSqrtInterval(totalCount);

        // Disable auto-detect changes for performance
        var originalAutoDetectChanges = context.ChangeTracker.AutoDetectChangesEnabled;
        context.ChangeTracker.AutoDetectChangesEnabled = false;

        try
        {
            for (int i = 0; i < totalCount; i += batchSize)
            {
                var batch = entityList.Skip(i).Take(batchSize);
                
                await context.AddRangeAsync(batch, cancellationToken);
                await context.SaveChangesAsync(cancellationToken);
                
                // Clear change tracker after each batch to prevent memory buildup
                context.ChangeTracker.Clear();
            }
        }
        finally
        {
            context.ChangeTracker.AutoDetectChangesEnabled = originalAutoDetectChanges;
        }
    }

    /// <summary>
    /// Performs bulk update with √n buffering for memory efficiency
    /// </summary>
    public static async Task BulkUpdateWithSqrtNBufferingAsync<T>(
        this DbContext context,
        IEnumerable<T> entities,
        CancellationToken cancellationToken = default) where T : class
    {
        ArgumentNullException.ThrowIfNull(context);
        ArgumentNullException.ThrowIfNull(entities);

        var entityList = entities as IList<T> ?? entities.ToList();
        var totalCount = entityList.Count;
        var batchSize = SpaceTimeCalculator.CalculateSqrtInterval(totalCount);

        // Disable auto-detect changes for performance
        var originalAutoDetectChanges = context.ChangeTracker.AutoDetectChangesEnabled;
        context.ChangeTracker.AutoDetectChangesEnabled = false;

        try
        {
            for (int i = 0; i < totalCount; i += batchSize)
            {
                var batch = entityList.Skip(i).Take(batchSize);
                
                context.UpdateRange(batch);
                await context.SaveChangesAsync(cancellationToken);
                
                // Clear change tracker after each batch
                context.ChangeTracker.Clear();
            }
        }
        finally
        {
            context.ChangeTracker.AutoDetectChangesEnabled = originalAutoDetectChanges;
        }
    }

    /// <summary>
    /// Performs bulk delete with √n buffering for memory efficiency
    /// </summary>
    public static async Task BulkDeleteWithSqrtNBufferingAsync<T>(
        this DbContext context,
        IEnumerable<T> entities,
        CancellationToken cancellationToken = default) where T : class
    {
        ArgumentNullException.ThrowIfNull(context);
        ArgumentNullException.ThrowIfNull(entities);

        var entityList = entities as IList<T> ?? entities.ToList();
        var totalCount = entityList.Count;
        var batchSize = SpaceTimeCalculator.CalculateSqrtInterval(totalCount);

        // Disable auto-detect changes for performance
        var originalAutoDetectChanges = context.ChangeTracker.AutoDetectChangesEnabled;
        context.ChangeTracker.AutoDetectChangesEnabled = false;

        try
        {
            for (int i = 0; i < totalCount; i += batchSize)
            {
                var batch = entityList.Skip(i).Take(batchSize);
                
                context.RemoveRange(batch);
                await context.SaveChangesAsync(cancellationToken);
                
                // Clear change tracker after each batch
                context.ChangeTracker.Clear();
            }
        }
        finally
        {
            context.ChangeTracker.AutoDetectChangesEnabled = originalAutoDetectChanges;
        }
    }
}