package store

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

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

func TestSecurityEventCredentialReferenceSwitchesRemainCrashRecoverable(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()

	connectionID := uuid.NewString()
	pushReference := "ssf-push-" + connectionID
	managementReference := "ssf-management-" + connectionID
	nextReference := "ssf-push-next-" + uuid.NewString()
	newManagementReference := "ssf-management-next-" + uuid.NewString()
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`,
			[]string{pushReference, managementReference, nextReference, newManagementReference})
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}

	for _, reference := range []string{pushReference, managementReference} {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
	})
	if err != nil {
		t.Fatal(err)
	}
	assertIdentitySecretQueueState(t, ctx, db, pushReference, false)
	assertIdentitySecretQueueState(t, ctx, db, managementReference, false)

	if err := db.QueueIdentitySecretCleanup(ctx, newManagementReference, time.Now().Add(10*time.Minute)); err != nil {
		t.Fatal(err)
	}
	if _, err := db.SetSecurityEventManagementCredential(ctx, connectionID, "machine-client", newManagementReference, connection.Version); err != nil {
		t.Fatal(err)
	}
	assertIdentitySecretQueueState(t, ctx, db, newManagementReference, false)
	assertIdentitySecretQueueState(t, ctx, db, managementReference, true)

	if err := db.QueueIdentitySecretCleanup(ctx, nextReference, time.Now().Add(10*time.Minute)); err != nil {
		t.Fatal(err)
	}
	connection, err = db.SecurityEventConnection(ctx)
	if err != nil {
		t.Fatal(err)
	}
	connection, err = db.SetSecurityEventNextCredential(ctx, connectionID, nextReference, connection.Version)
	if err != nil {
		t.Fatal(err)
	}
	assertIdentitySecretQueueState(t, ctx, db, nextReference, false)
	if _, err := db.PromoteSecurityEventCredential(ctx, connection.ConnectionID, nextReference); err != nil {
		t.Fatal(err)
	}
	assertIdentitySecretQueueState(t, ctx, db, pushReference, true)
}

func TestDiscardPreparedSecurityEventConnectionQueuesAllSecretsAtomically(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	ownerKey := "identity-pairing-ssf-" + uuid.NewString()
	pushReference := "ssf-push-" + connectionID
	managementReference := "ssf-management-" + connectionID
	nextReference := "ssf-push-next-" + uuid.NewString()
	references := []string{pushReference, managementReference, nextReference}
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	for _, reference := range []string{pushReference, managementReference} {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: ownerKey,
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := db.QueueIdentitySecretCleanup(ctx, nextReference, time.Now().Add(10*time.Minute)); err != nil {
		t.Fatal(err)
	}
	connection, err = db.SetSecurityEventNextCredential(ctx, connectionID, nextReference, connection.Version)
	if err != nil {
		t.Fatal(err)
	}

	if err := db.DiscardPreparedSecurityEventConnection(ctx, connectionID, ownerKey, connection.Version); err != nil {
		t.Fatal(err)
	}
	if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
		t.Fatalf("discarded prepared connection still exists: %v", err)
	}
	for _, reference := range references {
		assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
	}
}

func TestConcurrentPreparedDiscardAndBindPreserveTheWinningGeneration(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	ownerKey := "identity-pairing-ssf-" + uuid.NewString()
	pushReference := "ssf-push-" + connectionID
	managementReference := "ssf-management-" + connectionID
	references := []string{pushReference, managementReference}
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	for _, reference := range references {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: ownerKey,
	})
	if err != nil {
		t.Fatal(err)
	}

	type result struct {
		operation string
		err       error
	}
	start := make(chan struct{})
	results := make(chan result, 2)
	audience := "urn:easyai:ssf:receiver:" + uuid.NewString()
	streamID := uuid.NewString()
	go func() {
		<-start
		_, err := db.BindSecurityEventConnection(ctx, connectionID, audience, streamID, "verifying", connection.Version)
		results <- result{operation: "bind", err: err}
	}()
	go func() {
		<-start
		err := db.DiscardPreparedSecurityEventConnection(ctx, connectionID, ownerKey, connection.Version)
		results <- result{operation: "discard", err: err}
	}()
	close(start)

	outcomes := map[string]error{}
	for range 2 {
		outcome := <-results
		outcomes[outcome.operation] = outcome.err
	}
	bindWon := outcomes["bind"] == nil
	discardWon := outcomes["discard"] == nil
	if bindWon == discardWon {
		t.Fatalf("concurrent outcomes bind=%v discard=%v, want exactly one winner", outcomes["bind"], outcomes["discard"])
	}
	loser := outcomes["bind"]
	if bindWon {
		loser = outcomes["discard"]
	}
	if !errors.Is(loser, ErrSecurityEventConnectionConflict) {
		t.Fatalf("concurrent loser error=%v, want connection conflict", loser)
	}
	if bindWon {
		persisted, err := db.SecurityEventConnection(ctx)
		if err != nil || persisted.StreamID == nil || *persisted.StreamID != streamID {
			t.Fatalf("winning bound generation was not preserved: connection=%+v err=%v", persisted, err)
		}
		for _, reference := range references {
			assertIdentitySecretQueueState(t, ctx, db, reference, false)
		}
		return
	}
	if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
		t.Fatalf("winning discard left a connection behind: %v", err)
	}
	for _, reference := range references {
		assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
	}
}

func TestConcurrentSecurityEventNextCredentialsAdoptOnlyOneStagedSecret(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	pushReference := "ssf-push-" + connectionID
	managementReference := "ssf-management-" + connectionID
	candidates := []string{
		"ssf-push-next-" + uuid.NewString(),
		"ssf-push-next-" + uuid.NewString(),
	}
	allReferences := append([]string{pushReference, managementReference}, candidates...)
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, allReferences)
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	for _, reference := range []string{pushReference, managementReference} {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
	})
	if err != nil {
		t.Fatal(err)
	}
	for _, reference := range candidates {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}

	type result struct {
		reference string
		err       error
	}
	start := make(chan struct{})
	results := make(chan result, len(candidates))
	var workers sync.WaitGroup
	for _, reference := range candidates {
		workers.Add(1)
		go func(reference string) {
			defer workers.Done()
			<-start
			_, err := db.SetSecurityEventNextCredential(ctx, connectionID, reference, connection.Version)
			results <- result{reference: reference, err: err}
		}(reference)
	}
	close(start)
	workers.Wait()
	close(results)

	var adopted, pending string
	for result := range results {
		switch {
		case result.err == nil:
			if adopted != "" {
				t.Fatal("more than one concurrent next credential was adopted")
			}
			adopted = result.reference
		case errors.Is(result.err, ErrSecurityEventConnectionConflict):
			pending = result.reference
		default:
			t.Fatalf("concurrent next credential returned unexpected error: %v", result.err)
		}
	}
	if adopted == "" || pending == "" {
		t.Fatalf("concurrent next credential outcomes adopted=%t pending=%t", adopted != "", pending != "")
	}
	connection, err = db.SecurityEventConnection(ctx)
	if err != nil || connection.NextCredentialRef == nil || *connection.NextCredentialRef != adopted {
		t.Fatalf("connection did not retain the sole adopted next credential: %v", err)
	}
	assertIdentitySecretQueueState(t, ctx, db, adopted, false)
	assertIdentitySecretQueueStatus(t, ctx, db, pending, "pending")
}

func TestRetiringSecurityEventConnectionCannotAdoptManagementSecret(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	pushReference := "ssf-push-" + connectionID
	managementReference := "ssf-management-" + connectionID
	newManagementReference := "ssf-management-next-" + uuid.NewString()
	allReferences := []string{pushReference, managementReference, newManagementReference}
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, allReferences)
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	for _, reference := range []string{pushReference, managementReference} {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
	})
	if err != nil {
		t.Fatal(err)
	}
	connection, err = db.TransitionSecurityEventConnectionLifecycle(
		ctx, connectionID, connection.LifecycleStatus, connection.Version, "retiring", nil,
	)
	if err != nil {
		t.Fatal(err)
	}
	if err := db.QueueIdentitySecretCleanup(ctx, newManagementReference, time.Now().Add(10*time.Minute)); err != nil {
		t.Fatal(err)
	}

	if _, err := db.SetSecurityEventManagementCredential(ctx, connectionID, "machine-client", newManagementReference, connection.Version); !errors.Is(err, ErrSecurityEventConnectionConflict) {
		t.Fatalf("retiring connection management adoption error=%v, want conflict", err)
	}
	connection, err = db.SecurityEventConnection(ctx)
	if err != nil || connection.ManagementCredentialRef == nil || *connection.ManagementCredentialRef != managementReference {
		t.Fatalf("retiring connection changed its management credential: %v", err)
	}
	assertIdentitySecretQueueStatus(t, ctx, db, newManagementReference, "pending")
}

func TestSecurityEventRetirementRequiresTheCurrentConnectionGeneration(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	pushReference := "ssf-push-" + connectionID
	managementReference := "ssf-management-" + connectionID
	references := []string{pushReference, managementReference}
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	for _, reference := range references {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
	})
	if err != nil {
		t.Fatal(err)
	}

	if _, err := db.TransitionSecurityEventConnectionLifecycle(
		ctx, connectionID, "connecting", connection.Version-1, "retiring", nil,
	); !errors.Is(err, ErrSecurityEventConnectionConflict) {
		t.Fatalf("stale retirement transition error=%v, want conflict", err)
	}
	current, err := db.SecurityEventConnection(ctx)
	if err != nil || current.LifecycleStatus != "connecting" || current.Version != connection.Version {
		t.Fatalf("stale transition changed connection generation: lifecycle=%q version=%d err=%v",
			current.LifecycleStatus, current.Version, err)
	}
	current, err = db.TransitionSecurityEventConnectionLifecycle(
		ctx, connectionID, "connecting", connection.Version, "retiring", nil,
	)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := db.TransitionSecurityEventConnectionLifecycle(
		ctx, connectionID, "retiring", current.Version, "error", nil,
	); !errors.Is(err, ErrSecurityEventConnectionConflict) {
		t.Fatalf("retiring generation revival error=%v, want conflict", err)
	}
	if err := db.FinalizeRetiringSecurityEventConnection(ctx, connectionID, connection.Version); !errors.Is(err, ErrSecurityEventConnectionConflict) {
		t.Fatalf("stale retirement deletion error=%v, want conflict", err)
	}
	if persisted, err := db.SecurityEventConnection(ctx); err != nil || persisted.LifecycleStatus != "retiring" || persisted.Version != current.Version {
		t.Fatalf("stale retirement deletion removed current generation: lifecycle=%q version=%d err=%v",
			persisted.LifecycleStatus, persisted.Version, err)
	}
	if err := db.FinalizeRetiringSecurityEventConnection(ctx, connectionID, current.Version); err != nil {
		t.Fatal(err)
	}
	if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
		t.Fatalf("current retiring generation was not deleted: %v", err)
	}
	for _, reference := range references {
		assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
	}
}

func TestConcurrentSecurityEventRetirementFinalizationQueuesOneGenerationAtomically(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	pushReference := "ssf-push-retire-worker-" + connectionID
	managementReference := "ssf-management-retire-worker-" + connectionID
	nextReference := "ssf-push-next-retire-worker-" + uuid.NewString()
	references := []string{pushReference, managementReference, nextReference}
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
	})
	if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
		t.Fatal(err)
	}
	for _, reference := range []string{pushReference, managementReference} {
		if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
	})
	if err != nil {
		t.Fatal(err)
	}
	if err := db.QueueIdentitySecretCleanup(ctx, nextReference, time.Now().Add(10*time.Minute)); err != nil {
		t.Fatal(err)
	}
	connection, err = db.SetSecurityEventNextCredential(ctx, connectionID, nextReference, connection.Version)
	if err != nil {
		t.Fatal(err)
	}
	connection, err = db.TransitionSecurityEventConnectionLifecycle(
		ctx, connectionID, connection.LifecycleStatus, connection.Version, "retiring", nil,
	)
	if err != nil {
		t.Fatal(err)
	}

	start := make(chan struct{})
	results := make(chan error, 2)
	var workers sync.WaitGroup
	for range 2 {
		workers.Add(1)
		go func() {
			defer workers.Done()
			<-start
			results <- db.FinalizeRetiringSecurityEventConnection(ctx, connectionID, connection.Version)
		}()
	}
	close(start)
	workers.Wait()
	close(results)
	succeeded, conflicted := 0, 0
	for err := range results {
		switch {
		case err == nil:
			succeeded++
		case errors.Is(err, ErrSecurityEventConnectionConflict):
			conflicted++
		default:
			t.Fatalf("unexpected concurrent finalization error: %v", err)
		}
	}
	if succeeded != 1 || conflicted != 1 {
		t.Fatalf("concurrent finalization succeeded=%d conflicted=%d, want 1/1", succeeded, conflicted)
	}
	if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
		t.Fatalf("finalized connection still exists: %v", err)
	}
	for _, reference := range references {
		assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
	}
}

func TestClaimedStagingSecretsCannotBeAdoptedByAConnection(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	connectionID := uuid.NewString()
	pushReference := "ssf-push-claimed-" + uuid.NewString()
	managementReference := "ssf-management-claimed-" + uuid.NewString()
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`,
			[]string{pushReference, managementReference})
	})

	if err := db.QueueIdentitySecretCleanup(ctx, pushReference, time.Now().Add(10*time.Minute)); err != nil {
		t.Fatal(err)
	}
	if err := db.QueueIdentitySecretCleanup(ctx, managementReference, time.Now().Add(-time.Minute)); err != nil {
		t.Fatal(err)
	}
	claims, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
	if err != nil || len(claims) != 1 || claims[0].Reference != managementReference {
		t.Fatalf("claim staging management Secret: count=%d reference_matches=%t err=%v",
			len(claims), len(claims) == 1 && claims[0].Reference == managementReference, err)
	}

	_, err = db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
		ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
		EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
		ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
		IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
	})
	if !errors.Is(err, ErrIdentitySecretCleanupConflict) {
		t.Fatalf("claimed Secret adoption error=%v, want cleanup conflict", err)
	}
	if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
		t.Fatalf("failed adoption committed connection: %v", err)
	}
	assertIdentitySecretQueueStatus(t, ctx, db, pushReference, "pending")
	assertIdentitySecretQueueStatus(t, ctx, db, managementReference, "claimed")
}

func TestIdentitySecretCleanupCompletionRejectsExpiredClaimABA(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	reference := "ssf-management-aba-" + uuid.NewString()
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1`, reference)
	})
	if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(-time.Minute)); err != nil {
		t.Fatal(err)
	}
	first, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
	if err != nil || len(first) != 1 {
		t.Fatalf("first claim count=%d err=%v", len(first), err)
	}
	if _, err := db.pool.Exec(ctx, `UPDATE gateway_identity_secret_cleanup_queue SET lease_expires_at=now()-interval '1 second' WHERE secret_ref=$1`, reference); err != nil {
		t.Fatal(err)
	}
	second, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
	if err != nil || len(second) != 1 || second[0].ClaimToken == first[0].ClaimToken {
		t.Fatalf("second claim count=%d token_rotated=%t err=%v",
			len(second), len(second) == 1 && second[0].ClaimToken != first[0].ClaimToken, err)
	}
	completed, err := db.CompleteIdentitySecretCleanup(ctx, reference, first[0].ClaimToken)
	if err != nil || completed {
		t.Fatalf("stale completion completed=%t err=%v", completed, err)
	}
	assertIdentitySecretQueueStatus(t, ctx, db, reference, "claimed")
	completed, err = db.CompleteIdentitySecretCleanup(ctx, reference, second[0].ClaimToken)
	if err != nil || !completed {
		t.Fatalf("current completion completed=%t err=%v", completed, err)
	}
}

func TestConcurrentIdentitySecretCleanupClaimsAreDisjoint(t *testing.T) {
	db := newSecurityEventSecretCleanupPostgresStore(t)
	ctx := context.Background()
	references := make([]string, 20)
	for index := range references {
		references[index] = "ssf-cleanup-concurrent-" + uuid.NewString()
		if err := db.QueueIdentitySecretCleanup(ctx, references[index], time.Now().Add(-time.Minute)); err != nil {
			t.Fatal(err)
		}
	}
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
	})

	start := make(chan struct{})
	results := make(chan []IdentitySecretCleanupClaim, 2)
	errorsCh := make(chan error, 2)
	var workers sync.WaitGroup
	for range 2 {
		workers.Add(1)
		go func() {
			defer workers.Done()
			<-start
			claims, err := db.ClaimIdentitySecretCleanups(ctx, len(references), time.Minute)
			results <- claims
			errorsCh <- err
		}()
	}
	close(start)
	workers.Wait()
	close(results)
	close(errorsCh)
	for err := range errorsCh {
		if err != nil {
			t.Fatal(err)
		}
	}
	claimed := map[string]string{}
	for claims := range results {
		for _, claim := range claims {
			if previous := claimed[claim.Reference]; previous != "" {
				t.Fatalf("Secret %q was claimed by more than one worker", claim.Reference)
			}
			claimed[claim.Reference] = claim.ClaimToken
		}
	}
	if len(claimed) != len(references) {
		t.Fatalf("claimed %d Secrets, want %d", len(claimed), len(references))
	}
}

func newSecurityEventSecretCleanupPostgresStore(t *testing.T) *Store {
	t.Helper()
	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
	if databaseURL == "" {
		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run identity Secret cleanup PostgreSQL integration tests")
	}
	ctx := context.Background()
	probe, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	var databaseName string
	if err := probe.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
		probe.Close()
		t.Fatal(err)
	}
	probe.Close()
	if !strings.Contains(strings.ToLower(databaseName), "test") {
		t.Fatalf("refusing to migrate non-test database %q", databaseName)
	}
	applyOIDCJITTestMigrations(t, ctx, databaseURL)
	db, err := Connect(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(db.Close)
	return db
}

func assertIdentitySecretQueueState(t *testing.T, ctx context.Context, db *Store, reference string, expected bool) {
	t.Helper()
	var exists bool
	if err := db.pool.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1)`, reference).Scan(&exists); err != nil {
		t.Fatal(err)
	}
	if exists != expected {
		t.Fatalf("Secret cleanup reference %q exists=%t want=%t", reference, exists, expected)
	}
}

func assertIdentitySecretQueueStatus(t *testing.T, ctx context.Context, db *Store, reference, expected string) {
	t.Helper()
	var status string
	if err := db.pool.QueryRow(ctx, `SELECT status FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1`, reference).Scan(&status); err != nil {
		t.Fatal(err)
	}
	if status != expected {
		t.Fatalf("Secret cleanup reference %q status=%q want=%q", reference, status, expected)
	}
}
