package datastore

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/semantic"
)

func newSemanticTask(taskID, resourceID string, version int64, availableAt, createdAt time.Time) *semantic.Task {
	return newSemanticTaskOfType(semantic.TaskTypeEmbed, taskID, resourceID, version, availableAt, createdAt)
}

func newSemanticTaskOfType(taskType semantic.TaskType, taskID, resourceID string, version int64, availableAt, createdAt time.Time) *semantic.Task {
	return &semantic.Task{
		TaskID:          taskID,
		TaskType:        taskType,
		ResourceID:      resourceID,
		ResourceVersion: version,
		Status:          semantic.TaskQueued,
		MaxAttempts:     3,
		AvailableAt:     availableAt.UTC(),
		CreatedAt:       createdAt.UTC(),
		UpdatedAt:       createdAt.UTC(),
		PayloadJSON:     []byte(`{"resource_id":"` + resourceID + `"}`),
	}
}

func TestSemanticTaskLifecycle(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	created, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now))
	if err != nil {
		t.Fatal(err)
	}
	if !created {
		t.Fatal("expected task to be created")
	}

	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find queued task")
	}
	if claimed.TaskID != "task-1" || claimed.AttemptCount != 1 {
		t.Fatalf("unexpected claim result: %+v", claimed)
	}
	if claimed.Receipt == "" || claimed.LeasedAt == nil || claimed.LeaseUntil == nil {
		t.Fatalf("claim did not populate lease fields: %+v", claimed)
	}

	if err := s.AckSemanticTask(ctx, claimed.TaskID, claimed.Receipt); err != nil {
		t.Fatal(err)
	}

	task := mustGetSemanticTask(t, s, claimed.TaskID)
	if task.Status != semantic.TaskSucceeded {
		t.Fatalf("status=%q, want %q", task.Status, semantic.TaskSucceeded)
	}
	if task.CompletedAt == nil {
		t.Fatal("expected completed_at after ack")
	}
	if task.Receipt != "" || task.LeasedAt != nil || task.LeaseUntil != nil {
		t.Fatalf("ack should clear lease fields: %+v", task)
	}
}

func TestUpdateSemanticTaskPayloadRequiresLiveLease(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if _, err := s.EnqueueSemanticTask(context.Background(), newSemanticTask("checkpoint-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(context.Background(), now.Add(time.Second), time.Minute)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}
	if err := s.UpdateSemanticTaskPayload(context.Background(), claimed.TaskID, claimed.Receipt, []byte(`{"segment":2}`)); err != nil {
		t.Fatalf("update payload: %v", err)
	}
	got := mustGetSemanticTask(t, s, claimed.TaskID)
	if string(got.PayloadJSON) != `{"segment": 2}` {
		t.Fatalf("payload=%s", got.PayloadJSON)
	}
	if err := s.UpdateSemanticTaskPayload(context.Background(), claimed.TaskID, "stale", nil); !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("stale update err=%v", err)
	}
	if err := s.UpdateSemanticTaskPayload(context.Background(), claimed.TaskID, claimed.Receipt, make([]byte, MaxSemanticTaskPayloadBytes+1)); !errors.Is(err, ErrSemanticTaskPayloadTooLarge) {
		t.Fatalf("oversize err=%v", err)
	}
}

func TestSemanticTaskDuplicateEnqueueReturnsFalse(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Unix(1711600100, 0).UTC()

	if created, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 7, now, now)); err != nil {
		t.Fatal(err)
	} else if !created {
		t.Fatal("first enqueue should create task")
	}

	created, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-2", "file-1", 7, now.Add(time.Second), now.Add(time.Second)))
	if err != nil {
		t.Fatal(err)
	}
	if created {
		t.Fatal("duplicate enqueue should not create a second task")
	}

	if count := countSemanticTasks(t, s); count != 1 {
		t.Fatalf("semantic task count=%d, want 1", count)
	}
}

func TestSemanticTaskAckWrongReceiptFails(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}

	err = s.AckSemanticTask(ctx, claimed.TaskID, "wrong-receipt")
	if !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("ack error=%v, want %v", err, semantic.ErrTaskLeaseMismatch)
	}
}

func TestDiscardSemanticTaskRequiresLiveLeaseAndDeletesRow(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTaskOfType(semantic.TaskTypeVideoExtractVisual, "discard-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeVideoExtractVisual)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}
	if err := s.DiscardSemanticTask(ctx, claimed.TaskID, "stale-receipt"); !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("stale discard err=%v, want lease mismatch", err)
	}
	if err := s.DiscardSemanticTask(ctx, claimed.TaskID, claimed.Receipt); err != nil {
		t.Fatal(err)
	}
	var count int
	if err := s.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM semantic_tasks WHERE task_id = ?`, claimed.TaskID).Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 0 {
		t.Fatalf("task count=%d, want deleted", count)
	}
}

func TestDiscardFreshClaimableSemanticTasksByTypeIsBoundedDueAndUnattemptedOnly(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	tasks := []*semantic.Task{
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "img-processing", "file-processing", 1, now, now),
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "img-due-1", "file-due-1", 1, now, now.Add(time.Millisecond)),
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "img-due-2", "file-due-2", 1, now, now.Add(2*time.Millisecond)),
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "img-future", "file-future", 1, now.Add(time.Hour), now),
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "img-retried", "file-retried", 1, now, now),
		newSemanticTaskOfType(semantic.TaskTypeAudioExtractText, "audio-due", "file-audio", 1, now, now),
	}
	for _, task := range tasks {
		if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := s.DB().ExecContext(ctx, `UPDATE semantic_tasks SET attempt_count = 1 WHERE task_id = 'img-retried'`); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeImgExtractText)
	if err != nil || !found || claimed.TaskID != "img-processing" {
		t.Fatalf("claim found=%v task=%+v err=%v", found, claimed, err)
	}

	deleted, err := s.DiscardFreshClaimableSemanticTasksByType(ctx, semantic.TaskTypeImgExtractText, now.Add(time.Second), 1)
	if err != nil {
		t.Fatal(err)
	}
	if deleted != 1 {
		t.Fatalf("deleted=%d, want 1", deleted)
	}
	if got := countSemanticTasks(t, s); got != 5 {
		t.Fatalf("task count=%d, want 5 after bounded delete", got)
	}
	deleted, err = s.DiscardFreshClaimableSemanticTasksByType(ctx, semantic.TaskTypeImgExtractText, now.Add(time.Second), 100)
	if err != nil {
		t.Fatal(err)
	}
	if deleted != 1 {
		t.Fatalf("second deleted=%d, want remaining due image task", deleted)
	}
	for _, taskID := range []string{"img-processing", "img-future", "img-retried", "audio-due"} {
		var count int
		if err := s.DB().QueryRowContext(ctx, `SELECT COUNT(*) FROM semantic_tasks WHERE task_id = ?`, taskID).Scan(&count); err != nil {
			t.Fatal(err)
		}
		if count != 1 {
			t.Fatalf("task %s count=%d, want preserved", taskID, count)
		}
	}
}

func TestSemanticTaskDeadLetterWithoutRetry(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}

	if err := s.DeadLetterSemanticTask(ctx, claimed.TaskID, claimed.Receipt, "permanent audio extract error"); err != nil {
		t.Fatal(err)
	}

	task := mustGetSemanticTask(t, s, claimed.TaskID)
	if task.Status != semantic.TaskDeadLettered {
		t.Fatalf("status=%q, want %q", task.Status, semantic.TaskDeadLettered)
	}
	if task.AttemptCount != 1 {
		t.Fatalf("attempt_count=%d, want 1", task.AttemptCount)
	}
	if task.LastError != "permanent audio extract error" {
		t.Fatalf("last_error=%q", task.LastError)
	}
	if task.CompletedAt == nil {
		t.Fatal("expected completed_at after dead-letter")
	}
	if task.Receipt != "" || task.LeasedAt != nil || task.LeaseUntil != nil {
		t.Fatalf("dead-letter should clear lease fields: %+v", task)
	}
}

func TestSemanticTaskDeferDoesNotConsumeClaimAttempt(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-defer", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil || !found {
		t.Fatalf("ClaimSemanticTask: found=%v err=%v", found, err)
	}
	if err := s.DeferSemanticTask(ctx, claimed.TaskID, claimed.Receipt, now.Add(2*time.Second), "global_disabled"); err != nil {
		t.Fatal(err)
	}

	task := mustGetSemanticTask(t, s, claimed.TaskID)
	if task.Status != semantic.TaskQueued || task.AttemptCount != 0 {
		t.Fatalf("status/attempt_count = %s/%d, want queued/0", task.Status, task.AttemptCount)
	}
	if task.Receipt != "" || task.LeasedAt != nil || task.LeaseUntil != nil {
		t.Fatalf("defer should clear lease fields: %+v", task)
	}
	if task.LastError != "global_disabled" {
		t.Fatalf("last_error = %q, want global_disabled", task.LastError)
	}
}

func TestSemanticTaskRenewExtendsLease(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}

	renewedUntil, err := s.RenewSemanticTask(ctx, claimed.TaskID, claimed.Receipt, 2*time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !renewedUntil.After(*claimed.LeaseUntil) {
		t.Fatalf("renewed lease_until=%s, want after %s", renewedUntil.UTC(), claimed.LeaseUntil.UTC())
	}

	task := mustGetSemanticTask(t, s, claimed.TaskID)
	if task.Status != semantic.TaskProcessing {
		t.Fatalf("status=%q, want %q", task.Status, semantic.TaskProcessing)
	}
	if task.AttemptCount != 1 {
		t.Fatalf("attempt_count=%d, want 1", task.AttemptCount)
	}
	if !task.AvailableAt.Equal(claimed.AvailableAt) {
		t.Fatalf("available_at=%s, want %s", task.AvailableAt.UTC(), claimed.AvailableAt.UTC())
	}
	if task.LeaseUntil == nil {
		t.Fatal("stored lease_until should not be nil after renew")
	}
	if !task.LeaseUntil.After(*claimed.LeaseUntil) {
		t.Fatalf("stored lease_until=%s, want after %s", task.LeaseUntil.UTC(), claimed.LeaseUntil.UTC())
	}
	if delta := task.LeaseUntil.Sub(renewedUntil); delta < -time.Second || delta > time.Second {
		t.Fatalf("stored lease_until=%s, renewed_until=%s, delta=%s", task.LeaseUntil.UTC(), renewedUntil.UTC(), delta)
	}
}

func TestSemanticTaskRenewWrongReceiptFails(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}

	_, err = s.RenewSemanticTask(ctx, claimed.TaskID, "wrong-receipt", time.Minute)
	if !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("renew error=%v, want %v", err, semantic.ErrTaskLeaseMismatch)
	}
}

func TestSemanticTaskRenewExpiredLeaseFails(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	availableAt := now.Add(-time.Second)

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, availableAt, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, time.Now().UTC(), 20*time.Millisecond)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}

	time.Sleep(40 * time.Millisecond)

	_, err = s.RenewSemanticTask(ctx, claimed.TaskID, claimed.Receipt, time.Minute)
	if !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("renew error=%v, want %v", err, semantic.ErrTaskLeaseMismatch)
	}
}

func TestSemanticTaskRenewAfterAckFails(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}
	if err := s.AckSemanticTask(ctx, claimed.TaskID, claimed.Receipt); err != nil {
		t.Fatal(err)
	}

	_, err = s.RenewSemanticTask(ctx, claimed.TaskID, claimed.Receipt, time.Minute)
	if !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("renew error=%v, want %v", err, semantic.ErrTaskLeaseMismatch)
	}
}

func TestSemanticTaskRenewNotFoundFails(t *testing.T) {
	s := newTestStore(t)
	_, err := s.RenewSemanticTask(context.Background(), "missing-task", "missing-receipt", time.Minute)
	if !errors.Is(err, semantic.ErrTaskNotFound) {
		t.Fatalf("renew error=%v, want %v", err, semantic.ErrTaskNotFound)
	}
}

func TestSemanticTaskClaimOrderByAvailableAtThenCreatedAt(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600300, 0).UTC()

	tasks := []*semantic.Task{
		newSemanticTask("task-late", "file-late", 1, base.Add(10*time.Second), base.Add(10*time.Second)),
		newSemanticTask("task-b", "file-b", 1, base, base.Add(2*time.Second)),
		newSemanticTask("task-a", "file-a", 1, base, base.Add(time.Second)),
	}
	for _, task := range tasks {
		if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
			t.Fatal(err)
		}
	}

	for _, want := range []string{"task-a", "task-b", "task-late"} {
		claimed, found, err := s.ClaimSemanticTask(ctx, base.Add(20*time.Second), time.Minute)
		if err != nil {
			t.Fatal(err)
		}
		if !found {
			t.Fatalf("expected claim to find %s", want)
		}
		if claimed.TaskID != want {
			t.Fatalf("claimed task=%q, want %q", claimed.TaskID, want)
		}
	}
}

func TestClaimSemanticTaskSkipsOlderDisallowedTask(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600350, 0).UTC()

	tasks := []*semantic.Task{
		newSemanticTaskOfType(semantic.TaskTypeEmbed, "task-embed", "file-embed", 1, base, base),
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "task-img", "file-img", 1, base.Add(time.Second), base.Add(time.Second)),
	}
	for _, task := range tasks {
		if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
			t.Fatal(err)
		}
	}

	claimed, found, err := s.ClaimSemanticTask(ctx, base.Add(5*time.Second), time.Minute, semantic.TaskTypeImgExtractText)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected filtered claim to find image task")
	}
	if claimed.TaskID != "task-img" {
		t.Fatalf("claimed task=%q, want %q", claimed.TaskID, "task-img")
	}

	embed := mustGetSemanticTask(t, s, "task-embed")
	if embed.Status != semantic.TaskQueued {
		t.Fatalf("embed task status=%q, want %q", embed.Status, semantic.TaskQueued)
	}
}

func TestClaimSemanticTaskReturnsNotFoundWhenOnlyDisallowedTasksExist(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600360, 0).UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTaskOfType(semantic.TaskTypeEmbed, "task-embed", "file-embed", 1, base, base)); err != nil {
		t.Fatal(err)
	}

	claimed, found, err := s.ClaimSemanticTask(ctx, base.Add(time.Second), time.Minute, semantic.TaskTypeImgExtractText)
	if err != nil {
		t.Fatal(err)
	}
	if found {
		t.Fatalf("unexpected claimed task: %+v", claimed)
	}

	task := mustGetSemanticTask(t, s, "task-embed")
	if task.Status != semantic.TaskQueued {
		t.Fatalf("task status=%q, want %q", task.Status, semantic.TaskQueued)
	}
}

func TestClaimSemanticTaskSupportsMultipleTaskTypes(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600370, 0).UTC()

	tasks := []*semantic.Task{
		newSemanticTaskOfType(semantic.TaskTypeEmbed, "task-embed", "file-embed", 1, base, base),
		newSemanticTaskOfType(semantic.TaskTypeImgExtractText, "task-img", "file-img", 1, base.Add(time.Second), base.Add(time.Second)),
	}
	for _, task := range tasks {
		if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
			t.Fatal(err)
		}
	}

	claimed, found, err := s.ClaimSemanticTask(ctx, base.Add(5*time.Second), time.Minute, semantic.TaskTypeImgExtractText, semantic.TaskTypeEmbed)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected filtered claim to find a task")
	}
	if claimed.TaskID != "task-embed" {
		t.Fatalf("claimed task=%q, want %q", claimed.TaskID, "task-embed")
	}
}

func TestClaimSemanticTaskWithoutFiltersClaimsAnyTask(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600380, 0).UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, base, base)); err != nil {
		t.Fatal(err)
	}

	claimed, found, err := s.ClaimSemanticTask(ctx, base.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected unfiltered claim to find task")
	}
	if claimed.TaskID != "task-1" {
		t.Fatalf("claimed task=%q, want %q", claimed.TaskID, "task-1")
	}
}

func TestSemanticTaskRetryAndRecoverExpired(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, base, base)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, base.Add(time.Second), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}

	retryAt := base.Add(5 * time.Second)
	if err := s.RetrySemanticTask(ctx, claimed.TaskID, claimed.Receipt, retryAt, "temporary failure"); err != nil {
		t.Fatal(err)
	}

	if _, found, err := s.ClaimSemanticTask(ctx, base.Add(4*time.Second), time.Second); err != nil {
		t.Fatal(err)
	} else if found {
		t.Fatal("task should not be claimable before retry_at")
	}

	reclaimed, found, err := s.ClaimSemanticTask(ctx, retryAt.Add(time.Second), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected task to be claimable at retry_at")
	}
	if reclaimed.AttemptCount != 2 {
		t.Fatalf("attempt_count=%d, want 2", reclaimed.AttemptCount)
	}

	recoveredAt := retryAt.Add(3 * time.Second)
	recovered, err := s.RecoverExpiredSemanticTasks(ctx, recoveredAt, 0)
	if err != nil {
		t.Fatal(err)
	}
	if recovered != 1 {
		t.Fatalf("recovered=%d, want 1", recovered)
	}

	recoveredTask, found, err := s.ClaimSemanticTask(ctx, recoveredAt.Add(time.Second), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected recovered task to become claimable again")
	}
	if recoveredTask.AttemptCount != 3 {
		t.Fatalf("attempt_count=%d, want 3", recoveredTask.AttemptCount)
	}
	if err := s.RetrySemanticTask(ctx, recoveredTask.TaskID, recoveredTask.Receipt, retryAt.Add(3*time.Second), "permanent failure"); err != nil {
		t.Fatal(err)
	}

	task := mustGetSemanticTask(t, s, recoveredTask.TaskID)
	if task.Status != semantic.TaskDeadLettered {
		t.Fatalf("status=%q, want %q", task.Status, semantic.TaskDeadLettered)
	}
	if task.CompletedAt == nil {
		t.Fatal("expected completed_at for dead-lettered task")
	}
}

func TestSemanticTaskRecoverAfterReopen(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600500, 0).UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, base, base)); err != nil {
		t.Fatal(err)
	}
	if _, found, err := s.ClaimSemanticTask(ctx, base, time.Second); err != nil {
		t.Fatal(err)
	} else if !found {
		t.Fatal("expected initial claim to succeed")
	}
	if err := s.Close(); err != nil {
		t.Fatal(err)
	}

	reopened, err := Open(testDSN)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = reopened.Close() })

	recovered, err := reopened.RecoverExpiredSemanticTasks(ctx, base.Add(2*time.Second), 0)
	if err != nil {
		t.Fatal(err)
	}
	if recovered != 1 {
		t.Fatalf("recovered=%d, want 1", recovered)
	}

	claimed, found, err := reopened.ClaimSemanticTask(ctx, base.Add(2*time.Second), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected recovered task after reopen")
	}
	if claimed.TaskID != "task-1" || claimed.AttemptCount != 2 {
		t.Fatalf("unexpected reclaimed task: %+v", claimed)
	}
}

func TestEnsureSemanticTaskQueuedRequeuesSucceededTask(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-1", "file-1", 1, now, now)); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected claim to find task")
	}
	if err := s.AckSemanticTask(ctx, claimed.TaskID, claimed.Receipt); err != nil {
		t.Fatal(err)
	}

	queued, err := s.EnsureSemanticTaskQueued(ctx, &semantic.Task{
		TaskID:          "task-requeue",
		TaskType:        semantic.TaskTypeEmbed,
		ResourceID:      "file-1",
		ResourceVersion: 1,
		Status:          semantic.TaskQueued,
		MaxAttempts:     3,
		AvailableAt:     now.Add(2 * time.Second),
		CreatedAt:       now.Add(2 * time.Second),
		UpdatedAt:       now.Add(2 * time.Second),
	})
	if err != nil {
		t.Fatal(err)
	}
	if !queued {
		t.Fatal("expected ensure to requeue succeeded task")
	}

	task := mustGetSemanticTask(t, s, claimed.TaskID)
	if task.Status != semantic.TaskQueued {
		t.Fatalf("status=%q, want %q", task.Status, semantic.TaskQueued)
	}
	if task.CompletedAt != nil {
		t.Fatal("completed_at should be cleared when re-queued")
	}
	if task.Receipt != "" || task.LeasedAt != nil || task.LeaseUntil != nil {
		t.Fatalf("lease fields should be cleared when re-queued: %+v", task)
	}
}

func TestEnsureSemanticTaskQueuedSupersedesProcessingEmbedFromOlderGeneration(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	oldTask := newSemanticTask("task-old", "file-1", 1, now, now)
	oldTask.PayloadJSON = []byte(`{"generation":1}`)
	if _, err := s.EnqueueSemanticTask(ctx, oldTask); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}

	newTask := newSemanticTask("task-new", "file-1", 1, now.Add(2*time.Second), now.Add(2*time.Second))
	newTask.PayloadJSON = []byte(`{"generation":2}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, newTask)
	if err != nil {
		t.Fatal(err)
	}
	if !queued {
		t.Fatal("new generation did not supersede processing embed task")
	}

	got := mustGetSemanticTask(t, s, claimed.TaskID)
	var payload semantic.EmbedTaskPayload
	if err := json.Unmarshal(got.PayloadJSON, &payload); err != nil {
		t.Fatalf("decode superseded payload: %v", err)
	}
	if got.Status != semantic.TaskQueued || payload.Generation != 2 {
		t.Fatalf("superseded task=%+v payload=%s", got, got.PayloadJSON)
	}
	if got.Receipt != "" || got.LeasedAt != nil || got.LeaseUntil != nil {
		t.Fatalf("superseded task retained lease: %+v", got)
	}
	if err := s.AckSemanticTask(ctx, claimed.TaskID, claimed.Receipt); !errors.Is(err, semantic.ErrTaskLeaseMismatch) {
		t.Fatalf("old worker ack err=%v, want lease mismatch", err)
	}
}

func TestEnsureSemanticTaskQueuedKeepsProcessingEmbedForSameGeneration(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	oldTask := newSemanticTask("task-old", "file-1", 1, now, now)
	oldTask.PayloadJSON = []byte(`{"generation":1}`)
	if _, err := s.EnqueueSemanticTask(ctx, oldTask); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}

	duplicate := newSemanticTask("task-duplicate", "file-1", 1, now.Add(2*time.Second), now.Add(2*time.Second))
	duplicate.PayloadJSON = []byte(`{"generation":1}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, duplicate)
	if err != nil {
		t.Fatal(err)
	}
	if queued {
		t.Fatal("same generation should not supersede processing embed task")
	}

	got := mustGetSemanticTask(t, s, claimed.TaskID)
	if got.Status != semantic.TaskProcessing || got.Receipt != claimed.Receipt {
		t.Fatalf("processing task lease changed: %+v", got)
	}
}

func TestEnsureSemanticTaskQueuedRejectsOlderTenantEmbeddingConfig(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	current := newSemanticTask("task-current", "file-1", 1, now, now)
	current.PayloadJSON = []byte(`{"generation":2,"config_revision":2,"tenant_managed":true}`)
	if _, err := s.EnqueueSemanticTask(ctx, current); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim current task found=%v err=%v", found, err)
	}

	stale := newSemanticTask("task-stale", "file-1", 1, now.Add(time.Second), now.Add(time.Second))
	stale.PayloadJSON = []byte(`{"generation":1,"config_revision":1,"tenant_managed":true}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, stale)
	if err != nil {
		t.Fatal(err)
	}
	if queued {
		t.Fatal("older tenant embedding config replaced current work")
	}

	got := mustGetSemanticTask(t, s, current.TaskID)
	var payload semantic.EmbedTaskPayload
	if err := json.Unmarshal(got.PayloadJSON, &payload); err != nil {
		t.Fatal(err)
	}
	if payload.Generation != 2 || payload.ConfigRevision != 2 || !payload.TenantManaged {
		t.Fatalf("payload=%+v, want current tenant config", payload)
	}
	if got.Status != semantic.TaskProcessing || got.Receipt != claimed.Receipt {
		t.Fatalf("status/receipt=%s/%q, want leased current task %q", got.Status, got.Receipt, claimed.Receipt)
	}
}

func TestEnsureSemanticTaskQueuedRestoresAttemptsForNewEmbeddingConfig(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	failed := newSemanticTask("task-failed", "file-1", 1, now, now)
	failed.PayloadJSON = []byte(`{"generation":1,"config_revision":1,"tenant_managed":true}`)
	if _, err := s.EnqueueSemanticTask(ctx, failed); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().ExecContext(ctx, `UPDATE semantic_tasks
		SET status = ?, attempt_count = max_attempts, completed_at = ? WHERE task_id = ?`,
		semantic.TaskDeadLettered, now, failed.TaskID); err != nil {
		t.Fatal(err)
	}

	repaired := newSemanticTask("task-repaired", "file-1", 1, now.Add(time.Second), now.Add(time.Second))
	repaired.PayloadJSON = []byte(`{"generation":1,"config_revision":2,"tenant_managed":true}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, repaired)
	if err != nil {
		t.Fatal(err)
	}
	if !queued {
		t.Fatal("new tenant embedding config did not revive failed work")
	}

	got := mustGetSemanticTask(t, s, failed.TaskID)
	if got.Status != semantic.TaskQueued || got.AttemptCount != 0 {
		t.Fatalf("status/attempt_count=%s/%d, want queued/0", got.Status, got.AttemptCount)
	}
}

func TestEnsureSemanticTaskQueuedRestoresAttemptsForLegacyDeadLetterWithoutPayload(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	failed := newSemanticTask("task-failed-legacy", "file-1", 1, now, now)
	failed.PayloadJSON = []byte(`{"generation":1}`)
	if _, err := s.EnqueueSemanticTask(ctx, failed); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().ExecContext(ctx, `UPDATE semantic_tasks
		SET status = ?, attempt_count = max_attempts, payload_json = NULL, completed_at = ? WHERE task_id = ?`,
		semantic.TaskDeadLettered, now, failed.TaskID); err != nil {
		t.Fatal(err)
	}

	repaired := newSemanticTask("task-repaired-legacy", "file-1", 1, now.Add(time.Second), now.Add(time.Second))
	repaired.PayloadJSON = []byte(`{"generation":2,"config_revision":2,"tenant_managed":true}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, repaired)
	if err != nil {
		t.Fatal(err)
	}
	if !queued {
		t.Fatal("new tenant embedding config did not revive legacy dead-letter work")
	}

	got := mustGetSemanticTask(t, s, failed.TaskID)
	if got.Status != semantic.TaskQueued || got.AttemptCount != 0 {
		t.Fatalf("status/attempt_count=%s/%d, want queued/0", got.Status, got.AttemptCount)
	}
}

func TestEnsureSemanticTaskQueuedKeepsAttemptsForSameEmbeddingConfig(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	failed := newSemanticTask("task-failed-same-config", "file-1", 1, now, now)
	failed.PayloadJSON = []byte(`{"generation":1,"config_revision":2,"tenant_managed":true}`)
	if _, err := s.EnqueueSemanticTask(ctx, failed); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().ExecContext(ctx, `UPDATE semantic_tasks
		SET status = ?, attempt_count = max_attempts, completed_at = ? WHERE task_id = ?`,
		semantic.TaskDeadLettered, now, failed.TaskID); err != nil {
		t.Fatal(err)
	}

	same := newSemanticTask("task-same-config", "file-1", 1, now.Add(time.Second), now.Add(time.Second))
	same.PayloadJSON = []byte(`{"generation":1,"config_revision":2,"tenant_managed":true}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, same)
	if err != nil || !queued {
		t.Fatalf("same config requeue queued=%v err=%v", queued, err)
	}
	got := mustGetSemanticTask(t, s, failed.TaskID)
	if got.AttemptCount != got.MaxAttempts {
		t.Fatalf("attempt_count=%d, want preserved %d", got.AttemptCount, got.MaxAttempts)
	}
}

func TestEnsureSemanticTaskQueuedRejectsDefaultOverTenantConfig(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	custom := newSemanticTask("task-custom", "file-1", 1, now, now)
	custom.PayloadJSON = []byte(`{"generation":1,"config_revision":1,"tenant_managed":true}`)
	if _, err := s.EnqueueSemanticTask(ctx, custom); err != nil {
		t.Fatal(err)
	}

	staleDefault := newSemanticTask("task-default", "file-1", 1, now.Add(time.Second), now.Add(time.Second))
	staleDefault.PayloadJSON = []byte(`{"generation":99,"config_revision":99}`)
	queued, err := s.EnsureSemanticTaskQueued(ctx, staleDefault)
	if err != nil {
		t.Fatal(err)
	}
	if queued {
		t.Fatal("stale process default replaced tenant-managed work")
	}
}

func TestSemanticTaskRetryDeadLetterPreservesEmbedConfigIdentity(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	task := newSemanticTask("task-embed-dead", "file-1", 1, now, now)
	task.MaxAttempts = 1
	task.PayloadJSON = []byte(`{"generation":1,"config_revision":1,"tenant_managed":true,"secret":"must-not-survive"}`)
	if _, err := s.EnqueueSemanticTask(ctx, task); err != nil {
		t.Fatal(err)
	}
	claimed, found, err := s.ClaimSemanticTask(ctx, now.Add(time.Second), time.Minute, semantic.TaskTypeEmbed)
	if err != nil || !found {
		t.Fatalf("claim found=%v err=%v", found, err)
	}
	if err := s.RetrySemanticTask(ctx, claimed.TaskID, claimed.Receipt, now.Add(2*time.Second), "provider failure"); err != nil {
		t.Fatal(err)
	}

	got := mustGetSemanticTask(t, s, task.TaskID)
	if got.Status != semantic.TaskDeadLettered {
		t.Fatalf("status=%s, want dead_lettered", got.Status)
	}
	var payload semantic.EmbedTaskPayload
	if err := json.Unmarshal(got.PayloadJSON, &payload); err != nil {
		t.Fatal(err)
	}
	if payload.Generation != 1 || payload.ConfigRevision != 1 || !payload.TenantManaged {
		t.Fatalf("payload=%+v, want preserved embed config identity", payload)
	}
	if bytes.Contains(got.PayloadJSON, []byte("must-not-survive")) {
		t.Fatalf("dead-letter payload retained fields outside embed config identity: %s", got.PayloadJSON)
	}
}

func TestObserveSemanticTasksSummarizesQueueState(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Now().UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-queued", "file-queued", 1, base.Add(-3*time.Second), base)); err != nil {
		t.Fatal(err)
	}
	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-future", "file-future", 1, base.Add(30*time.Second), base.Add(time.Second))); err != nil {
		t.Fatal(err)
	}
	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-processing", "file-processing", 1, base.Add(-5*time.Second), base.Add(2*time.Second))); err != nil {
		t.Fatal(err)
	}
	processingTask, found, err := s.ClaimSemanticTask(ctx, base.Add(time.Second), time.Minute)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected processing task to be claimed")
	}

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-dead", "file-dead", 1, base.Add(-4*time.Second), base.Add(3*time.Second))); err != nil {
		t.Fatal(err)
	}
	if _, err := s.DB().Exec(`UPDATE semantic_tasks SET max_attempts = 1 WHERE task_id = ?`, "task-dead"); err != nil {
		t.Fatal(err)
	}
	deadTask, found, err := s.ClaimSemanticTask(ctx, base.Add(2*time.Second), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	if !found {
		t.Fatal("expected dead-letter task to be claimed")
	}
	if err := s.RetrySemanticTask(ctx, deadTask.TaskID, deadTask.Receipt, base.Add(3*time.Second), "permanent failure"); err != nil {
		t.Fatal(err)
	}
	deadAfterRetry := mustGetSemanticTask(t, s, deadTask.TaskID)
	if deadAfterRetry.PayloadJSON != nil {
		t.Fatalf("dead-letter retry payload_json = %q, want nil", deadAfterRetry.PayloadJSON)
	}

	obs, err := s.ObserveSemanticTasks(ctx, base.Add(5*time.Second))
	if err != nil {
		t.Fatal(err)
	}
	if obs.Queued != 2 {
		t.Fatalf("queued=%d, want 2", obs.Queued)
	}
	if obs.Processing != 1 {
		t.Fatalf("processing=%d, want 1", obs.Processing)
	}
	if obs.DeadLettered != 1 {
		t.Fatalf("dead_lettered=%d, want 1", obs.DeadLettered)
	}
	if obs.OldestClaimableAvailableAt == nil {
		t.Fatal("expected oldest claimable available_at")
	}
	wantOldest := base.Add(-3 * time.Second).Round(time.Millisecond)
	if !obs.OldestClaimableAvailableAt.Equal(wantOldest) {
		t.Fatalf("oldest claimable=%s, want %s", obs.OldestClaimableAvailableAt.UTC(), wantOldest.UTC())
	}

	task := mustGetSemanticTask(t, s, processingTask.TaskID)
	if task.Status != semantic.TaskProcessing {
		t.Fatalf("processing task status=%q, want %q", task.Status, semantic.TaskProcessing)
	}
}

func TestObserveSemanticTasksSkipsFutureOnlyQueueLag(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	base := time.Unix(1711600700, 0).UTC()

	if _, err := s.EnqueueSemanticTask(ctx, newSemanticTask("task-future", "file-future", 1, base.Add(10*time.Second), base)); err != nil {
		t.Fatal(err)
	}

	obs, err := s.ObserveSemanticTasks(ctx, base)
	if err != nil {
		t.Fatal(err)
	}
	if obs.Queued != 1 {
		t.Fatalf("queued=%d, want 1", obs.Queued)
	}
	if obs.Processing != 0 {
		t.Fatalf("processing=%d, want 0", obs.Processing)
	}
	if obs.DeadLettered != 0 {
		t.Fatalf("dead_lettered=%d, want 0", obs.DeadLettered)
	}
	if obs.OldestClaimableAvailableAt != nil {
		t.Fatalf("oldest claimable=%s, want nil", obs.OldestClaimableAvailableAt.UTC())
	}
}

func mustGetSemanticTask(t *testing.T, s *Store, taskID string) *semantic.Task {
	t.Helper()
	row := s.DB().QueryRow(`SELECT task_id, task_type, resource_id, resource_version, status,
		attempt_count, max_attempts, receipt, leased_at, lease_until, available_at,
		payload_json, last_error, created_at, updated_at, completed_at
		FROM semantic_tasks WHERE task_id = ?`, taskID)
	task, err := scanSemanticTask(row)
	if err != nil {
		t.Fatalf("get semantic task %s: %v", taskID, err)
	}
	return task
}

func countSemanticTasks(t *testing.T, s *Store) int {
	t.Helper()
	var count int
	if err := s.DB().QueryRow(`SELECT COUNT(*) FROM semantic_tasks`).Scan(&count); err != nil {
		t.Fatalf("count semantic tasks: %v", err)
	}
	return count
}