package v1

import (
	"context"
	"encoding/json"
	"sort"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/rs/zerolog"

	"github.com/hatchet-dev/hatchet/pkg/integrations/metrics/prometheus"
	v1 "github.com/hatchet-dev/hatchet/pkg/repository"
	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
	"github.com/hatchet-dev/hatchet/pkg/telemetry"
)

type Queuer struct {
	repo       v1.QueueRepository
	optimistic v1.OptimisticSchedulingRepository
	tenantId   uuid.UUID
	queueName  string

	l *zerolog.Logger

	s *Scheduler

	lastReplenished *time.Time

	limit int

	resultsCh chan<- *QueueResults

	notifyQueueCh chan map[string]string

	queueMu mutex

	cleanup func()

	isCleanedUp bool

	unackedMu rwMutex
	unacked   map[int64]struct{}

	unassigned   map[int64]*sqlcv1.V1QueueItem
	unassignedMu mutex

	hasRateLimits bool
}

func newQueuer(conf *sharedConfig, tenantId uuid.UUID, queueName string, s *Scheduler, resultsCh chan<- *QueueResults) *Queuer {
	defaultLimit := 100

	if conf.singleQueueLimit > 0 {
		defaultLimit = conf.singleQueueLimit
	}

	queueRepo := conf.repo.QueueFactory().NewQueue(tenantId, queueName)

	notifyQueueCh := make(chan map[string]string, 1)

	queueLogger := conf.l.With().
		Str("tenant_id", tenantId.String()).
		Str("queue_name", queueName).
		Logger()

	q := &Queuer{
		repo:          queueRepo,
		optimistic:    conf.repo.Optimistic(),
		tenantId:      tenantId,
		queueName:     queueName,
		l:             &queueLogger,
		s:             s,
		limit:         defaultLimit,
		resultsCh:     resultsCh,
		notifyQueueCh: notifyQueueCh,
		queueMu:       newMu(&queueLogger),
		unackedMu:     newRWMu(&queueLogger),
		unacked:       make(map[int64]struct{}),
		unassigned:    make(map[int64]*sqlcv1.V1QueueItem),
		unassignedMu:  newMu(&queueLogger),
	}

	ctx, cancel := context.WithCancel(context.Background())
	cleanupMu := sync.Mutex{}
	q.cleanup = func() {
		cleanupMu.Lock()
		defer cleanupMu.Unlock()

		if q.isCleanedUp {
			return
		}

		q.isCleanedUp = true
		cancel()

		queueRepo.Cleanup()
	}

	go q.loopQueue(ctx)

	return q
}

func (q *Queuer) Cleanup() {
	q.cleanup()
}

func (q *Queuer) queue(ctx context.Context) {
	if ok := q.queueMu.TryLock(); !ok {
		return
	}

	go func() {
		defer q.queueMu.Unlock()

		telemetryCtx, span := telemetry.NewSpan(ctx, "notify-queue")
		defer span.End()

		telemetry.WithAttributes(span,
			telemetry.AttributeKV{Key: "tenant.id", Value: q.tenantId.String()},
			telemetry.AttributeKV{Key: "queue.name", Value: q.queueName},
		)

		q.notifyQueueCh <- telemetry.GetCarrier(telemetryCtx)
	}()
}

func (q *Queuer) loopQueue(ctx context.Context) {
	ticker := time.NewTicker(1 * time.Second)

	q.l.Debug().Ctx(ctx).Int("limit", q.limit).Msg("starting queue loop")

	for {
		var carrier map[string]string

		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		case carrier = <-q.notifyQueueCh:
		}

		q.l.Debug().Ctx(ctx).Msg("queue loop tick")

		prometheus.QueueInvocations.Inc()
		prometheus.TenantQueueInvocations.WithLabelValues(q.tenantId.String()).Inc()

		ctx, span := telemetry.NewSpanWithCarrier(ctx, "queue", carrier)

		telemetry.WithAttributes(span,
			telemetry.AttributeKV{Key: "queue.name", Value: q.queueName},
			telemetry.AttributeKV{Key: "tenant.id", Value: q.tenantId.String()},
		)

		start := time.Now()
		checkpoint := start
		var err error

		if q.hasRateLimits {
			_, err := q.repo.RequeueRateLimitedItems(ctx, q.tenantId, q.queueName)

			if err != nil {
				q.l.Error().Ctx(ctx).Err(err).Msg("error requeuing rate limited items")
			}
		}

		qis, err := q.refillQueue(ctx)

		if err != nil {
			span.RecordError(err)
			span.End()
			q.l.Error().Ctx(ctx).Err(err).Msg("error refilling queue")
			continue
		}

		q.l.Debug().Ctx(ctx).Int("refilled_items", len(qis)).Msg("refilled queue")

		// NOTE: we don't terminate early out of this loop because calling `tryAssign` is necessary
		// for calling the scheduling extensions.

		refillTime := time.Since(checkpoint)
		checkpoint = time.Now()

		rls, err := q.repo.GetTaskRateLimits(ctx, nil, qis)

		if err != nil {
			span.RecordError(err)
			span.End()

			q.l.Error().Ctx(ctx).Err(err).Msg("error getting rate limits")

			q.unackedToUnassigned(qis)
			continue
		}

		if len(rls) > 0 {
			q.hasRateLimits = true
		}

		rateLimitTime := time.Since(checkpoint)
		checkpoint = time.Now()

		stepIds := make([]uuid.UUID, 0, len(qis))
		taskIdToDesiredLabelsFromTrigger := make(map[int64][]*sqlcv1.GetDesiredLabelsRow)

		for _, qi := range qis {
			stepIds = append(stepIds, qi.StepID)

			if len(qi.DesiredWorkerLabel) > 0 {
				var desiredLabels []*sqlcv1.GetDesiredLabelsRow

				err = json.Unmarshal(qi.DesiredWorkerLabel, &desiredLabels)

				if err != nil {
					q.l.Error().Ctx(ctx).Err(err).Int64("queue_item_id", qi.ID).Msg("error unmarshalling desired worker labels for queue item")
				}

				taskIdToDesiredLabelsFromTrigger[qi.TaskID] = desiredLabels
			}
		}

		labels, err := q.repo.GetDesiredLabels(ctx, nil, stepIds)

		if err != nil {
			span.RecordError(err)
			span.End()
			q.l.Error().Ctx(ctx).Err(err).Msg("error getting desired labels")

			q.unackedToUnassigned(qis)
			continue
		}

		desiredLabelsTime := time.Since(checkpoint)
		checkpoint = time.Now()

		stepRequests, err := q.repo.GetStepSlotRequests(ctx, nil, stepIds)

		if err != nil {
			span.RecordError(err)
			span.End()
			q.l.Error().Ctx(ctx).Err(err).Msg("error getting step slot requests")

			q.unackedToUnassigned(qis)
			continue
		}

		getSlotRequestsTime := time.Since(checkpoint)
		checkpoint = time.Now()

		assignCh := q.s.tryAssign(ctx, qis, labels, stepRequests, rls, taskIdToDesiredLabelsFromTrigger)
		count := 0

		countMu := sync.Mutex{}
		wg := sync.WaitGroup{}

		startingQiLength := len(qis)
		processedQiLength := 0

		for r := range assignCh {
			wg.Add(1)

			// asynchronously flush to database
			go func(ar *assignResults) {
				defer wg.Done()

				startFlush := time.Now()

				numFlushed := q.flushToDatabase(ctx, ar)

				countMu.Lock()
				count += numFlushed
				processedQiLength += len(ar.assigned) + len(ar.unassigned) + len(ar.schedulingTimedOut) + len(ar.rateLimited) + len(ar.rateLimitedToMove)
				countMu.Unlock()

				if sinceStart := time.Since(startFlush); sinceStart > 100*time.Millisecond {
					q.l.Warn().Ctx(ctx).Msgf("flushing items to database took longer than 100ms (%d items in %s)", numFlushed, time.Since(startFlush))
				}

				if numFlushed > 0 {
					now := time.Now()

					for _, assignedItem := range ar.assigned {
						prometheus.AssignedTasks.Inc()
						prometheus.TenantAssignedTasks.WithLabelValues(q.tenantId.String()).Inc()

						qi := assignedItem.QueueItem

						// we only track queue times if this is the first retry, because the TaskInsertedAt will
						// match the time the first queue item was created. we should eventually write an InsertedAt
						// column to the queue item table to track this better.
						if qi.RetryCount != 0 || !qi.TaskInsertedAt.Valid {
							continue
						}

						prometheus.QueuedToAssigned.Inc()
						prometheus.TenantQueuedToAssigned.WithLabelValues(q.tenantId.String()).Inc()

						timeInQueueSeconds := now.Sub(qi.TaskInsertedAt.Time).Seconds()

						if timeInQueueSeconds > 0 {
							prometheus.QueuedToAssignedTimeBuckets.Observe(timeInQueueSeconds)
							prometheus.TenantQueuedToAssignedTimeBuckets.WithLabelValues(q.tenantId.String()).Observe(timeInQueueSeconds)
						}
					}
				}

				for range ar.schedulingTimedOut {
					prometheus.SchedulingTimedOut.Inc()
					prometheus.TenantSchedulingTimedOut.WithLabelValues(q.tenantId.String()).Inc()
				}

				for range ar.rateLimited {
					prometheus.RateLimited.Inc()
					prometheus.TenantRateLimited.WithLabelValues(q.tenantId.String()).Inc()
				}
			}(r)
		}

		assignTime := time.Since(checkpoint)
		elapsed := time.Since(start)

		if elapsed > 100*time.Millisecond {
			q.l.Warn().Ctx(ctx).Dur(
				"refill_time", refillTime,
			).Dur(
				"rate_limit_time", rateLimitTime,
			).Dur(
				"desired_labels_time", desiredLabelsTime,
			).Dur(
				"get_slot_requests_time", getSlotRequestsTime,
			).Dur(
				"assign_time", assignTime,
			).Dur("elapsed", elapsed).Int("item_count", len(qis)).Msg("queue processing took longer than 100ms")
		}

		// if we processed all queue items, queue again
		prevQis := qis

		go func(originalStart time.Time) {
			wg.Wait()
			span.End()

			countMu.Lock()
			if len(prevQis) > 0 && count == len(prevQis) {
				q.queue(context.Background())
			}

			if startingQiLength != processedQiLength {
				q.l.Error().Ctx(ctx).Int("starting", startingQiLength).Int("processed", processedQiLength).Msg("queue items processed mismatch")
			}

			countMu.Unlock()

			if sinceStart := time.Since(originalStart); sinceStart > 100*time.Millisecond {
				q.l.Warn().Ctx(ctx).Dur(
					"duration", sinceStart,
				).Int("item_count", len(prevQis)).Msg("queue took longer than 100ms to process and flush items")
			}
		}(start)
	}
}

func (q *Queuer) refillQueue(ctx context.Context) ([]*sqlcv1.V1QueueItem, error) {
	q.unackedMu.Lock()
	defer q.unackedMu.Unlock()

	q.unassignedMu.Lock()
	defer q.unassignedMu.Unlock()

	curr := make([]*sqlcv1.V1QueueItem, 0, len(q.unassigned))

	for _, qi := range q.unassigned {
		curr = append(curr, qi)
	}

	// determine whether we need to replenish with the following cases:
	// - we last replenished more than 1 second ago
	// - if we are at less than 50% of the limit, we always attempt to replenish
	replenish := false

	if len(curr) < q.limit {
		replenish = true
	} else if q.lastReplenished != nil {
		if time.Since(*q.lastReplenished) > 990*time.Millisecond {
			replenish = true
		}
	}

	if replenish {
		now := time.Now()
		q.lastReplenished = &now
		limit := 2 * q.limit

		var err error
		curr, err = q.repo.ListQueueItems(ctx, limit)

		if err != nil {
			return nil, err
		}
	}

	newCurr := make([]*sqlcv1.V1QueueItem, 0, len(curr))

	for _, qi := range curr {
		if _, ok := q.unacked[qi.ID]; !ok {
			newCurr = append(newCurr, qi)
		}
	}

	// add all newCurr to unacked so we don't assign them again
	for _, qi := range newCurr {
		q.unacked[qi.ID] = struct{}{}
	}

	sort.Slice(newCurr, func(i, j int) bool {
		if newCurr[i].Priority == newCurr[j].Priority {
			return newCurr[i].ID < newCurr[j].ID
		}
		return newCurr[i].Priority > newCurr[j].Priority
	})

	return newCurr, nil
}

type QueueResults struct {
	TenantId uuid.UUID
	Assigned []*v1.AssignedItem

	Unassigned         []*sqlcv1.V1QueueItem
	SchedulingTimedOut []*sqlcv1.V1QueueItem
	RateLimited        []*v1.RateLimitResult
}

func (q *Queuer) ack(r *assignResults) {
	q.unackedMu.Lock()
	defer q.unackedMu.Unlock()

	q.unassignedMu.Lock()
	defer q.unassignedMu.Unlock()

	for _, assignedItem := range r.assigned {
		delete(q.unacked, assignedItem.QueueItem.ID)
		delete(q.unassigned, assignedItem.QueueItem.ID)
	}

	for _, unassignedItem := range r.unassigned {
		delete(q.unacked, unassignedItem.ID)
		q.unassigned[unassignedItem.ID] = unassignedItem
	}

	for _, schedulingTimedOutItem := range r.schedulingTimedOut {
		delete(q.unacked, schedulingTimedOutItem.ID)
		delete(q.unassigned, schedulingTimedOutItem.ID)
	}

	for _, rateLimitedItemToMove := range r.rateLimitedToMove {
		delete(q.unacked, rateLimitedItemToMove.qi.ID)
		delete(q.unassigned, rateLimitedItemToMove.qi.ID)
	}

	for _, rateLimitedItem := range r.rateLimited {
		delete(q.unacked, rateLimitedItem.qi.ID)
		q.unassigned[rateLimitedItem.qi.ID] = rateLimitedItem.qi
	}
}

func (q *Queuer) unackedToUnassigned(items []*sqlcv1.V1QueueItem) {
	q.unackedMu.Lock()
	defer q.unackedMu.Unlock()

	q.unassignedMu.Lock()
	defer q.unassignedMu.Unlock()

	for _, item := range items {
		delete(q.unacked, item.ID)
		q.unassigned[item.ID] = item
	}
}

func (q *Queuer) flushToDatabase(ctx context.Context, r *assignResults) int {
	// no matter what, we always ack the items in the queuer
	defer q.ack(r)

	ctx, span := telemetry.NewSpan(ctx, "flush-to-database")
	defer span.End()

	telemetry.WithAttributes(span,
		telemetry.AttributeKV{Key: "tenant.id", Value: q.tenantId.String()},
		telemetry.AttributeKV{Key: "queue.name", Value: q.queueName},
	)

	begin := time.Now()

	q.l.Debug().Ctx(ctx).Int("assigned", len(r.assigned)).Int("unassigned", len(r.unassigned)).Int("scheduling_timed_out", len(r.schedulingTimedOut)).Msg("flushing to database")

	if len(r.assigned) == 0 && len(r.unassigned) == 0 && len(r.schedulingTimedOut) == 0 && len(r.rateLimited) == 0 && len(r.rateLimitedToMove) == 0 {
		return 0
	}

	opts := &v1.AssignResults{
		Assigned:           make([]*v1.AssignedItem, 0, len(r.assigned)),
		Unassigned:         r.unassigned,
		SchedulingTimedOut: r.schedulingTimedOut,
		RateLimited:        make([]*v1.RateLimitResult, 0, len(r.rateLimited)),
		RateLimitedToMove:  make([]*v1.RateLimitResult, 0, len(r.rateLimitedToMove)),
	}

	stepRunIdsToAcks := make(map[int64]int, len(r.assigned))

	for _, assignedItem := range r.assigned {
		stepRunIdsToAcks[assignedItem.QueueItem.TaskID] = assignedItem.AckId

		opts.Assigned = append(opts.Assigned, &v1.AssignedItem{
			WorkerId:  assignedItem.WorkerId,
			QueueItem: assignedItem.QueueItem,
		})
	}

	for _, rateLimitedItem := range r.rateLimited {
		opts.RateLimited = append(opts.RateLimited, &v1.RateLimitResult{
			V1QueueItem:    rateLimitedItem.qi,
			ExceededKey:    rateLimitedItem.exceededKey,
			ExceededUnits:  rateLimitedItem.exceededUnits,
			ExceededVal:    rateLimitedItem.exceededVal,
			NextRefillAt:   rateLimitedItem.nextRefillAt,
			TaskId:         rateLimitedItem.qi.TaskID,
			TaskInsertedAt: rateLimitedItem.qi.TaskInsertedAt,
			RetryCount:     rateLimitedItem.qi.RetryCount,
		})
	}

	for _, rateLimitedItemToMove := range r.rateLimitedToMove {
		opts.RateLimitedToMove = append(opts.RateLimitedToMove, &v1.RateLimitResult{
			V1QueueItem:    rateLimitedItemToMove.qi,
			ExceededKey:    rateLimitedItemToMove.exceededKey,
			ExceededUnits:  rateLimitedItemToMove.exceededUnits,
			ExceededVal:    rateLimitedItemToMove.exceededVal,
			NextRefillAt:   rateLimitedItemToMove.nextRefillAt,
			TaskId:         rateLimitedItemToMove.qi.TaskID,
			TaskInsertedAt: rateLimitedItemToMove.qi.TaskInsertedAt,
			RetryCount:     rateLimitedItemToMove.qi.RetryCount,
		})
	}

	var succeeded []*v1.AssignedItem
	var failed []*v1.AssignedItem
	var err error

	succeeded, failed, err = q.repo.MarkQueueItemsProcessed(ctx, opts)

	if err != nil {
		q.l.Error().Ctx(ctx).Err(err).Msg("error marking queue items processed")

		nackIds := make([]int, 0, len(r.assigned))

		for _, assignedItem := range r.assigned {
			nackIds = append(nackIds, assignedItem.AckId)
		}

		q.s.nack(nackIds)

		return 0
	}

	writeDuration := time.Since(begin)
	checkpoint := time.Now()

	nackIds := make([]int, 0, len(failed))
	ackIds := make([]int, 0, len(succeeded))

	for _, failedItem := range failed {
		nackId := stepRunIdsToAcks[failedItem.QueueItem.TaskID]
		nackIds = append(nackIds, nackId)
	}

	for _, assignedItem := range succeeded {
		ackId := stepRunIdsToAcks[assignedItem.QueueItem.TaskID]
		ackIds = append(ackIds, ackId)
	}

	q.s.nack(nackIds)

	nackDuration := time.Since(checkpoint)
	checkpoint = time.Now()

	q.s.ack(ackIds)

	ackDuration := time.Since(checkpoint)
	checkpoint = time.Now()

	q.resultsCh <- &QueueResults{
		TenantId:           q.tenantId,
		Assigned:           succeeded,
		SchedulingTimedOut: r.schedulingTimedOut,
		RateLimited:        append(opts.RateLimited, opts.RateLimitedToMove...),
		Unassigned:         r.unassigned,
	}

	chWriteDuration := time.Since(checkpoint)

	q.l.Debug().Ctx(ctx).Int("succeeded", len(succeeded)).Int("failed", len(failed)).Msg("flushed to database")

	if time.Since(begin) > 100*time.Millisecond {
		q.l.Warn().Ctx(ctx).Dur(
			"write_duration", writeDuration,
		).Dur(
			"nack_duration", nackDuration,
		).Dur(
			"ack_duration", ackDuration,
		).Dur(
			"ch_write_duration", chWriteDuration,
		).Msgf("flushing %d items to database took longer than 100ms", len(r.assigned)+len(r.unassigned)+len(r.schedulingTimedOut))
	}

	return len(succeeded) + len(r.schedulingTimedOut)
}

func (q *Queuer) runOptimisticQueue(
	ctx context.Context,
	tx *v1.OptimisticTx,
	qis []*sqlcv1.V1QueueItem,
	localWorkerIds map[uuid.UUID]struct{},
) ([]*v1.AssignedItem, []*QueueResults, error) {
	rls, err := q.repo.GetTaskRateLimits(ctx, tx, qis)

	if err != nil {
		return nil, nil, err
	}

	stepIds := make([]uuid.UUID, 0, len(qis))
	taskIdToDesiredLabelsFromTrigger := make(map[int64][]*sqlcv1.GetDesiredLabelsRow)

	for _, qi := range qis {
		stepIds = append(stepIds, qi.StepID)

		if len(qi.DesiredWorkerLabel) > 0 {
			var desiredLabels []*sqlcv1.GetDesiredLabelsRow

			err = json.Unmarshal(qi.DesiredWorkerLabel, &desiredLabels)

			if err != nil {
				q.l.Error().Ctx(ctx).Err(err).Int64("queue_item_id", qi.ID).Msg("error unmarshalling desired worker labels for queue item")
			}

			taskIdToDesiredLabelsFromTrigger[qi.TaskID] = desiredLabels
		}
	}

	labels, err := q.repo.GetDesiredLabels(ctx, tx, stepIds)

	if err != nil {
		return nil, nil, err
	}

	stepRequests, err := q.repo.GetStepSlotRequests(ctx, tx, stepIds)
	if err != nil {
		return nil, nil, err
	}

	assignCh := q.s.tryAssign(ctx, qis, labels, stepRequests, rls, taskIdToDesiredLabelsFromTrigger)

	var allLocalAssigned []*v1.AssignedItem
	var allQueueResults []*QueueResults

	for r := range assignCh {
		// no parallelization since we're in a transaction
		assigned, queueResults, err := q.flushToDatabaseOptimistic(ctx, r, tx, localWorkerIds)

		if err != nil {
			return nil, nil, err
		}

		allLocalAssigned = append(allLocalAssigned, assigned...)
		allQueueResults = append(allQueueResults, queueResults)
	}

	return allLocalAssigned, allQueueResults, nil
}

func (q *Queuer) flushToDatabaseOptimistic(
	ctx context.Context,
	r *assignResults,
	tx *v1.OptimisticTx,
	localWorkerIds map[uuid.UUID]struct{},
) ([]*v1.AssignedItem, *QueueResults, error) {
	begin := time.Now()

	ctx, span := telemetry.NewSpan(ctx, "Queuer.flushToDatabaseOptimistic")
	defer span.End()

	telemetry.WithAttributes(span,
		telemetry.AttributeKV{Key: "tenant.id", Value: q.tenantId.String()},
		telemetry.AttributeKV{Key: "queue.name", Value: q.queueName},
	)

	q.l.Debug().Ctx(ctx).Int("assigned", len(r.assigned)).Int("unassigned", len(r.unassigned)).Int("scheduling_timed_out", len(r.schedulingTimedOut)).Msg("flushing to database")

	if len(r.assigned) == 0 && len(r.unassigned) == 0 && len(r.schedulingTimedOut) == 0 && len(r.rateLimited) == 0 && len(r.rateLimitedToMove) == 0 {
		return nil, nil, nil
	}

	opts := &v1.AssignResults{
		Assigned:           make([]*v1.AssignedItem, 0, len(r.assigned)),
		Unassigned:         r.unassigned,
		SchedulingTimedOut: r.schedulingTimedOut,
		RateLimited:        make([]*v1.RateLimitResult, 0, len(r.rateLimited)),
		RateLimitedToMove:  make([]*v1.RateLimitResult, 0, len(r.rateLimitedToMove)),
	}

	stepRunIdsToAcks := make(map[int64]int, len(r.assigned))

	for _, assignedItem := range r.assigned {
		stepRunIdsToAcks[assignedItem.QueueItem.TaskID] = assignedItem.AckId

		opts.Assigned = append(opts.Assigned, &v1.AssignedItem{
			WorkerId:  assignedItem.WorkerId,
			QueueItem: assignedItem.QueueItem,
		})
	}

	for _, rateLimitedItem := range r.rateLimited {
		opts.RateLimited = append(opts.RateLimited, &v1.RateLimitResult{
			V1QueueItem:    rateLimitedItem.qi,
			ExceededKey:    rateLimitedItem.exceededKey,
			ExceededUnits:  rateLimitedItem.exceededUnits,
			ExceededVal:    rateLimitedItem.exceededVal,
			NextRefillAt:   rateLimitedItem.nextRefillAt,
			TaskId:         rateLimitedItem.qi.TaskID,
			TaskInsertedAt: rateLimitedItem.qi.TaskInsertedAt,
			RetryCount:     rateLimitedItem.qi.RetryCount,
		})
	}

	for _, rateLimitedItemToMove := range r.rateLimitedToMove {
		opts.RateLimitedToMove = append(opts.RateLimitedToMove, &v1.RateLimitResult{
			V1QueueItem:    rateLimitedItemToMove.qi,
			ExceededKey:    rateLimitedItemToMove.exceededKey,
			ExceededUnits:  rateLimitedItemToMove.exceededUnits,
			ExceededVal:    rateLimitedItemToMove.exceededVal,
			NextRefillAt:   rateLimitedItemToMove.nextRefillAt,
			TaskId:         rateLimitedItemToMove.qi.TaskID,
			TaskInsertedAt: rateLimitedItemToMove.qi.TaskInsertedAt,
			RetryCount:     rateLimitedItemToMove.qi.RetryCount,
		})
	}

	var succeeded []*v1.AssignedItem
	var failed []*v1.AssignedItem
	var err error

	succeeded, failed, err = q.optimistic.MarkQueueItemsProcessed(ctx, tx, q.tenantId, opts)

	if err != nil {
		q.l.Error().Ctx(ctx).Err(err).Msg("error marking queue items processed")

		nackIds := make([]int, 0, len(r.assigned))

		for _, assignedItem := range r.assigned {
			nackIds = append(nackIds, assignedItem.AckId)
		}

		q.s.nack(nackIds)

		return nil, nil, err
	}

	writeDuration := time.Since(begin)
	checkpoint := time.Now()

	nackIds := make([]int, 0, len(failed))
	ackIds := make([]int, 0, len(succeeded))

	for _, failedItem := range failed {
		nackId := stepRunIdsToAcks[failedItem.QueueItem.TaskID]
		nackIds = append(nackIds, nackId)
	}

	for _, assignedItem := range succeeded {
		ackId := stepRunIdsToAcks[assignedItem.QueueItem.TaskID]
		ackIds = append(ackIds, ackId)
	}

	q.s.nack(nackIds)

	nackDuration := time.Since(checkpoint)
	checkpoint = time.Now()

	q.s.ack(ackIds)

	ackDuration := time.Since(checkpoint)
	checkpoint = time.Now()

	succeededLocal := make([]*v1.AssignedItem, 0, len(succeeded))

	for _, assignedItem := range succeeded {
		workerId := assignedItem.WorkerId

		if _, ok := localWorkerIds[workerId]; ok {
			assignedItem.IsAssignedLocally = true
			succeededLocal = append(succeededLocal, assignedItem)
		}
	}

	chWriteDuration := time.Since(checkpoint)

	q.l.Debug().Ctx(ctx).Int("succeeded", len(succeeded)).Int("failed", len(failed)).Msg("flushed to database")

	if time.Since(begin) > 100*time.Millisecond {
		q.l.Warn().Ctx(ctx).Dur(
			"write_duration", writeDuration,
		).Dur(
			"nack_duration", nackDuration,
		).Dur(
			"ack_duration", ackDuration,
		).Dur(
			"ch_write_duration", chWriteDuration,
		).Msgf("flushing %d items to database took longer than 100ms", len(r.assigned)+len(r.unassigned)+len(r.schedulingTimedOut))
	}

	return succeededLocal, &QueueResults{
		TenantId:           q.tenantId,
		Assigned:           succeeded,
		SchedulingTimedOut: r.schedulingTimedOut,
		RateLimited:        append(opts.RateLimited, opts.RateLimitedToMove...),
		Unassigned:         r.unassigned,
	}, nil
}
