using System.Data;
using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Storage;

namespace Jiaowu.Api.Infrastructure.Persistence;

public static class RetriableTransactionExtensions
{
    public static Task<TResult> ExecuteInRetriableTransactionAsync<TResult>(
        this DbContext db,
        Func<IDbContextTransaction, Task<TResult>> operation,
        CancellationToken cancellationToken,
        IsolationLevel? isolationLevel = null)
    {
        var strategy = db.Database.CreateExecutionStrategy();
        return strategy.ExecuteAsync(async () =>
        {
            await using var transaction = isolationLevel.HasValue
                ? await db.Database.BeginTransactionAsync(
                    isolationLevel.Value,
                    cancellationToken)
                : await db.Database.BeginTransactionAsync(cancellationToken);
            return await operation(transaction);
        });
    }

    public static Task ExecuteInRetriableTransactionAsync(
        this DbContext db,
        Func<IDbContextTransaction, Task> operation,
        CancellationToken cancellationToken,
        IsolationLevel? isolationLevel = null)
    {
        var strategy = db.Database.CreateExecutionStrategy();
        return strategy.ExecuteAsync(async () =>
        {
            await using var transaction = isolationLevel.HasValue
                ? await db.Database.BeginTransactionAsync(
                    isolationLevel.Value,
                    cancellationToken)
                : await db.Database.BeginTransactionAsync(cancellationToken);
            await operation(transaction);
        });
    }
}