package case_driver

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-e2e/cases"
	"github.com/tidbcloud/dataflow-service-e2e/internal/config"
	"github.com/tidbcloud/dataflow-service-e2e/internal/log"
	. "github.com/tidbcloud/dataflow-service-e2e/internal/model"
)

type mockCase struct {
	cases.BaseCase
	name       string
	reqs       map[string]DependencyRequirement
	assignedMu sync.Mutex
	assigned   map[string]DependencyEntity
	assignedCh chan string
	onFulfill  func(string, DependencyEntity)
}

func (m *mockCase) Name() string                                                    { return m.name }
func (m *mockCase) DeclareDependencyRequirements() map[string]DependencyRequirement { return m.reqs }
func (m *mockCase) FulfillDependencyRequirement(reqName string, entity DependencyEntity) {
	m.BaseCase.FulfillDependencyRequirement(reqName, entity)
	m.assignedMu.Lock()
	if m.assigned == nil {
		m.assigned = make(map[string]DependencyEntity)
	}
	m.assigned[reqName] = entity
	assignedCh := m.assignedCh
	onFulfill := m.onFulfill
	m.assignedMu.Unlock()
	if assignedCh != nil {
		assignedCh <- reqName
	}
	if onFulfill != nil {
		onFulfill(reqName, entity)
	}
}
func (m *mockCase) Run()                    {}
func (m *mockCase) Tags() map[string]string { return nil }

type dummyRequirement struct {
	typ        string
	dimensions map[string]string
	dedicated  bool
}

func (r *dummyRequirement) Type() string                  { return r.typ }
func (r *dummyRequirement) Dimensions() map[string]string { return r.dimensions }
func (r *dummyRequirement) IsDedicated() bool             { return r.dedicated }

type dummyEntity struct {
	typ             string
	dimensions      map[string]string
	createdByDriver bool
}

func (e *dummyEntity) Type() string                  { return e.typ }
func (e *dummyEntity) Dimensions() map[string]string { return e.dimensions }
func (e *dummyEntity) IsCreatedByDriver() bool       { return e.createdByDriver }

type streamDependency struct {
	typ         string
	existed     []DependencyEntity
	provisionFn func(context.Context, EntitiesStats) (<-chan DependencyEntity, error)
}

func (m *streamDependency) Type() string { return m.typ }

func (m *streamDependency) GetExistingUsableEntities() ([]DependencyEntity, error) {
	return m.existed, nil
}

func (m *streamDependency) ProvisionEntities(ctx context.Context, stats EntitiesStats) (<-chan DependencyEntity, error) {
	if m.provisionFn == nil {
		ch := make(chan DependencyEntity)
		close(ch)
		return ch, nil
	}
	return m.provisionFn(ctx, stats)
}

func (m *streamDependency) Delete(...DependencyEntity) error { return nil }

// observableDependency is a no-I/O Dependency implementation used to exercise
// the manager, dispatcher, and cleanup loop together. Its Delete calls are
// recorded under a mutex so tests remain race-detector safe.
type observableDependency struct {
	streamDependency

	mu          sync.Mutex
	deleteCalls []DependencyEntity
	deleteFn    func(DependencyEntity) error
	deleteStart chan DependencyEntity
}

type recordingLogger struct {
	mu     sync.Mutex
	errors []string
}

func (*recordingLogger) Info(string, ...any) {}
func (*recordingLogger) Warn(string, ...any) {}
func (l *recordingLogger) Error(format string, a ...any) {
	l.mu.Lock()
	defer l.mu.Unlock()
	l.errors = append(l.errors, fmt.Sprintf(format, a...))
}
func (*recordingLogger) Sync() {}

func (l *recordingLogger) lastError(t *testing.T) string {
	t.Helper()
	l.mu.Lock()
	defer l.mu.Unlock()
	require.NotEmpty(t, l.errors)
	return l.errors[len(l.errors)-1]
}

func (d *observableDependency) Delete(entities ...DependencyEntity) error {
	for _, entity := range entities {
		d.mu.Lock()
		d.deleteCalls = append(d.deleteCalls, entity)
		d.mu.Unlock()

		if d.deleteStart != nil {
			d.deleteStart <- entity
		}
		if d.deleteFn != nil {
			if err := d.deleteFn(entity); err != nil {
				return err
			}
		}
	}
	return nil
}

func (d *observableDependency) deleteCount(entity DependencyEntity) int {
	d.mu.Lock()
	defer d.mu.Unlock()

	count := 0
	for _, deleted := range d.deleteCalls {
		if deleted == entity {
			count++
		}
	}
	return count
}

func makeEntity(region string) *TiDBDedicatedEntity {
	return &TiDBDedicatedEntity{
		TiDBDedicatedDimensions: TiDBDedicatedDimensions{
			CloudProvider: AWS,
			Region:        region,
		},
		CreatedByDriver: true,
	}
}

func makeDedicatedReq(region string) *TiDBDedicatedRequirement {
	return &TiDBDedicatedRequirement{
		Dedicated: true,
		TiDBDedicatedDimensions: TiDBDedicatedDimensions{
			CloudProvider: AWS,
			Region:        region,
		},
	}
}

func makeSharedReq(region string) *TiDBDedicatedRequirement {
	return &TiDBDedicatedRequirement{
		Dedicated: false,
		TiDBDedicatedDimensions: TiDBDedicatedDimensions{
			CloudProvider: AWS,
			Region:        region,
		},
	}
}

func newTestMgr(lgr log.SimpleLogger, casesMap map[string]cases.E2ECase, dep Dependency) *dependenciesManager {
	return NewDependenciesManager(lgr, casesMap, dep)
}

// registerManagerCleanup closes every cleanup loop created by a test. Tests
// that need to observe suite teardown can invoke the returned function early;
// sync.Once makes the registered test cleanup a no-op afterwards.
func registerManagerCleanup(t *testing.T, mgr *dependenciesManager) func() error {
	t.Helper()

	var once sync.Once
	var cleanupErr error
	cleanup := func() error {
		once.Do(func() {
			// Mirror casesDriver.TearDown: the suite cleanup snapshot is valid
			// only after all provisioned entities have reached the manager.
			select {
			case <-mgr.allProvisionDone:
			case <-time.After(2 * time.Second):
				cleanupErr = errors.New("timed out waiting for dependency provisioning")
				return
			}
			cleanupErr = mgr.Cleanup()
		})
		return cleanupErr
	}
	t.Cleanup(func() {
		require.NoError(t, cleanup())
	})
	return cleanup
}

func waitForCleanup(t *testing.T, mgr *dependenciesManager) {
	t.Helper()

	done := make(chan struct{})
	go func() {
		mgr.cleanupWG.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for case-scoped dependency cleanup")
	}
}

func requireDeleteStart(t *testing.T, dep *observableDependency, entity DependencyEntity) {
	t.Helper()
	select {
	case deleted := <-dep.deleteStart:
		require.Same(t, entity, deleted)
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for dependency deletion")
	}
}

func waitErr(ch <-chan error, d time.Duration) error {
	select {
	case err := <-ch:
		return err
	case <-time.After(d):
		return fmt.Errorf("timeout waiting for case readiness")
	}
}

func TestAsyncDeps_ReserveFirstEntityForShared(t *testing.T) {
	lgr := log.NewNoopLogger()
	stream := make(chan DependencyEntity)

	sharedCase := &mockCase{
		name: "sharedCase",
		reqs: map[string]DependencyRequirement{
			"tidb": makeSharedReq("us-east-1"),
		},
	}
	dedicatedCase := &mockCase{
		name: "dedicatedCase",
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	casesMap := map[string]cases.E2ECase{
		"sharedCase":    sharedCase,
		"dedicatedCase": dedicatedCase,
	}
	dep := &streamDependency{
		typ: TiDBDedicated,
		provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
			return stream, nil
		},
	}
	mgr := newTestMgr(lgr, casesMap, dep)
	registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	entity1 := makeEntity("us-east-1")
	entity2 := makeEntity("us-east-1")
	stream <- entity1
	assert.NoError(t, waitErr(mgr.caseDepResult["sharedCase"], 2*time.Second))
	select {
	case err := <-mgr.caseDepResult["dedicatedCase"]:
		assert.FailNow(t, "dedicated should not be ready after first(shared) entity", err)
	default:
	}
	stream <- entity2
	close(stream)

	assert.NoError(t, waitErr(mgr.caseDepResult["dedicatedCase"], 2*time.Second))
	assert.Same(t, entity1, sharedCase.TiDBDedicated("tidb"))
	assert.Same(t, entity2, dedicatedCase.TiDBDedicated("tidb"))
}

func TestAsyncDeps_IndependentByDimensions(t *testing.T) {
	lgr := log.NewNoopLogger()
	stream := make(chan DependencyEntity)

	eastCase := &mockCase{
		name: "eastCase",
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	westCase := &mockCase{
		name: "westCase",
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-west-2"),
		},
	}
	casesMap := map[string]cases.E2ECase{
		"eastCase": eastCase,
		"westCase": westCase,
	}
	dep := &streamDependency{
		typ: TiDBDedicated,
		provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
			return stream, nil
		},
	}
	mgr := newTestMgr(lgr, casesMap, dep)
	registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	eastEntity := makeEntity("us-east-1")
	westEntity := makeEntity("us-west-2")
	stream <- eastEntity
	assert.NoError(t, waitErr(mgr.caseDepResult["eastCase"], 2*time.Second))

	select {
	case err := <-mgr.caseDepResult["westCase"]:
		assert.FailNow(t, "west case should still be waiting", err)
	default:
	}

	stream <- westEntity
	close(stream)
	assert.NoError(t, waitErr(mgr.caseDepResult["westCase"], 2*time.Second))
	assert.Same(t, eastEntity, eastCase.TiDBDedicated("tidb"))
	assert.Same(t, westEntity, westCase.TiDBDedicated("tidb"))
}

func TestAsyncDeps_FailCaseWhenEntitiesNotEnough(t *testing.T) {
	lgr := log.NewNoopLogger()
	stream := make(chan DependencyEntity)

	sharedCase := &mockCase{
		name: "sharedCase",
		reqs: map[string]DependencyRequirement{
			"tidb": makeSharedReq("us-east-1"),
		},
	}
	dedicatedCase := &mockCase{
		name: "dedicatedCase",
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	casesMap := map[string]cases.E2ECase{
		"sharedCase":    sharedCase,
		"dedicatedCase": dedicatedCase,
	}
	dep := &streamDependency{
		typ: TiDBDedicated,
		provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
			return stream, nil
		},
	}
	mgr := newTestMgr(lgr, casesMap, dep)
	registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	stream <- makeEntity("us-east-1") // reserved for shared
	close(stream)                     // no dedicated entity arrives

	assert.NoError(t, waitErr(mgr.caseDepResult["sharedCase"], 2*time.Second))
	err := waitErr(mgr.caseDepResult["dedicatedCase"], 2*time.Second)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "unsatisfied dependency")
}

func TestDependenciesManagerCleanup_DeletesCreatedExclusiveEntityAfterAssignedCase(t *testing.T) {
	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	caseName := "exclusive"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(entity), nil
			},
		},
		deleteStart: make(chan DependencyEntity, 1),
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult[caseName], 2*time.Second))
	require.Same(t, entity, theCase.TiDBDedicated("tidb"))
	mgr.ReleaseCase(caseName)
	requireDeleteStart(t, dep, entity)
	waitForCleanup(t, mgr)

	require.NoError(t, cleanup())
	assert.Equal(t, 1, dep.deleteCount(entity))
}

func TestDependenciesManagerCleanup_DeletesCreatedSharedEntityAfterFinalAssignedCase(t *testing.T) {
	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	first := &mockCase{name: "first", reqs: map[string]DependencyRequirement{"tidb": makeSharedReq("us-east-1")}}
	second := &mockCase{name: "second", reqs: map[string]DependencyRequirement{"tidb": makeSharedReq("us-east-1")}}
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(entity), nil
			},
		},
		deleteStart: make(chan DependencyEntity, 1),
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{"first": first, "second": second}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult["first"], 2*time.Second))
	require.NoError(t, waitErr(mgr.caseDepResult["second"], 2*time.Second))
	require.Same(t, entity, first.TiDBDedicated("tidb"))
	require.Same(t, entity, second.TiDBDedicated("tidb"))

	mgr.ReleaseCase("first")
	waitForCleanup(t, mgr)
	assert.Zero(t, dep.deleteCount(entity))

	mgr.ReleaseCase("second")
	requireDeleteStart(t, dep, entity)
	waitForCleanup(t, mgr)
	require.NoError(t, cleanup())
	assert.Equal(t, 1, dep.deleteCount(entity))
}

func TestDependenciesManagerCleanup_NeverDeletesExistingEntity(t *testing.T) {
	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	entity.CreatedByDriver = false
	caseName := "existing"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	dep := &observableDependency{
		streamDependency: streamDependency{typ: TiDBDedicated, existed: []DependencyEntity{entity}},
		deleteStart:      make(chan DependencyEntity, 1),
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult[caseName], 2*time.Second))
	mgr.ReleaseCase(caseName)
	waitForCleanup(t, mgr)
	require.NoError(t, cleanup())
	assert.Zero(t, dep.deleteCount(entity))
}

func TestDependenciesManagerCleanup_RetriesCaseDeleteFailureAtSuiteTeardown(t *testing.T) {
	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	caseName := "retry"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	var attempts int
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(entity), nil
			},
		},
		deleteStart: make(chan DependencyEntity, 2),
		deleteFn: func(DependencyEntity) error {
			attempts++
			if attempts == 1 {
				return errors.New("case-scoped delete failed")
			}
			return nil
		},
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult[caseName], 2*time.Second))
	mgr.ReleaseCase(caseName)
	requireDeleteStart(t, dep, entity)
	waitForCleanup(t, mgr)
	require.NoError(t, cleanup())
	requireDeleteStart(t, dep, entity)
	assert.Equal(t, 2, dep.deleteCount(entity))
}

func TestDependenciesManagerCleanup_LogsRedactedDeleteErrorAndRetriesAtSuiteTeardown(t *testing.T) {
	lgr := &recordingLogger{}
	entity := makeEntity("us-east-1")
	caseName := "retry-with-error-log"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	var attempts int
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(entity), nil
			},
		},
		deleteFn: func(DependencyEntity) error {
			attempts++
			if attempts == 1 {
				return errors.New("delete rejected: token=top-secret request-id=req-123")
			}
			return nil
		},
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult[caseName], 2*time.Second))
	mgr.ReleaseCase(caseName)
	waitForCleanup(t, mgr)

	logMessage := lgr.lastError(t)
	assert.Contains(t, logMessage, "request-id=req-123")
	assert.Contains(t, logMessage, "token=<redacted>")
	assert.NotContains(t, logMessage, "top-secret")
	assert.Contains(t, logMessage, "retrying at suite teardown")
	require.NoError(t, cleanup())
	assert.Equal(t, 2, dep.deleteCount(entity))
}

func TestRedactDeleteError(t *testing.T) {
	tests := map[string]struct {
		input    string
		contains []string
		absent   []string
	}{
		"actual failed entity JSON": {
			input: `failed clusters: [{"ClusterID":"cluster-1","Password":"password-secret","User":"test-user","CACert":"certificate-data"}], request-id=req-123`,
			contains: []string{
				`"ClusterID":"cluster-1"`,
				"request-id=req-123",
				`"Password":<redacted>`,
				`"User":<redacted>`,
				`"CACert":<redacted>`,
			},
			absent: []string{"password-secret", "test-user", "certificate-data"},
		},
		"quoted and unquoted key values": {
			input:    `password="password-secret", token=token-secret; API_KEY: 'api-secret'; request-id=req-456`,
			contains: []string{"password=<redacted>", "token=<redacted>", "API_KEY: <redacted>", "request-id=req-456"},
			absent:   []string{"password-secret", "token-secret", "api-secret"},
		},
		"compound secret fields": {
			input: `{"SecretKey":"secret-key-value","secret_key":"secret-key-snake-value","clientSecret":"client-secret-value","client_secret":"client-secret-snake-value","SecretAccessKey":"secret-access-key-value","secret-access-key":"kebab-secret-value","sasToken":"sas-token-value"}, request-id=req-567`,
			contains: []string{
				`"SecretKey":<redacted>`,
				`"secret_key":<redacted>`,
				`"clientSecret":<redacted>`,
				`"client_secret":<redacted>`,
				`"SecretAccessKey":<redacted>`,
				`"secret-access-key":<redacted>`,
				`"sasToken":<redacted>`,
				"request-id=req-567",
			},
			absent: []string{
				"secret-key-value",
				"secret-key-snake-value",
				"client-secret-value",
				"client-secret-snake-value",
				"secret-access-key-value",
				"kebab-secret-value",
				"sas-token-value",
			},
		},
		"bearer and basic credentials": {
			input:    `authorization: Bearer bearer-secret, proxy=Basic basic-secret, request-id=req-789`,
			contains: []string{"authorization: <redacted>", "Basic <redacted>", "request-id=req-789"},
			absent:   []string{"bearer-secret", "basic-secret"},
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			redacted := redactDeleteError(errors.New(tc.input))
			for _, expected := range tc.contains {
				assert.Contains(t, redacted, expected)
			}
			for _, secret := range tc.absent {
				assert.NotContains(t, redacted, secret)
			}
		})
	}
}

func TestDependenciesManagerCleanup_DeletesUnassignedCreatedEntityAtSuiteTeardown(t *testing.T) {
	lgr := log.NewNoopLogger()
	assigned := makeEntity("us-east-1")
	unassigned := makeEntity("us-east-1")
	caseName := "assigned"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(assigned, unassigned), nil
			},
		},
		deleteStart: make(chan DependencyEntity, 2),
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult[caseName], 2*time.Second))
	mgr.ReleaseCase(caseName)
	requireDeleteStart(t, dep, assigned)
	waitForCleanup(t, mgr)
	assert.Zero(t, dep.deleteCount(unassigned))

	require.NoError(t, cleanup())
	requireDeleteStart(t, dep, unassigned)
	assert.Equal(t, 1, dep.deleteCount(assigned))
	assert.Equal(t, 1, dep.deleteCount(unassigned))
}

func TestDependenciesManagerCleanup_WaitsForInFlightCaseDeleteBeforeSuiteTeardown(t *testing.T) {
	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	caseName := "blocking-delete"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	deleteContinue := make(chan struct{})
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(entity), nil
			},
		},
		deleteStart: make(chan DependencyEntity, 1),
		deleteFn: func(DependencyEntity) error {
			<-deleteContinue
			return nil
		},
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	var unblockDelete sync.Once
	unblock := func() { unblockDelete.Do(func() { close(deleteContinue) }) }
	// Cleanup callbacks run LIFO. Register the unblocker after manager cleanup
	// so a failed assertion cannot strand the manager cleanup waiting on Delete.
	t.Cleanup(unblock)
	mgr.ProvisionAsync(context.Background())

	require.NoError(t, waitErr(mgr.caseDepResult[caseName], 2*time.Second))
	mgr.ReleaseCase(caseName)
	requireDeleteStart(t, dep, entity)

	cleanupDone := make(chan error, 1)
	go func() { cleanupDone <- cleanup() }()
	select {
	case err := <-cleanupDone:
		t.Fatalf("suite cleanup returned before in-flight delete finished: %v", err)
	case <-time.After(100 * time.Millisecond):
	}
	unblock()

	select {
	case err := <-cleanupDone:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for suite cleanup")
	}
	assert.Equal(t, 1, dep.deleteCount(entity))
}

func TestDependenciesManagerCleanup_RecordsSharedAssignmentsBeforeCaseReadiness(t *testing.T) {
	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	secondFulfilled := make(chan struct{})
	allowSecondFulfillment := make(chan struct{})
	var allowSecondFulfillmentOnce sync.Once
	allowFulfillment := func() { allowSecondFulfillmentOnce.Do(func() { close(allowSecondFulfillment) }) }
	fulfillmentCount := 0
	fulfillBarrier := func(string, DependencyEntity) {
		fulfillmentCount++
		if fulfillmentCount == 2 {
			close(secondFulfilled)
			<-allowSecondFulfillment
		}
	}
	first := &mockCase{
		name:      "first",
		reqs:      map[string]DependencyRequirement{"tidb": makeSharedReq("us-east-1")},
		onFulfill: fulfillBarrier,
	}
	second := &mockCase{
		name:      "second",
		reqs:      map[string]DependencyRequirement{"tidb": makeSharedReq("us-east-1")},
		onFulfill: fulfillBarrier,
	}
	stream := make(chan DependencyEntity, 1)
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return stream, nil
			},
		},
		deleteStart: make(chan DependencyEntity, 1),
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{"first": first, "second": second}, dep)
	cleanup := registerManagerCleanup(t, mgr)
	// The dispatcher goroutine is intentionally paused inside the second actual
	// case fulfillment. Unblock it before manager cleanup if this test fails.
	t.Cleanup(allowFulfillment)
	mgr.ProvisionAsync(context.Background())
	stream <- entity
	select {
	case <-secondFulfilled:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for second shared fulfillment")
	}
	for caseName, resultCh := range map[string]<-chan error{
		"first":  mgr.caseDepResult["first"],
		"second": mgr.caseDepResult["second"],
	} {
		select {
		case <-resultCh:
			allowFulfillment()
			t.Fatalf("case %s became ready before the assignment batch completed", caseName)
		default:
		}
	}
	allowFulfillment()
	close(stream)
	select {
	case <-mgr.allProvisionDone:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out provisioning shared dependency")
	}
	require.NoError(t, waitErr(mgr.caseDepResult["first"], 2*time.Second))
	require.NoError(t, waitErr(mgr.caseDepResult["second"], 2*time.Second))

	mgr.ReleaseCase("first")
	waitForCleanup(t, mgr)
	assert.Zero(t, dep.deleteCount(entity), "first release must not delete the shared entity")
	mgr.ReleaseCase("second")
	requireDeleteStart(t, dep, entity)
	waitForCleanup(t, mgr)
	require.NoError(t, cleanup())
	assert.Equal(t, 1, dep.deleteCount(entity))
}

func TestDependenciesManagerCleanup_DeletesPartialAssignmentAfterDependencyFailure(t *testing.T) {
	const (
		lateType = "late"
		failType = "fail"
	)
	lgr := log.NewNoopLogger()
	assignedEntity := &dummyEntity{typ: lateType, dimensions: map[string]string{"region": "us-east-1"}, createdByDriver: true}
	unassignedEntity := &dummyEntity{typ: lateType, dimensions: map[string]string{"region": "us-east-1"}, createdByDriver: true}
	theCase := &mockCase{
		name:       "partial-failure",
		assignedCh: make(chan string, 1),
		reqs: map[string]DependencyRequirement{
			"late": &dummyRequirement{typ: lateType, dimensions: assignedEntity.Dimensions(), dedicated: true},
			"fail": &dummyRequirement{typ: failType, dimensions: map[string]string{"region": "us-east-1"}, dedicated: true},
		},
	}
	lateStream := make(chan DependencyEntity, 1)
	lateStarted := make(chan struct{})
	failureStarted := make(chan struct{})
	allowFailure := make(chan struct{})
	lateDep := &observableDependency{
		streamDependency: streamDependency{
			typ: lateType,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				close(lateStarted)
				return lateStream, nil
			},
		},
		deleteStart: make(chan DependencyEntity, 2),
	}
	failingDep := &observableDependency{
		streamDependency: streamDependency{
			typ: failType,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				close(failureStarted)
				<-allowFailure
				return nil, errors.New("simulated dependency failure")
			},
		},
	}
	mgr := NewDependenciesManager(lgr, map[string]cases.E2ECase{theCase.Name(): theCase}, lateDep, failingDep)
	cleanup := registerManagerCleanup(t, mgr)
	mgr.ProvisionAsync(context.Background())

	for _, started := range []<-chan struct{}{lateStarted, failureStarted} {
		select {
		case <-started:
		case <-time.After(2 * time.Second):
			t.Fatal("timed out starting dependency provisioning")
		}
	}
	lateStream <- assignedEntity
	select {
	case reqName := <-theCase.assignedCh:
		require.Equal(t, "late", reqName)
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for partial dependency assignment")
	}

	close(allowFailure)
	err := waitErr(mgr.caseDepResult[theCase.Name()], 2*time.Second)
	require.Error(t, err)
	mgr.ReleaseCase(theCase.Name())
	requireDeleteStart(t, lateDep, assignedEntity)
	waitForCleanup(t, mgr)

	// The additional entity arrives after the case has been closed by the
	// dependency error, so it has no actual user and remains suite-scoped.
	lateStream <- unassignedEntity
	close(lateStream)
	select {
	case <-mgr.allProvisionDone:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for late dependency provisioning")
	}

	// A case may release after one dependency is assigned but another fails.
	// The partial assignment must still be case-cleaned; post-error entities
	// are intentionally unassigned and are left for suite fallback cleanup.
	assert.Equal(t, 1, lateDep.deleteCount(assignedEntity))
	assert.Zero(t, lateDep.deleteCount(unassignedEntity))
	require.NoError(t, cleanup())
	requireDeleteStart(t, lateDep, unassignedEntity)
	assert.Equal(t, 1, lateDep.deleteCount(unassignedEntity))
}

func TestCasesDriver_RunSkipsCaseCleanupWhenMustCleanupIsFalse(t *testing.T) {
	// config.InitConfig reads config.common.yaml relative to the E2E module
	// root. Supplying a token avoids its console-login path, so this test stays
	// entirely local while exercising the same GlobalConfig guard as Run.
	t.Chdir("../..")
	configPath := t.TempDir() + "/must-cleanup-false.yaml"
	require.NoError(t, os.WriteFile(configPath, []byte("mustCleanup: false\ntidbCloudToken: dummy-token\n"), 0o600))
	previousConfig := config.GlobalConfig
	t.Cleanup(func() { config.GlobalConfig = previousConfig })
	config.InitConfig(config.LocalENV, "", "", configPath, "", "", "", config.SuitePlanDedicated)

	lgr := log.NewNoopLogger()
	entity := makeEntity("us-east-1")
	caseName := "must-cleanup-disabled"
	theCase := &mockCase{
		name: caseName,
		reqs: map[string]DependencyRequirement{
			"tidb": makeDedicatedReq("us-east-1"),
		},
	}
	dep := &observableDependency{
		streamDependency: streamDependency{
			typ: TiDBDedicated,
			provisionFn: func(context.Context, EntitiesStats) (<-chan DependencyEntity, error) {
				return closedEntities(entity), nil
			},
		},
		deleteStart: make(chan DependencyEntity, 1),
	}
	mgr := newTestMgr(lgr, map[string]cases.E2ECase{caseName: theCase}, dep)
	driver := &casesDriver{
		lgr:                 lgr,
		dependenciesManager: mgr,
		cases2run:           map[string]cases.E2ECase{caseName: theCase},
	}
	var teardownOnce sync.Once
	teardown := func() { teardownOnce.Do(driver.TearDown) }
	t.Cleanup(teardown)

	driver.SetUp()
	t.Run("driver", driver.Run)
	assert.Zero(t, dep.deleteCount(entity), "mustCleanup=false must skip per-case release")

	teardown()
	requireDeleteStart(t, dep, entity)
	assert.Equal(t, 1, dep.deleteCount(entity), "suite teardown must still clean driver-created dependencies")
}

func closedEntities(entities ...DependencyEntity) <-chan DependencyEntity {
	ch := make(chan DependencyEntity, len(entities))
	for _, entity := range entities {
		ch <- entity
	}
	close(ch)
	return ch
}