jobtrackingapp/JobTrackerApi/Services/UserOperationStore.cs

using JobTrackerApi.Data;
using JobTrackerApi.Models;
using Microsoft.EntityFrameworkCore;

namespace JobTrackerApi.Services;

public sealed record CreateUserOperation(
    string TaskType,
    string IdempotencyKey,
    string EntitlementDecision,
    string PrivacyPolicy,
    string? SubjectType = null,
    string? SubjectId = null,
    int Priority = 0,
    int MaxAttempts = 3,
    DateTime? DeadlineAtUtc = null);

public sealed record UserOperationCreation(UserOperation Operation, bool Created);
public sealed record UserOperationLease(Guid OperationId, string OwnerUserId, string LeaseToken, string TaskType, string PrivacyPolicy, string? SubjectType, string? SubjectId, int AttemptCount, DateTime? DeadlineAtUtc);

public sealed class UserOperationStore(JobTrackerContext db, TimeProvider timeProvider)
{
    private DateTime UtcNow => timeProvider.GetUtcNow().UtcDateTime;

    public Task<UserOperation?> GetAsync(Guid operationId, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        return db.UserOperations.AsNoTracking().FirstOrDefaultAsync(operation => operation.Id == operationId, cancellationToken);
    }

    public Task<List<UserOperation>> ListAsync(int limit, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        if (limit is < 1 or > 100) throw new ArgumentOutOfRangeException(nameof(limit));
        return db.UserOperations.AsNoTracking()
            .OrderByDescending(operation => operation.CreatedAtUtc)
            .Take(limit)
            .ToListAsync(cancellationToken);
    }

    public Task<UserOperation?> FindByIdempotencyAsync(string taskType, string idempotencyKey, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        return db.UserOperations.AsNoTracking().FirstOrDefaultAsync(
            operation => operation.TaskType == taskType && operation.IdempotencyKey == idempotencyKey,
            cancellationToken);
    }

    public async Task<UserOperationCreation> CreateAsync(CreateUserOperation request, CancellationToken cancellationToken)
    {
        var owner = db.CurrentUserId ?? throw new InvalidOperationException("Operation creation requires an authenticated owner scope.");
        Validate(request);
        var existing = await db.UserOperations.FirstOrDefaultAsync(
            operation => operation.TaskType == request.TaskType && operation.IdempotencyKey == request.IdempotencyKey,
            cancellationToken);
        if (existing is not null) return new UserOperationCreation(existing, false);

        var now = UtcNow;
        var operation = new UserOperation
        {
            Id = Guid.NewGuid(),
            OwnerUserId = owner,
            TaskType = request.TaskType,
            IdempotencyKey = request.IdempotencyKey,
            EntitlementDecision = request.EntitlementDecision,
            PrivacyPolicy = request.PrivacyPolicy,
            SubjectType = request.SubjectType,
            SubjectId = request.SubjectId,
            Priority = request.Priority,
            MaxAttempts = request.MaxAttempts,
            CreatedAtUtc = now,
            AvailableAtUtc = now,
            DeadlineAtUtc = request.DeadlineAtUtc,
        };
        db.UserOperations.Add(operation);
        try
        {
            await db.SaveChangesAsync(cancellationToken);
            return new UserOperationCreation(operation, true);
        }
        catch (DbUpdateException)
        {
            db.Entry(operation).State = EntityState.Detached;
            existing = await db.UserOperations.FirstOrDefaultAsync(
                item => item.TaskType == request.TaskType && item.IdempotencyKey == request.IdempotencyKey,
                cancellationToken);
            if (existing is not null) return new UserOperationCreation(existing, false);
            throw;
        }
    }

    public async Task<UserOperationLease?> ClaimNextAsync(
        TimeSpan leaseDuration,
        CancellationToken cancellationToken,
        IReadOnlyCollection<string>? allowedTaskTypes = null)
    {
        if (db.CurrentUserId is not null) throw new InvalidOperationException("Worker claims require a neutral background scope.");
        ValidateLeaseDuration(leaseDuration);

        var now = UtcNow;
        await RecoverExpiredLeasesAsync(now, cancellationToken);
        var candidates = db.UserOperations.IgnoreQueryFilters().AsNoTracking()
            .Where(operation =>
                (operation.Status == OperationStatuses.Queued || operation.Status == OperationStatuses.WaitingForRetry) &&
                operation.AvailableAtUtc <= now &&
                operation.CancellationRequestedAtUtc == null &&
                operation.AttemptCount < operation.MaxAttempts &&
                (operation.DeadlineAtUtc == null || operation.DeadlineAtUtc > now));
        if (allowedTaskTypes is { Count: > 0 })
            candidates = candidates.Where(operation => allowedTaskTypes.Contains(operation.TaskType));
        var candidateIds = await candidates
            .OrderByDescending(operation => operation.Priority)
            .ThenBy(operation => operation.CreatedAtUtc)
            .Select(operation => operation.Id)
            .Take(16)
            .ToListAsync(cancellationToken);

        foreach (var candidateId in candidateIds)
        {
            var leaseToken = Guid.NewGuid().ToString("N");
            var affected = await db.UserOperations.IgnoreQueryFilters()
                .Where(operation => operation.Id == candidateId &&
                    (operation.Status == OperationStatuses.Queued || operation.Status == OperationStatuses.WaitingForRetry) &&
                    operation.AvailableAtUtc <= now && operation.CancellationRequestedAtUtc == null &&
                    operation.AttemptCount < operation.MaxAttempts &&
                    (operation.DeadlineAtUtc == null || operation.DeadlineAtUtc > now))
                .ExecuteUpdateAsync(setters => setters
                    .SetProperty(operation => operation.Status, OperationStatuses.Running)
                    .SetProperty(operation => operation.LeaseToken, leaseToken)
                    .SetProperty(operation => operation.LeaseExpiresAtUtc, now.Add(leaseDuration))
                    .SetProperty(operation => operation.LastHeartbeatAtUtc, now)
                    .SetProperty(operation => operation.StartedAtUtc, operation => operation.StartedAtUtc ?? now)
                    .SetProperty(operation => operation.AttemptCount, operation => operation.AttemptCount + 1),
                    cancellationToken);
            if (affected != 1) continue;

            var claimed = await db.UserOperations.IgnoreQueryFilters().AsNoTracking()
                .SingleAsync(operation => operation.Id == candidateId && operation.LeaseToken == leaseToken, cancellationToken);
            return new UserOperationLease(claimed.Id, claimed.OwnerUserId, leaseToken, claimed.TaskType, claimed.PrivacyPolicy, claimed.SubjectType, claimed.SubjectId, claimed.AttemptCount, claimed.DeadlineAtUtc);
        }

        return null;
    }

    public Task<int> HeartbeatAsync(Guid operationId, string leaseToken, TimeSpan leaseDuration, string? progressStage, int? progressPercent, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        ValidateLeaseDuration(leaseDuration);
        if (progressPercent is < 0 or > 100) throw new ArgumentOutOfRangeException(nameof(progressPercent));
        ValidateOptional(progressStage, 64, nameof(progressStage));
        var now = UtcNow;
        return db.UserOperations
            .Where(operation => operation.Id == operationId && operation.Status == OperationStatuses.Running && operation.LeaseToken == leaseToken)
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(operation => operation.LeaseExpiresAtUtc, now.Add(leaseDuration))
                .SetProperty(operation => operation.LastHeartbeatAtUtc, now)
                .SetProperty(operation => operation.ProgressStage, progressStage)
                .SetProperty(operation => operation.ProgressPercent, progressPercent),
                cancellationToken);
    }

    public Task<int> CompleteAsync(Guid operationId, string leaseToken, string? resultReference, CancellationToken cancellationToken)
        => CompleteAsync(operationId, leaseToken, resultReference, null, null, null, cancellationToken);

    public async Task<int> CompleteAsync(
        Guid operationId,
        string leaseToken,
        string? resultReference,
        string? provider,
        string? model,
        string? completionStage,
        CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        ValidateOptional(resultReference, 256, nameof(resultReference));
        ValidateOptional(provider, 128, nameof(provider));
        ValidateOptional(model, 128, nameof(model));
        ValidateOptional(completionStage, 64, nameof(completionStage));
        var now = UtcNow;
        await using var transaction = await BeginTransactionAsync(cancellationToken);
        var operation = await db.UserOperations.AsNoTracking().FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken);
        if (operation is null) return 0;
        var affected = await db.UserOperations
            .Where(operation => operation.Id == operationId && operation.Status == OperationStatuses.Running &&
                operation.LeaseToken == leaseToken && operation.CancellationRequestedAtUtc == null)
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(operation => operation.Status, OperationStatuses.Succeeded)
                .SetProperty(operation => operation.ResultReference, resultReference)
                .SetProperty(operation => operation.Provider, provider)
                .SetProperty(operation => operation.Model, model)
                .SetProperty(operation => operation.ProgressStage, completionStage)
                .SetProperty(operation => operation.CompletedAtUtc, now)
                .SetProperty(operation => operation.ProgressPercent, 100)
                .SetProperty(operation => operation.LeaseToken, (string?)null)
                .SetProperty(operation => operation.LeaseExpiresAtUtc, (DateTime?)null),
                cancellationToken);
        if (affected == 1)
        {
            db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Succeeded, now));
            await db.SaveChangesAsync(cancellationToken);
            if (transaction is not null) await transaction.CommitAsync(cancellationToken);
        }
        return affected;
    }

    public Task<bool> FailAsync(Guid operationId, string leaseToken, bool retryable, string category, string message, TimeSpan retryDelay, CancellationToken cancellationToken)
        => FailAsync(operationId, leaseToken, retryable, category, message, retryDelay, null, null, null, cancellationToken);

    public async Task<bool> FailAsync(
        Guid operationId,
        string leaseToken,
        bool retryable,
        string category,
        string message,
        TimeSpan retryDelay,
        string? provider,
        string? model,
        string? progressStage,
        CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        ValidateRequired(category, 64, nameof(category));
        ValidateRequired(message, 512, nameof(message));
        ValidateOptional(provider, 128, nameof(provider));
        ValidateOptional(model, 128, nameof(model));
        ValidateOptional(progressStage, 64, nameof(progressStage));
        if (retryDelay < TimeSpan.Zero || retryDelay > TimeSpan.FromHours(1)) throw new ArgumentOutOfRangeException(nameof(retryDelay));
        await using var transaction = await BeginTransactionAsync(cancellationToken);
        var operation = await db.UserOperations.AsNoTracking()
            .FirstOrDefaultAsync(item => item.Id == operationId && item.Status == OperationStatuses.Running && item.LeaseToken == leaseToken, cancellationToken);
        if (operation is null) return false;

        var now = UtcNow;
        var canRetry = retryable && operation.AttemptCount < operation.MaxAttempts && (operation.DeadlineAtUtc is null || operation.DeadlineAtUtc > now);
        var retryAt = now.Add(retryDelay);
        var affected = await db.UserOperations
            .Where(item => item.Id == operationId && item.Status == OperationStatuses.Running && item.LeaseToken == leaseToken)
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(item => item.Status, canRetry ? OperationStatuses.WaitingForRetry : OperationStatuses.Failed)
                .SetProperty(item => item.AvailableAtUtc, item => canRetry ? retryAt : item.AvailableAtUtc)
                .SetProperty(item => item.CompletedAtUtc, canRetry ? null : now)
                .SetProperty(item => item.FailureCategory, category)
                .SetProperty(item => item.FailureMessage, message)
                .SetProperty(item => item.Provider, provider)
                .SetProperty(item => item.Model, model)
                .SetProperty(item => item.ProgressStage, progressStage)
                .SetProperty(item => item.LeaseToken, (string?)null)
                .SetProperty(item => item.LeaseExpiresAtUtc, (DateTime?)null),
                cancellationToken);
        if (affected == 1 && !canRetry)
        {
            db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Failed, now));
            await db.SaveChangesAsync(cancellationToken);
        }
        if (affected == 1 && transaction is not null) await transaction.CommitAsync(cancellationToken);
        return affected == 1;
    }

    public async Task<bool> RequestCancellationAsync(Guid operationId, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        var now = UtcNow;
        await using var transaction = await BeginTransactionAsync(cancellationToken);
        var operation = await db.UserOperations.AsNoTracking().FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken);
        if (operation is null || OperationStatuses.IsTerminal(operation.Status)) return false;
        var cancelled = await db.UserOperations
            .Where(item => item.Id == operationId &&
                (item.Status == OperationStatuses.Queued || item.Status == OperationStatuses.WaitingForRetry || item.Status == OperationStatuses.WaitingForExternalFallback))
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(item => item.Status, OperationStatuses.Cancelled)
                .SetProperty(item => item.CompletedAtUtc, now),
                cancellationToken);
        if (cancelled == 1)
        {
            db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Cancelled, now));
            await db.SaveChangesAsync(cancellationToken);
            if (transaction is not null) await transaction.CommitAsync(cancellationToken);
            return true;
        }

        var requested = await db.UserOperations
            .Where(item => item.Id == operationId && item.Status == OperationStatuses.Running && item.CancellationRequestedAtUtc == null)
            .ExecuteUpdateAsync(setters => setters.SetProperty(item => item.CancellationRequestedAtUtc, now), cancellationToken);
        if (requested == 1 && transaction is not null) await transaction.CommitAsync(cancellationToken);
        return requested == 1;
    }

    public async Task<int> AcknowledgeCancellationAsync(Guid operationId, string leaseToken, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        var now = UtcNow;
        await using var transaction = await BeginTransactionAsync(cancellationToken);
        var operation = await db.UserOperations.AsNoTracking().FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken);
        if (operation is null) return 0;
        var affected = await db.UserOperations
            .Where(operation => operation.Id == operationId && operation.Status == OperationStatuses.Running &&
                operation.LeaseToken == leaseToken && operation.CancellationRequestedAtUtc != null)
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(operation => operation.Status, OperationStatuses.Cancelled)
                .SetProperty(operation => operation.CompletedAtUtc, now)
                .SetProperty(operation => operation.LeaseToken, (string?)null)
                .SetProperty(operation => operation.LeaseExpiresAtUtc, (DateTime?)null),
                cancellationToken);
        if (affected == 1)
        {
            db.UserNotifications.Add(CreateTerminalNotification(operation, OperationStatuses.Cancelled, now));
            await db.SaveChangesAsync(cancellationToken);
            if (transaction is not null) await transaction.CommitAsync(cancellationToken);
        }
        return affected;
    }

    public async Task<bool> RetryAsync(Guid operationId, CancellationToken cancellationToken)
    {
        EnsureOwnerScope();
        await using var transaction = await BeginTransactionAsync(cancellationToken);
        var operation = await db.UserOperations.FirstOrDefaultAsync(item => item.Id == operationId, cancellationToken);
        if (operation is null || operation.Status is not (OperationStatuses.Failed or OperationStatuses.Cancelled)) return false;
        operation.Status = OperationStatuses.Queued;
        operation.AttemptCount = 0;
        operation.AvailableAtUtc = UtcNow;
        operation.StartedAtUtc = null;
        operation.CompletedAtUtc = null;
        operation.CancellationRequestedAtUtc = null;
        operation.FailureCategory = null;
        operation.FailureMessage = null;
        operation.ResultReference = null;
        operation.Provider = null;
        operation.Model = null;
        operation.ProgressStage = null;
        operation.ProgressPercent = null;
        var existingNotification = await db.UserNotifications.FirstOrDefaultAsync(item => item.OperationId == operationId, cancellationToken);
        if (existingNotification is not null) db.UserNotifications.Remove(existingNotification);
        await db.SaveChangesAsync(cancellationToken);
        if (transaction is not null) await transaction.CommitAsync(cancellationToken);
        return true;
    }

    private async Task RecoverExpiredLeasesAsync(DateTime now, CancellationToken cancellationToken)
    {
        var cancelled = await db.UserOperations.IgnoreQueryFilters().AsNoTracking()
            .Where(operation => operation.Status == OperationStatuses.Running && operation.LeaseExpiresAtUtc <= now && operation.CancellationRequestedAtUtc != null)
            .ToListAsync(cancellationToken);
        foreach (var operation in cancelled)
            await FinalizeRecoveredAsync(operation, RecoveryTerminal.CancelledLease, OperationStatuses.Cancelled, "cancelled", "The operation was cancelled.", now, cancellationToken);

        var exhausted = await db.UserOperations.IgnoreQueryFilters().AsNoTracking()
            .Where(operation => operation.Status == OperationStatuses.Running && operation.LeaseExpiresAtUtc <= now &&
                operation.CancellationRequestedAtUtc == null &&
                (operation.AttemptCount >= operation.MaxAttempts || operation.DeadlineAtUtc <= now))
            .ToListAsync(cancellationToken);
        foreach (var operation in exhausted)
            await FinalizeRecoveredAsync(operation, RecoveryTerminal.ExpiredLease, OperationStatuses.Failed, "lease_expired", "The operation could not be recovered after its final worker attempt.", now, cancellationToken);

        var deadlineExpired = await db.UserOperations.IgnoreQueryFilters().AsNoTracking()
            .Where(operation =>
                (operation.Status == OperationStatuses.Queued || operation.Status == OperationStatuses.WaitingForRetry || operation.Status == OperationStatuses.WaitingForExternalFallback) &&
                operation.DeadlineAtUtc <= now)
            .ToListAsync(cancellationToken);
        foreach (var operation in deadlineExpired)
            await FinalizeRecoveredAsync(operation, RecoveryTerminal.QueuedDeadline, OperationStatuses.Failed, "deadline_exceeded", "The operation deadline elapsed before work could complete.", now, cancellationToken);

        await db.UserOperations.IgnoreQueryFilters()
            .Where(operation => operation.Status == OperationStatuses.Running && operation.LeaseExpiresAtUtc <= now &&
                operation.CancellationRequestedAtUtc == null && operation.AttemptCount < operation.MaxAttempts &&
                (operation.DeadlineAtUtc == null || operation.DeadlineAtUtc > now))
            .ExecuteUpdateAsync(setters => setters
                .SetProperty(operation => operation.Status, OperationStatuses.WaitingForRetry)
                .SetProperty(operation => operation.AvailableAtUtc, now)
                .SetProperty(operation => operation.FailureCategory, "lease_expired")
                .SetProperty(operation => operation.FailureMessage, "The worker stopped before completing this operation; it will be retried.")
                .SetProperty(operation => operation.LeaseToken, (string?)null)
                .SetProperty(operation => operation.LeaseExpiresAtUtc, (DateTime?)null),
                cancellationToken);
    }

    private async Task FinalizeRecoveredAsync(UserOperation operation, RecoveryTerminal reason, string status, string category, string message, DateTime now, CancellationToken cancellationToken)
    {
        await using var transaction = await BeginTransactionAsync(cancellationToken);
        var query = db.UserOperations.IgnoreQueryFilters().Where(item => item.Id == operation.Id);
        query = reason switch
        {
            RecoveryTerminal.CancelledLease => query.Where(item => item.Status == OperationStatuses.Running && item.LeaseExpiresAtUtc <= now && item.CancellationRequestedAtUtc != null),
            RecoveryTerminal.ExpiredLease => query.Where(item => item.Status == OperationStatuses.Running && item.LeaseExpiresAtUtc <= now && item.CancellationRequestedAtUtc == null && (item.AttemptCount >= item.MaxAttempts || item.DeadlineAtUtc <= now)),
            _ => query.Where(item => (item.Status == OperationStatuses.Queued || item.Status == OperationStatuses.WaitingForRetry || item.Status == OperationStatuses.WaitingForExternalFallback) && item.DeadlineAtUtc <= now),
        };
        var affected = await query.ExecuteUpdateAsync(setters => setters
            .SetProperty(item => item.Status, status)
            .SetProperty(item => item.CompletedAtUtc, now)
            .SetProperty(item => item.FailureCategory, category)
            .SetProperty(item => item.FailureMessage, message)
            .SetProperty(item => item.LeaseToken, (string?)null)
            .SetProperty(item => item.LeaseExpiresAtUtc, (DateTime?)null), cancellationToken);
        if (affected == 1)
        {
            db.UserNotifications.Add(CreateTerminalNotification(operation, status, now));
            await db.SaveChangesAsync(cancellationToken);
            if (transaction is not null) await transaction.CommitAsync(cancellationToken);
        }
    }

    private async Task<Microsoft.EntityFrameworkCore.Storage.IDbContextTransaction?> BeginTransactionAsync(CancellationToken cancellationToken)
    {
        if (!db.Database.IsRelational()) return null;
        return await db.Database.BeginTransactionAsync(cancellationToken);
    }

    private static UserNotification CreateTerminalNotification(UserOperation operation, string status, DateTime now)
    {
        var (kind, title, message) = status switch
        {
            OperationStatuses.Succeeded => ("operation_succeeded", "Operation completed", "Your background operation completed."),
            OperationStatuses.Cancelled => ("operation_cancelled", "Operation cancelled", "Your background operation was cancelled."),
            _ => ("operation_failed", "Operation failed", "A background operation failed. Review it for details."),
        };
        return new UserNotification
        {
            Id = Guid.NewGuid(),
            OwnerUserId = operation.OwnerUserId,
            OperationId = operation.Id,
            Kind = kind,
            Title = title,
            Message = message,
            CreatedAtUtc = now,
        };
    }

    private enum RecoveryTerminal { CancelledLease, ExpiredLease, QueuedDeadline }

    private static void Validate(CreateUserOperation request)
    {
        ValidateRequired(request.TaskType, 64, nameof(request.TaskType));
        ValidateRequired(request.IdempotencyKey, 128, nameof(request.IdempotencyKey));
        ValidateRequired(request.EntitlementDecision, 32, nameof(request.EntitlementDecision));
        ValidateRequired(request.PrivacyPolicy, 32, nameof(request.PrivacyPolicy));
        ValidateOptional(request.SubjectType, 64, nameof(request.SubjectType));
        ValidateOptional(request.SubjectId, 128, nameof(request.SubjectId));
        if (request.MaxAttempts is < 1 or > 10) throw new ArgumentOutOfRangeException(nameof(request.MaxAttempts));
        if (request.DeadlineAtUtc is { Kind: not DateTimeKind.Utc }) throw new ArgumentException("Operation deadlines must be UTC.", nameof(request.DeadlineAtUtc));
    }

    private void EnsureOwnerScope()
    {
        if (db.CurrentUserId is null) throw new InvalidOperationException("Operation mutation requires an explicit owner scope.");
    }

    private static void ValidateLeaseDuration(TimeSpan leaseDuration)
    {
        if (leaseDuration < TimeSpan.FromSeconds(5) || leaseDuration > TimeSpan.FromMinutes(30))
            throw new ArgumentOutOfRangeException(nameof(leaseDuration));
    }

    private static void ValidateRequired(string value, int maxLength, string name)
    {
        if (string.IsNullOrWhiteSpace(value) || value.Length > maxLength) throw new ArgumentException($"{name} is required and limited to {maxLength} characters.", name);
    }

    private static void ValidateOptional(string? value, int maxLength, string name)
    {
        if (value?.Length > maxLength) throw new ArgumentException($"{name} is limited to {maxLength} characters.", name);
    }
}