AI-Health/backend/src/Health.Infrastructure/AI/EfAiWriteConfirmationStore.cs

using Health.Application.AI;
using Health.Infrastructure.Data.Records;

namespace Health.Infrastructure.AI;

public sealed class EfAiWriteConfirmationStore(AppDbContext db) : IAiWriteConfirmationStore
{
    private readonly AppDbContext _db = db;

    public async Task<PendingAiWriteCommand> CreateAsync(
        Guid userId,
        string toolName,
        string arguments,
        TimeSpan lifetime,
        CancellationToken ct)
    {
        var now = DateTime.UtcNow;
        if (IsInMemory())
        {
            var expired = await _db.AiWriteCommands.Where(x => x.Status == "Pending" && x.ExpiresAt <= now).ToListAsync(ct);
            foreach (var item in expired) { item.Status = "Expired"; item.UpdatedAt = now; }
        }
        else
        {
            await _db.AiWriteCommands
                .Where(x => x.Status == "Pending" && x.ExpiresAt <= now)
                .ExecuteUpdateAsync(setters => setters
                    .SetProperty(x => x.Status, "Expired")
                    .SetProperty(x => x.UpdatedAt, now), ct);
        }

        var command = new AiWriteCommandRecord
        {
            Id = Guid.NewGuid(),
            UserId = userId,
            ToolName = toolName,
            Arguments = arguments,
            Status = "Pending",
            ExpiresAt = now.Add(lifetime),
            CreatedAt = now,
            UpdatedAt = now,
        };
        await _db.AiWriteCommands.AddAsync(command, ct);
        await _db.SaveChangesAsync(ct);
        return ToCommand(command);
    }

    public async Task<PendingAiWriteCommand?> TakeAsync(Guid commandId, Guid userId, CancellationToken ct)
    {
        var now = DateTime.UtcNow;
        if (IsInMemory())
        {
            var tracked = await _db.AiWriteCommands.FirstOrDefaultAsync(
                x => x.Id == commandId && x.UserId == userId && x.Status == "Pending" && x.ExpiresAt > now, ct);
            if (tracked == null) return null;
            tracked.Status = "Processing";
            tracked.UpdatedAt = now;
            await _db.SaveChangesAsync(ct);
        }
        else
        {
            var updated = await _db.AiWriteCommands
                .Where(x => x.Id == commandId && x.UserId == userId && x.Status == "Pending" && x.ExpiresAt > now)
                .ExecuteUpdateAsync(setters => setters
                    .SetProperty(x => x.Status, "Processing")
                    .SetProperty(x => x.UpdatedAt, now), ct);
            if (updated == 0) return null;
        }

        var record = await _db.AiWriteCommands.AsNoTracking().FirstAsync(x => x.Id == commandId, ct);
        return ToCommand(record);
    }

    public async Task CompleteAsync(PendingAiWriteCommand command, CancellationToken ct)
    {
        if (IsInMemory())
        {
            var tracked = await _db.AiWriteCommands.FirstOrDefaultAsync(
                x => x.Id == command.Id && x.UserId == command.UserId && x.Status == "Processing", ct);
            if (tracked != null) { tracked.Status = "Completed"; tracked.UpdatedAt = DateTime.UtcNow; await _db.SaveChangesAsync(ct); }
            return;
        }
        await _db.AiWriteCommands
            .Where(x => x.Id == command.Id && x.UserId == command.UserId && x.Status == "Processing")
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(x => x.Status, "Completed")
                .SetProperty(x => x.UpdatedAt, DateTime.UtcNow), ct);
    }

    private bool IsInMemory() => _db.Database.ProviderName == "Microsoft.EntityFrameworkCore.InMemory";

    private static PendingAiWriteCommand ToCommand(AiWriteCommandRecord record) => new(
        record.Id,
        record.UserId,
        record.ToolName,
        record.Arguments,
        record.ExpiresAt);
}