package store

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

func TestDistributedTaskAdmissionQueueAcrossStores(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run task admission PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)

	first, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect first store: %v", err)
	}
	defer first.Close()
	second, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect second store: %v", err)
	}
	defer second.Close()
	var databaseName string
	if err := first.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
	platform, err := first.CreatePlatform(ctx, CreatePlatformInput{
		Provider:    "queue-test",
		PlatformKey: "queue-test-" + suffix,
		Name:        "Queue Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	var platformModelID string
	if err := first.pool.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, model_name, provider_model_name, model_alias, model_type,
  display_name, pricing_mode, enabled
)
VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID); err != nil {
		t.Fatalf("create platform model: %v", err)
	}
	taskIDs := make([]string, 0, 4)
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		for _, taskID := range taskIDs {
			_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
		}
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = $1::uuid`, platformModelID)
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id LIKE $1`, "queue-test-"+suffix+"%")
	})

	createTask := func(async bool) GatewayTask {
		t.Helper()
		task, createErr := first.CreateTask(ctx, CreateTaskInput{
			Kind:    "images.generate",
			Model:   "queue-model-" + suffix,
			RunMode: "production",
			Async:   async,
			Request: map[string]any{"prompt": "queue admission test"},
		}, &auth.User{ID: "queue-test-user-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create task: %v", createErr)
		}
		taskIDs = append(taskIDs, task.ID)
		return task
	}
	scope := AdmissionScope{
		ScopeType:       "platform_model",
		ScopeKey:        platformModelID,
		ScopeName:       "queue-model",
		ConcurrentLimit: 1,
		Amount:          1,
		LeaseTTLSeconds: 137,
		QueueLimit:      2,
		MaxWaitSeconds:  600,
	}
	inputFor := func(task GatewayTask, priority int, waiterID string) TaskAdmissionInput {
		mode := "sync"
		if task.AsyncMode {
			mode = "async"
		}
		return TaskAdmissionInput{
			TaskID:          task.ID,
			PlatformID:      platform.ID,
			PlatformModelID: platformModelID,
			QueueKey:        "queue-test:" + suffix,
			Mode:            mode,
			Priority:        priority,
			WaiterID:        waiterID,
			Scopes:          []AdmissionScope{scope},
		}
	}

	running := createTask(false)
	result, err := first.TryTaskAdmission(ctx, inputFor(running, 100, "waiter-running"))
	if err != nil || !result.Admitted || len(result.Leases) != 1 {
		t.Fatalf("first task admission = %+v, err=%v", result, err)
	}
	reloadedLeases, err := second.ActiveTaskAdmissionLeases(ctx, running.ID)
	if err != nil || len(reloadedLeases) != 1 || reloadedLeases[0].TTL != 137*time.Second {
		t.Fatalf("reloaded admission leases = %+v, err=%v, want one 137s lease", reloadedLeases, err)
	}
	lowerPriority := createTask(false)
	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
	if err != nil || result.Admitted {
		t.Fatalf("second task should wait: result=%+v err=%v", result, err)
	}
	higherPriorityAsync := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
	if err != nil || result.Admitted {
		t.Fatalf("third task should wait: result=%+v err=%v", result, err)
	}
	rejected := createTask(true)
	_, err = second.TryTaskAdmission(ctx, inputFor(rejected, 30, ""))
	var limitErr *RateLimitExceededError
	if !errors.As(err, &limitErr) || limitErr.Metric != "queue_size" || limitErr.Reason != "queue_full" ||
		limitErr.QueueDepth != 2 || limitErr.QueueLimit != 2 {
		t.Fatalf("fourth task error = %#v, want queue_full depth 2/2", err)
	}

	var queueCounters int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_rate_limit_counters
WHERE metric = 'queue_size'
  AND scope_key = $1`, platformModelID).Scan(&queueCounters); err != nil {
		t.Fatalf("count queue counters: %v", err)
	}
	if queueCounters != 0 {
		t.Fatalf("queue_size wrote %d fixed-window counters", queueCounters)
	}
	var attempts int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_task_attempts
WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil {
		t.Fatalf("count attempts: %v", err)
	}
	if attempts != 0 {
		t.Fatalf("waiting admissions created %d task attempts", attempts)
	}

	if err := first.DeleteTaskAdmission(ctx, running.ID); err != nil {
		t.Fatalf("release running admission: %v", err)
	}
	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
	if err != nil {
		t.Fatalf("try lower-priority task after release: %v", err)
	}
	if result.Admitted {
		t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter")
	}
	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err)
	}
	for _, taskID := range taskIDs {
		_ = first.DeleteTaskAdmission(ctx, taskID)
	}

	concurrentTasks := []GatewayTask{createTask(false), createTask(true), createTask(false), createTask(true)}
	type concurrentResult struct {
		result TaskAdmissionResult
		err    error
	}
	start := make(chan struct{})
	results := make(chan concurrentResult, len(concurrentTasks))
	var waitGroup sync.WaitGroup
	for index, task := range concurrentTasks {
		waitGroup.Add(1)
		go func(index int, task GatewayTask) {
			defer waitGroup.Done()
			<-start
			target := first
			if index%2 == 1 {
				target = second
			}
			waiterID := ""
			if !task.AsyncMode {
				waiterID = "concurrent-waiter-" + task.ID
			}
			admission, admissionErr := target.TryTaskAdmission(ctx, inputFor(task, 100, waiterID))
			results <- concurrentResult{result: admission, err: admissionErr}
		}(index, task)
	}
	close(start)
	waitGroup.Wait()
	close(results)
	admittedCount := 0
	waitingCount := 0
	queueFullCount := 0
	for item := range results {
		if item.err == nil && item.result.Admitted {
			admittedCount++
			continue
		}
		if item.err == nil {
			waitingCount++
			continue
		}
		if errors.As(item.err, &limitErr) && limitErr.Metric == "queue_size" {
			queueFullCount++
			continue
		}
		t.Fatalf("unexpected concurrent admission result: %+v err=%v", item.result, item.err)
	}
	if admittedCount != 1 || waitingCount != 2 || queueFullCount != 1 {
		t.Fatalf("concurrent outcomes admitted=%d waiting=%d queue_full=%d, want 1/2/1", admittedCount, waitingCount, queueFullCount)
	}
	var activeLeases, waitingAdmissions int
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
  AND scope_key = $1
  AND released_at IS NULL
  AND expires_at > now()`, platformModelID).Scan(&activeLeases); err != nil {
		t.Fatalf("count active leases: %v", err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND status = 'waiting'`, platformModelID).Scan(&waitingAdmissions); err != nil {
		t.Fatalf("count waiting admissions: %v", err)
	}
	if activeLeases != 1 || waitingAdmissions != 2 {
		t.Fatalf("database bounds active=%d waiting=%d, want 1/2", activeLeases, waitingAdmissions)
	}

	var deadlineTaskID string
	if err := first.pool.QueryRow(ctx, `
SELECT task_id::text
FROM gateway_task_admissions
WHERE platform_model_id = $1::uuid
  AND status = 'waiting'
ORDER BY task_id
LIMIT 1`, platformModelID).Scan(&deadlineTaskID); err != nil {
		t.Fatalf("select deadline waiter: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET enqueued_at = now() - interval '2 seconds',
    wait_deadline_at = now() - interval '1 second'
WHERE task_id = $1::uuid`, deadlineTaskID); err != nil {
		t.Fatalf("expire queue deadline: %v", err)
	}
	reaped, err := second.ReapExpiredTaskAdmissions(ctx, 10)
	if err != nil || reaped.ExpiredDeadlines < 1 {
		t.Fatalf("deadline reap = %+v, err=%v", reaped, err)
	}
	var taskStatus, taskErrorCode string
	if err := first.pool.QueryRow(ctx, `
SELECT status, COALESCE(error_code, '')
FROM gateway_tasks
WHERE id = $1::uuid`, deadlineTaskID).Scan(&taskStatus, &taskErrorCode); err != nil {
		t.Fatalf("read deadline task: %v", err)
	}
	if taskStatus != "failed" || taskErrorCode != "queue_timeout" {
		t.Fatalf("deadline task status=%s code=%s, want failed/queue_timeout", taskStatus, taskErrorCode)
	}

	for _, taskID := range taskIDs {
		_ = first.DeleteTaskAdmission(ctx, taskID)
	}
	leaseHolder := createTask(false)
	if result, err := first.TryTaskAdmission(ctx, inputFor(leaseHolder, 100, "lease-holder")); err != nil || !result.Admitted {
		t.Fatalf("admit lease holder: result=%+v err=%v", result, err)
	}
	disconnectedWaiter := createTask(false)
	if result, err := second.TryTaskAdmission(ctx, inputFor(disconnectedWaiter, 100, "disconnected-waiter")); err != nil || result.Admitted {
		t.Fatalf("queue disconnected waiter: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_task_admissions
SET waiter_lease_expires_at = now() - interval '1 second'
WHERE task_id = $1::uuid`, disconnectedWaiter.ID); err != nil {
		t.Fatalf("expire waiter lease: %v", err)
	}
	reaped, err = second.ReapExpiredTaskAdmissions(ctx, 10)
	if err != nil || reaped.ExpiredWaiters < 1 {
		t.Fatalf("waiter reap = %+v, err=%v", reaped, err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT status, COALESCE(error_code, '')
FROM gateway_tasks
WHERE id = $1::uuid`, disconnectedWaiter.ID).Scan(&taskStatus, &taskErrorCode); err != nil {
		t.Fatalf("read disconnected task: %v", err)
	}
	if taskStatus != "cancelled" || taskErrorCode != "client_disconnected" {
		t.Fatalf("disconnected task status=%s code=%s, want cancelled/client_disconnected", taskStatus, taskErrorCode)
	}
	if err := first.DeleteTaskAdmission(ctx, leaseHolder.ID); err != nil {
		t.Fatalf("release lease holder admission: %v", err)
	}

	preSubmission := createTask(false)
	claimed, err := first.ClaimTaskPreparation(ctx, preSubmission.ID, uuid.NewString(), time.Minute)
	if err != nil {
		t.Fatalf("claim pre-submission task: %v", err)
	}
	result, err = first.TryTaskAdmission(ctx, inputFor(claimed, 100, "pre-submission"))
	if err != nil || !result.Admitted {
		t.Fatalf("admit pre-submission task: result=%+v err=%v", result, err)
	}
	var attemptID string
	if err := first.pool.QueryRow(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
  upstream_submission_status
)
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'not_submitted')
RETURNING id::text`,
		preSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
	).Scan(&attemptID); err != nil {
		t.Fatalf("create pre-submission attempt: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_param_preprocessing_logs (
  task_id, attempt_id, attempt_no, model_type
)
VALUES ($1::uuid, $2::uuid, 1, 'image_generate')`, preSubmission.ID, attemptID); err != nil {
		t.Fatalf("create pre-submission log: %v", err)
	}
	cancelled, changed, err := first.CancelTaskBeforeUpstreamSubmission(
		ctx,
		preSubmission.ID,
		claimed.ExecutionToken,
		"client disconnected before upstream submission",
	)
	if err != nil || !changed || cancelled.Status != "cancelled" || cancelled.ErrorCode != "client_disconnected" {
		t.Fatalf("cancel pre-submission task = %+v changed=%v err=%v", cancelled, changed, err)
	}
	var remainingAdmissions, remainingLeases, remainingAttempts, remainingPreprocessing int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL),
  (SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid)`,
		preSubmission.ID,
	).Scan(&remainingAdmissions, &remainingLeases, &remainingAttempts, &remainingPreprocessing); err != nil {
		t.Fatalf("read pre-submission cleanup: %v", err)
	}
	if remainingAdmissions != 0 || remainingLeases != 0 || remainingAttempts != 0 || remainingPreprocessing != 0 {
		t.Fatalf(
			"pre-submission cleanup admissions=%d leases=%d attempts=%d preprocessing=%d",
			remainingAdmissions,
			remainingLeases,
			remainingAttempts,
			remainingPreprocessing,
		)
	}

	postSubmission := createTask(false)
	claimed, err = first.ClaimTaskPreparation(ctx, postSubmission.ID, uuid.NewString(), time.Minute)
	if err != nil {
		t.Fatalf("claim post-submission task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
  upstream_submission_status
)
VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'submitting')`,
		postSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
	); err != nil {
		t.Fatalf("create submitted attempt: %v", err)
	}
	_, changed, err = first.CancelTaskBeforeUpstreamSubmission(
		ctx,
		postSubmission.ID,
		claimed.ExecutionToken,
		"client disconnected after upstream submission",
	)
	if err != nil || changed {
		t.Fatalf("post-submission cancellation changed=%v err=%v, want unchanged", changed, err)
	}

	queuedAtomicTask := createTask(true)
	queuedHookFailure := errors.New("synthetic queued hook failure")
	_, err = first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(pgx.Tx) error {
		return queuedHookFailure
	})
	if !errors.Is(err, queuedHookFailure) {
		t.Fatalf("queued hook failure = %v, want synthetic failure", err)
	}
	var queuedAdmissions, queuedLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		queuedAtomicTask.ID,
	).Scan(&queuedAdmissions, &queuedLeases); err != nil {
		t.Fatalf("read rolled back queued hook state: %v", err)
	}
	if queuedAdmissions != 0 || queuedLeases != 0 {
		t.Fatalf("failed queued hook left admissions=%d leases=%d", queuedAdmissions, queuedLeases)
	}
	const queuedSyntheticRiverJobID int64 = 987654320
	queuedAdmission, err := first.QueueTaskAdmissionWithHook(ctx, inputFor(queuedAtomicTask, 100, ""), func(tx pgx.Tx) error {
		_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, queuedAtomicTask.ID, queuedSyntheticRiverJobID)
		return hookErr
	})
	if err != nil || queuedAdmission.Status != "waiting" {
		t.Fatalf("atomic queued admission=%+v err=%v, want waiting", queuedAdmission, err)
	}
	var riverJobID int64
	if err := first.pool.QueryRow(ctx, `
SELECT
  river_job_id,
  (SELECT COUNT(*)
   FROM gateway_concurrency_leases
   WHERE task_id = $1::uuid AND released_at IS NULL)
FROM gateway_tasks
WHERE id = $1::uuid`, queuedAtomicTask.ID).Scan(&riverJobID, &queuedLeases); err != nil {
		t.Fatalf("read atomic queued River marker: %v", err)
	}
	if riverJobID != queuedSyntheticRiverJobID || queuedLeases != 0 {
		t.Fatalf("queued River marker=%d leases=%d, want %d/0", riverJobID, queuedLeases, queuedSyntheticRiverJobID)
	}
	waitingAsyncTaskIDs, err := first.ListWaitingAsyncAdmissionTaskIDs(ctx, 1000)
	if err != nil {
		t.Fatalf("list waiting async tasks with missing River job: %v", err)
	}
	waitingAsyncFound := false
	for _, taskID := range waitingAsyncTaskIDs {
		if taskID == queuedAtomicTask.ID {
			waitingAsyncFound = true
			break
		}
	}
	if !waitingAsyncFound {
		t.Fatal("waiting async task with missing River job was not dispatchable")
	}
	recoverableTasks, err := first.ListRecoverableAsyncTasks(ctx, 1000)
	if err != nil {
		t.Fatalf("list recoverable tasks while admission is waiting: %v", err)
	}
	for _, item := range recoverableTasks {
		if item.ID == queuedAtomicTask.ID {
			t.Fatal("generic River recovery claimed a task owned by the async admission dispatcher")
		}
	}
	recoveryGraceTask := createTask(true)
	recoverableTasks, err = first.ListRecoverableAsyncTasks(ctx, 1000)
	if err != nil {
		t.Fatalf("list fresh recoverable tasks: %v", err)
	}
	for _, item := range recoverableTasks {
		if item.ID == recoveryGraceTask.ID {
			t.Fatal("generic River recovery claimed a task still preparing admission")
		}
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET created_at = now() - interval '6 minutes'
WHERE id = $1::uuid`, recoveryGraceTask.ID); err != nil {
		t.Fatalf("age task beyond River recovery grace period: %v", err)
	}
	recoverableTasks, err = first.ListRecoverableAsyncTasks(ctx, 1000)
	if err != nil {
		t.Fatalf("list aged recoverable tasks: %v", err)
	}
	recoveryGraceFound := false
	for _, item := range recoverableTasks {
		if item.ID == recoveryGraceTask.ID {
			recoveryGraceFound = true
			break
		}
	}
	if !recoveryGraceFound {
		t.Fatal("generic River recovery did not claim a task beyond the recovery grace period")
	}

	batchScope := AdmissionScope{
		ScopeType:       "worker_capacity",
		ScopeKey:        "batch-" + suffix,
		ScopeName:       "batch capacity",
		ConcurrentLimit: 3,
		Amount:          1,
		LeaseTTLSeconds: 120,
		QueueLimit:      100,
		MaxWaitSeconds:  600,
	}
	batchTasks := []GatewayTask{createTask(true), createTask(true), createTask(true)}
	batchInputs := make([]TaskAdmissionInput, 0, len(batchTasks))
	batchJobIDs := make(map[string]int64, len(batchTasks))
	for index, task := range batchTasks {
		input := inputFor(task, 100+index, "")
		input.Scopes = []AdmissionScope{batchScope}
		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue batch task %d: %v", index, err)
		}
		batchInputs = append(batchInputs, input)
		batchJobIDs[task.ID] = int64(987654400 + index)
	}
	batchOutcomes, err := first.TryTaskAdmissionBatchWithAdmittedHook(
		ctx,
		batchInputs,
		func(tx pgx.Tx, input TaskAdmissionInput) error {
			_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, input.TaskID, batchJobIDs[input.TaskID])
			return hookErr
		},
	)
	if err != nil || len(batchOutcomes) != len(batchTasks) {
		t.Fatalf("batch admission outcomes=%+v err=%v", batchOutcomes, err)
	}
	for index, outcome := range batchOutcomes {
		if outcome.Err != nil || !outcome.Result.Admitted || !outcome.Result.NewlyAdmitted ||
			len(outcome.Result.Leases) != 1 {
			t.Fatalf("batch outcome %d=%+v", index, outcome)
		}
	}
	var batchAdmissions, batchLeases, batchRiverJobs int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*)
   FROM gateway_task_admissions
   WHERE task_id = ANY($1::uuid[]) AND status = 'admitted'),
  (SELECT count(*)
   FROM gateway_concurrency_leases
   WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL AND expires_at > now()),
  (SELECT count(*)
   FROM gateway_tasks
   WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
		[]string{batchTasks[0].ID, batchTasks[1].ID, batchTasks[2].ID},
	).Scan(&batchAdmissions, &batchLeases, &batchRiverJobs); err != nil {
		t.Fatalf("read committed batch admission: %v", err)
	}
	if batchAdmissions != 3 || batchLeases != 3 || batchRiverJobs != 3 {
		t.Fatalf(
			"batch admissions=%d leases=%d River jobs=%d, want 3/3/3",
			batchAdmissions,
			batchLeases,
			batchRiverJobs,
		)
	}
	for _, task := range batchTasks {
		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
			t.Fatalf("release batch task %s: %v", task.ID, err)
		}
	}

	rollbackScope := batchScope
	rollbackScope.ScopeKey = "batch-rollback-" + suffix
	rollbackScope.ConcurrentLimit = 2
	rollbackTasks := []GatewayTask{createTask(true), createTask(true)}
	rollbackInputs := make([]TaskAdmissionInput, 0, len(rollbackTasks))
	for index, task := range rollbackTasks {
		input := inputFor(task, 200+index, "")
		input.Scopes = []AdmissionScope{rollbackScope}
		if _, err := first.QueueTaskAdmissionWithHook(ctx, input, nil); err != nil {
			t.Fatalf("queue rollback batch task %d: %v", index, err)
		}
		rollbackInputs = append(rollbackInputs, input)
	}
	batchHookFailure := errors.New("synthetic batch hook failure")
	_, err = first.TryTaskAdmissionBatchWithAdmittedHook(
		ctx,
		rollbackInputs,
		func(tx pgx.Tx, input TaskAdmissionInput) error {
			if input.TaskID == rollbackTasks[1].ID {
				return batchHookFailure
			}
			_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = 987654499
WHERE id = $1::uuid`, input.TaskID)
			return hookErr
		},
	)
	if !errors.Is(err, batchHookFailure) {
		t.Fatalf("batch hook failure=%v, want synthetic failure", err)
	}
	var rollbackWaiting, rollbackLeases, rollbackRiverJobs int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*)
   FROM gateway_task_admissions
   WHERE task_id = ANY($1::uuid[]) AND status = 'waiting'),
  (SELECT count(*)
   FROM gateway_concurrency_leases
   WHERE task_id = ANY($1::uuid[]) AND released_at IS NULL),
  (SELECT count(*)
   FROM gateway_tasks
   WHERE id = ANY($1::uuid[]) AND river_job_id IS NOT NULL)`,
		[]string{rollbackTasks[0].ID, rollbackTasks[1].ID},
	).Scan(&rollbackWaiting, &rollbackLeases, &rollbackRiverJobs); err != nil {
		t.Fatalf("read rolled back batch admission: %v", err)
	}
	if rollbackWaiting != 2 || rollbackLeases != 0 || rollbackRiverJobs != 0 {
		t.Fatalf(
			"rolled back batch waiting=%d leases=%d River jobs=%d, want 2/0/0",
			rollbackWaiting,
			rollbackLeases,
			rollbackRiverJobs,
		)
	}
	for _, task := range rollbackTasks {
		if err := first.DeleteTaskAdmission(ctx, task.ID); err != nil {
			t.Fatalf("release rollback batch task %s: %v", task.ID, err)
		}
	}

	result, err = second.TryTaskAdmission(ctx, inputFor(queuedAtomicTask, 100, ""))
	if err != nil || !result.Admitted || len(result.Leases) != 1 {
		t.Fatalf("worker-time queued admission result=%+v err=%v", result, err)
	}
	if err := first.DeleteTaskAdmission(ctx, queuedAtomicTask.ID); err != nil {
		t.Fatalf("release queued atomic task: %v", err)
	}

	atomicTask := createTask(true)
	hookFailure := errors.New("synthetic admitted hook failure")
	_, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error {
		return hookFailure
	})
	if !errors.Is(err, hookFailure) {
		t.Fatalf("admitted hook failure = %v, want synthetic failure", err)
	}
	var atomicAdmissions, atomicLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		atomicTask.ID,
	).Scan(&atomicAdmissions, &atomicLeases); err != nil {
		t.Fatalf("read rolled back admitted hook state: %v", err)
	}
	if atomicAdmissions != 0 || atomicLeases != 0 {
		t.Fatalf("failed admitted hook left admissions=%d leases=%d", atomicAdmissions, atomicLeases)
	}
	const syntheticRiverJobID int64 = 987654321
	result, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(tx pgx.Tx) error {
		_, hookErr := tx.Exec(ctx, `
UPDATE gateway_tasks
SET river_job_id = $2
WHERE id = $1::uuid`, atomicTask.ID, syntheticRiverJobID)
		return hookErr
	})
	if err != nil || !result.Admitted {
		t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT river_job_id
FROM gateway_tasks
WHERE id = $1::uuid`, atomicTask.ID).Scan(&riverJobID); err != nil {
		t.Fatalf("read atomic River job marker: %v", err)
	}
	if riverJobID != syntheticRiverJobID {
		t.Fatalf("atomic River job marker=%d, want %d", riverJobID, syntheticRiverJobID)
	}
	if err := first.DeleteTaskAdmission(ctx, atomicTask.ID); err != nil {
		t.Fatalf("release atomic admitted hook task: %v", err)
	}

	recoveryTask := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(recoveryTask, 100, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("admit interrupted recovery task: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = $2::uuid,
    execution_lease_expires_at = now() - interval '1 second'
WHERE id = $1::uuid`, recoveryTask.ID, uuid.NewString()); err != nil {
		t.Fatalf("expire interrupted recovery task lease: %v", err)
	}
	recovery, err := first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("recover interrupted admitted task: %v", err)
	}
	if recovery.RequeuedAsyncTasks < 1 || recovery.CleanedTaskAdmissions < 1 {
		t.Fatalf("interrupted task recovery = %+v, want requeue and admission cleanup", recovery)
	}
	var recoveredStatus string
	var recoveredExecutionToken *string
	if err := first.pool.QueryRow(ctx, `
SELECT status, execution_token::text
FROM gateway_tasks
WHERE id = $1::uuid`, recoveryTask.ID).Scan(&recoveredStatus, &recoveredExecutionToken); err != nil {
		t.Fatalf("read recovered task state: %v", err)
	}
	if recoveredStatus != "queued" || recoveredExecutionToken != nil {
		t.Fatalf("recovered task status=%s execution_token=%v, want queued without token", recoveredStatus, recoveredExecutionToken)
	}
	var recoveredAdmissions, recoveredLeases int
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		recoveryTask.ID,
	).Scan(&recoveredAdmissions, &recoveredLeases); err != nil {
		t.Fatalf("read recovered task cleanup: %v", err)
	}
	if recoveredAdmissions != 0 || recoveredLeases != 0 {
		t.Fatalf("recovered task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
	}

	lockedRecoveryTask := createTask(true)
	unlockedRecoveryTask := createTask(true)
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = gen_random_uuid(),
    execution_lease_expires_at = now() - interval '1 second'
WHERE id = ANY($1::uuid[])`, []string{lockedRecoveryTask.ID, unlockedRecoveryTask.ID}); err != nil {
		t.Fatalf("prepare locked recovery tasks: %v", err)
	}
	lockTx, err := second.pool.Begin(ctx)
	if err != nil {
		t.Fatalf("begin task row lock: %v", err)
	}
	var lockedTaskID string
	if err := lockTx.QueryRow(ctx, `
SELECT id::text
FROM gateway_tasks
WHERE id = $1::uuid
FOR UPDATE`, lockedRecoveryTask.ID).Scan(&lockedTaskID); err != nil {
		rollbackTransaction(lockTx)
		t.Fatalf("lock interrupted task row: %v", err)
	}
	recoveryCtx, recoveryCancel := context.WithTimeout(ctx, 3*time.Second)
	recovery, err = first.RecoverInterruptedRuntimeState(recoveryCtx)
	recoveryCancel()
	if err != nil {
		rollbackTransaction(lockTx)
		t.Fatalf("recover unlocked task while another row is locked: %v", err)
	}
	if recovery.RequeuedAsyncTasks < 1 {
		rollbackTransaction(lockTx)
		t.Fatalf("SKIP LOCKED recovery = %+v, want at least one requeued task", recovery)
	}
	var lockedStatus, unlockedStatus string
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT status FROM gateway_tasks WHERE id = $1::uuid),
  (SELECT status FROM gateway_tasks WHERE id = $2::uuid)`,
		lockedRecoveryTask.ID,
		unlockedRecoveryTask.ID,
	).Scan(&lockedStatus, &unlockedStatus); err != nil {
		rollbackTransaction(lockTx)
		t.Fatalf("read SKIP LOCKED recovery states: %v", err)
	}
	if lockedStatus != "running" || unlockedStatus != "queued" {
		rollbackTransaction(lockTx)
		t.Fatalf("SKIP LOCKED recovery statuses = %s/%s, want running/queued", lockedStatus, unlockedStatus)
	}
	rollbackTransaction(lockTx)
	recovery, err = first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("recover previously locked task: %v", err)
	}
	if recovery.RequeuedAsyncTasks < 1 {
		t.Fatalf("unlocked recovery = %+v, want requeued task", recovery)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT status
FROM gateway_tasks
WHERE id = $1::uuid`, lockedRecoveryTask.ID).Scan(&lockedStatus); err != nil {
		t.Fatalf("read unlocked recovery status: %v", err)
	}
	if lockedStatus != "queued" {
		t.Fatalf("unlocked recovery status = %s, want queued", lockedStatus)
	}

	preSubmissionTask := createTask(true)
	preSubmissionToken := uuid.NewString()
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = $2::uuid,
    execution_lease_expires_at = now() + interval '5 minutes'
WHERE id = $1::uuid`,
		preSubmissionTask.ID,
		preSubmissionToken,
	); err != nil {
		t.Fatalf("prepare pre-submission lock timeout task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, queue_key, status, upstream_submission_status
)
VALUES ($1::uuid, 1, 'integration-test', 'running', 'not_submitted')`,
		preSubmissionTask.ID,
	); err != nil {
		t.Fatalf("create pre-submission attempt: %v", err)
	}
	requeuedTask, changed, err := first.RequeueTaskBeforeUpstreamSubmission(
		ctx,
		preSubmissionTask.ID,
		preSubmissionToken,
		time.Second,
	)
	if err != nil || !changed || requeuedTask.Status != "queued" {
		t.Fatalf("pre-submission lock timeout requeue task=%+v changed=%v err=%v", requeuedTask, changed, err)
	}
	var preSubmissionAttempts int
	if err := first.pool.QueryRow(ctx, `
SELECT count(*)
FROM gateway_task_attempts
WHERE task_id = $1::uuid`, preSubmissionTask.ID).Scan(&preSubmissionAttempts); err != nil {
		t.Fatalf("count cleaned pre-submission attempts: %v", err)
	}
	if preSubmissionAttempts != 0 {
		t.Fatalf("pre-submission lock timeout left %d attempts, want 0", preSubmissionAttempts)
	}

	ambiguousTask := createTask(true)
	ambiguousToken := uuid.NewString()
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
    execution_token = $2::uuid,
    execution_lease_expires_at = now() + interval '5 minutes'
WHERE id = $1::uuid`,
		ambiguousTask.ID,
		ambiguousToken,
	); err != nil {
		t.Fatalf("prepare ambiguous submission task: %v", err)
	}
	if _, err := first.pool.Exec(ctx, `
INSERT INTO gateway_task_attempts (
  task_id, attempt_no, queue_key, status, upstream_submission_status
)
VALUES ($1::uuid, 1, 'integration-test', 'running', 'submitting')`,
		ambiguousTask.ID,
	); err != nil {
		t.Fatalf("create ambiguous submission attempt: %v", err)
	}
	_, changed, err = first.RequeueTaskBeforeUpstreamSubmission(
		ctx,
		ambiguousTask.ID,
		ambiguousToken,
		time.Second,
	)
	if err != nil {
		t.Fatalf("guard ambiguous submission task requeue: %v", err)
	}
	if changed {
		t.Fatal("ambiguous submission task was requeued")
	}
	var ambiguousStatus string
	if err := first.pool.QueryRow(ctx, `
SELECT status
FROM gateway_tasks
WHERE id = $1::uuid`, ambiguousTask.ID).Scan(&ambiguousStatus); err != nil {
		t.Fatalf("read ambiguous submission task: %v", err)
	}
	if ambiguousStatus != "running" {
		t.Fatalf("ambiguous submission task status = %s, want running", ambiguousStatus)
	}

	terminalResidue := createTask(true)
	result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, ""))
	if err != nil || !result.Admitted {
		t.Fatalf("admit terminal residue task: result=%+v err=%v", result, err)
	}
	if _, err := first.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
    error_code = 'integration_test',
    finished_at = now(),
    updated_at = now()
WHERE id = $1::uuid`, terminalResidue.ID); err != nil {
		t.Fatalf("mark terminal residue task failed: %v", err)
	}
	recovery, err = first.RecoverInterruptedRuntimeState(ctx)
	if err != nil {
		t.Fatalf("clean terminal admission residue: %v", err)
	}
	if recovery.CleanedTaskAdmissions < 1 {
		t.Fatalf("terminal recovery = %+v, want admission cleanup", recovery)
	}
	if err := first.pool.QueryRow(ctx, `
SELECT
  (SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
  (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
		terminalResidue.ID,
	).Scan(&recoveredAdmissions, &recoveredLeases); err != nil {
		t.Fatalf("count terminal admission residue: %v", err)
	}
	if recoveredAdmissions != 0 || recoveredLeases != 0 {
		t.Fatalf("terminal task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases)
	}
}

func TestWaitingAdmissionLeadersDoNotCrossBlockIndependentCandidates(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run task admission PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()
	var databaseName string
	if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
	platform, err := db.CreatePlatform(ctx, CreatePlatformInput{
		Provider:    "leader-test",
		PlatformKey: "leader-test-" + suffix,
		Name:        "Admission Leader Test " + suffix,
		AuthType:    "none",
		Status:      "enabled",
	})
	if err != nil {
		t.Fatalf("create platform: %v", err)
	}
	modelIDs := make([]string, 0, 2)
	for index := 0; index < 2; index++ {
		modelName := fmt.Sprintf("leader-model-%d-%s", index, suffix)
		var modelID string
		if err := db.pool.QueryRow(ctx, `
INSERT INTO platform_models (
  platform_id, model_name, provider_model_name, model_alias, model_type,
  display_name, pricing_mode, enabled
)
VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
RETURNING id::text`, platform.ID, modelName).Scan(&modelID); err != nil {
			t.Fatalf("create platform model %d: %v", index, err)
		}
		modelIDs = append(modelIDs, modelID)
	}
	group, err := db.CreateUserGroup(ctx, UserGroupInput{
		GroupKey: "leader-test-" + suffix,
		Name:     "Admission Leader Test " + suffix,
		Status:   "active",
	})
	if err != nil {
		t.Fatalf("create user group: %v", err)
	}
	taskIDs := make([]string, 0, 4)
	t.Cleanup(func() {
		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cleanupCancel()
		for _, taskID := range taskIDs {
			_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
		}
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_user_groups WHERE id = $1::uuid`, group.ID)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = ANY($1::uuid[])`, modelIDs)
		_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
	})
	createTask := func(modelIndex int) GatewayTask {
		t.Helper()
		task, createErr := db.CreateTask(ctx, CreateTaskInput{
			Kind:    "images.generate",
			Model:   fmt.Sprintf("leader-model-%d-%s", modelIndex, suffix),
			RunMode: "production",
			Request: map[string]any{"prompt": "admission leader test"},
		}, &auth.User{ID: "leader-test-user-" + suffix, Source: "gateway"})
		if createErr != nil {
			t.Fatalf("create task: %v", createErr)
		}
		taskIDs = append(taskIDs, task.ID)
		return task
	}
	tasks := []GatewayTask{createTask(0), createTask(0), createTask(1), createTask(1)}
	insertWaiting := func(task GatewayTask, modelID string, groupID string, enqueuedAt time.Time) {
		t.Helper()
		if _, insertErr := db.pool.Exec(ctx, `
INSERT INTO gateway_task_admissions (
  task_id, platform_id, platform_model_id, user_group_id, queue_key, mode,
  status, priority, enqueued_at, wait_deadline_at, waiter_id,
  waiter_lease_expires_at
)
VALUES (
  $1::uuid, $2::uuid, $3::uuid, NULLIF($4, '')::uuid, $5, 'sync',
  'waiting', 100, $6, now() + interval '10 minutes', $7,
  now() + interval '15 seconds'
)`,
			task.ID,
			platform.ID,
			modelID,
			groupID,
			"leader-test:"+modelID,
			enqueuedAt,
			"waiter-"+task.ID,
		); insertErr != nil {
			t.Fatalf("insert waiting admission: %v", insertErr)
		}
	}
	baseTime := time.Now().Add(-time.Minute)
	insertWaiting(tasks[0], modelIDs[0], "", baseTime)
	insertWaiting(tasks[1], modelIDs[0], "", baseTime.Add(time.Second))
	insertWaiting(tasks[2], modelIDs[1], "", baseTime.Add(2*time.Second))
	insertWaiting(tasks[3], modelIDs[1], "", baseTime.Add(3*time.Second))
	leaders, err := db.ListWaitingTaskAdmissionIDs(ctx, 256)
	if err != nil {
		t.Fatalf("list independent candidate leaders: %v", err)
	}
	leaderSet := make(map[string]bool, len(leaders))
	for _, taskID := range leaders {
		leaderSet[taskID] = true
	}
	if !leaderSet[tasks[0].ID] || !leaderSet[tasks[2].ID] {
		t.Fatalf("independent candidate leaders = %v, want %s and %s", leaders, tasks[0].ID, tasks[2].ID)
	}
	if leaderSet[tasks[1].ID] || leaderSet[tasks[3].ID] {
		t.Fatalf("non-head independent candidate task was selected: %v", leaders)
	}

	if _, err := db.pool.Exec(ctx, `
DELETE FROM gateway_task_admissions
WHERE task_id = ANY($1::uuid[])`, taskIDs); err != nil {
		t.Fatalf("clear independent admissions: %v", err)
	}
	insertWaiting(tasks[0], modelIDs[0], group.ID, baseTime)
	insertWaiting(tasks[2], modelIDs[1], group.ID, baseTime.Add(time.Second))
	leaders, err = db.ListWaitingTaskAdmissionIDs(ctx, 256)
	if err != nil {
		t.Fatalf("list grouped candidate leaders: %v", err)
	}
	leaderSet = make(map[string]bool, len(leaders))
	for _, taskID := range leaders {
		leaderSet[taskID] = true
	}
	if !leaderSet[tasks[0].ID] || leaderSet[tasks[2].ID] {
		t.Fatalf("grouped candidate leaders = %v, want only group head %s", leaders, tasks[0].ID)
	}
}

func TestWorkerCapacityAllocationAndFailover(t *testing.T) {
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run worker registry PostgreSQL integration tests")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatalf("connect store: %v", err)
	}
	defer db.Close()
	var databaseName string
	if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		t.Fatalf("read test database name: %v", err)
	}
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to use non-test database %q", databaseName)
	}

	prefix := "queue-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
	firstID := prefix + "-a"
	secondID := prefix + "-b"
	defer db.pool.Exec(context.Background(), `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{firstID, secondID})

	first, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
		t.Fatalf("first allocation = %+v, err=%v", first, err)
	}
	second, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: secondID, DesiredCapacity: 5})
	if err != nil || second.Allocated != 2 || second.ActiveInstances != 2 {
		t.Fatalf("second allocation = %+v, err=%v", second, err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 3 || first.ActiveInstances != 2 {
		t.Fatalf("rebalanced first allocation = %+v, err=%v", first, err)
	}
	if _, err := db.pool.Exec(ctx, `
UPDATE gateway_worker_instances
SET heartbeat_at = now() - $2::interval
WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).String()); err != nil {
		t.Fatalf("expire second heartbeat: %v", err)
	}
	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
		t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
	}

	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
		InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 3,
	})
	if err != nil || first.Allocated != 3 || first.GlobalAllocated != 3 {
		t.Fatalf("bounded single-worker allocation = %+v, err=%v", first, err)
	}
	second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
		InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 2,
	})
	if err != nil || second.Allocated != 2 || second.GlobalAllocated != 5 || second.ActiveInstances != 2 {
		t.Fatalf("bounded two-worker allocation = %+v, err=%v", second, err)
	}
}