package tidbmanager

import (
	"context"
	"errors"
	"testing"

	"github.com/go-logr/logr"
	"github.com/tidbcloud/aws-shared-provider/pkg/keyspacemigration"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func TestAssignTidb_BlockedByMigrationLock(t *testing.T) {
	migrationLocker := keyspacemigration.New()
	release, ok := migrationLocker.TryAcquire("ks-a")
	if !ok {
		t.Fatal("setup: failed to acquire migration lock")
	}
	defer release()

	mgr := &tidbManager{
		Logger:          logr.New(log.NullLogSink{}),
		ctx:             context.Background(),
		migrationLocker: migrationLocker,
		assignCache:     &assignCache{},
		freeCache:       &freeCache{},
	}

	_, err := mgr.AssignTidb(context.Background(), "ks-a", &tidb.ActivateParameter{})
	if !errors.Is(err, utils.ErrKeyspaceMetaServiceGroupMigrating) {
		t.Fatalf("expected ErrKeyspaceMetaServiceGroupMigrating, got %v", err)
	}
}

// Token suffixes (export id, resource group) must still match the lock,
// which is keyed by the trimmed keyspace name.
func TestAssignTidb_BlockedByMigrationLock_TrimmedToken(t *testing.T) {
	migrationLocker := keyspacemigration.New()
	release, ok := migrationLocker.TryAcquire("ks-a")
	if !ok {
		t.Fatal("setup: failed to acquire migration lock")
	}
	defer release()

	mgr := &tidbManager{
		Logger:          logr.New(log.NullLogSink{}),
		ctx:             context.Background(),
		migrationLocker: migrationLocker,
		assignCache:     &assignCache{},
		freeCache:       &freeCache{},
	}

	// Resource-group-suffixed token; TrimToken should reduce it to "ks-a".
	_, err := mgr.AssignTidb(context.Background(), "ks-a/rg1", &tidb.ActivateParameter{})
	if !errors.Is(err, utils.ErrKeyspaceMetaServiceGroupMigrating) {
		t.Fatalf("expected ErrKeyspaceMetaServiceGroupMigrating for trimmed token, got %v", err)
	}
}

func TestAssignTidb_NoMigrationLockFallsThrough(t *testing.T) {
	migrationLocker := keyspacemigration.New()

	mgr := &tidbManager{
		Logger:          logr.New(log.NullLogSink{}),
		ctx:             context.Background(),
		migrationLocker: migrationLocker,
		assignCache:     &assignCache{},
		freeCache:       &freeCache{},
	}

	_, err := mgr.AssignTidb(context.Background(), "ks-fresh", &tidb.ActivateParameter{})
	if errors.Is(err, utils.ErrKeyspaceMetaServiceGroupMigrating) {
		t.Fatalf("did not expect ErrKeyspaceMetaServiceGroupMigrating; got %v", err)
	}
	if !errors.Is(err, utils.ErrNotFoundToken) {
		t.Fatalf("expected ErrNotFoundToken (token never assigned), got %v", err)
	}
}