package server

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	bractionmocks "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br/mocks"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor"
	executormocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/tioperator"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// Mock transactioner for testing
type mockCreateTiKVGroupTransactioner struct{}

func (m *mockCreateTiKVGroupTransactioner) Transaction(ctx context.Context, fn func(context.Context) error) error {
	return fn(ctx)
}

func TestValidateTikvGroupId(t *testing.T) {
	testCases := []struct {
		name          string
		tikvGroupId   string
		expectError   bool
		expectedError string
	}{
		{
			name:        "valid lowercase alphanumeric",
			tikvGroupId: "tikv1",
			expectError: false,
		},
		{
			name:        "valid all lowercase",
			tikvGroupId: "mytikvgroup",
			expectError: false,
		},
		{
			name:        "valid all numbers",
			tikvGroupId: "12345",
			expectError: false,
		},
		{
			name:          "empty string",
			tikvGroupId:   "",
			expectError:   true,
			expectedError: "tikv_group_id cannot be empty",
		},
		{
			name:          "contains hyphen",
			tikvGroupId:   "tikv-group",
			expectError:   true,
			expectedError: "tikv_group_id 'tikv-group' is invalid",
		},
		{
			name:          "contains underscore",
			tikvGroupId:   "tikv_group",
			expectError:   true,
			expectedError: "tikv_group_id 'tikv_group' is invalid",
		},
		{
			name:          "contains uppercase",
			tikvGroupId:   "TiKVGroup",
			expectError:   true,
			expectedError: "tikv_group_id 'TiKVGroup' is invalid",
		},
		{
			name:          "contains special characters",
			tikvGroupId:   "tikv@group",
			expectError:   true,
			expectedError: "tikv_group_id 'tikv@group' is invalid",
		},
		{
			name:          "contains space",
			tikvGroupId:   "tikv group",
			expectError:   true,
			expectedError: "tikv_group_id 'tikv group' is invalid",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateTikvGroupId(tc.tikvGroupId)
			if tc.expectError {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tc.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestValidateTiKVGroupNodeQuota(t *testing.T) {
	testCases := []struct {
		name          string
		nodeQuota     *entity.NodeQuota
		expectError   bool
		expectedError string
	}{
		{
			name:          "nil node quota",
			nodeQuota:     nil,
			expectError:   true,
			expectedError: "node_quota is required",
		},
		{
			name: "nil node spec",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec:  nil,
				NodeCount: 3,
			},
			expectError:   true,
			expectedError: "node_quota.node_spec is required",
		},
		{
			name: "zero cpu",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      0,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			expectError:   true,
			expectedError: "node_quota.node_spec.cpu is required",
		},
		{
			name: "zero memory",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 0,
				},
				NodeCount: 3,
			},
			expectError:   true,
			expectedError: "node_quota.node_spec.memory_gi is required",
		},
		{
			name: "zero node count",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 0,
			},
			expectError:   true,
			expectedError: "node_quota.node_count is required",
		},
		{
			name: "node count not multiple of 3",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 5,
			},
			expectError:   true,
			expectedError: "node_quota.node_count must be multiple of 3",
		},
		{
			name: "node count is 1",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 1,
			},
			expectError:   true,
			expectedError: "node_quota.node_count must be multiple of 3",
		},
		{
			name: "valid node quota with 3 nodes",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			expectError: false,
		},
		{
			name: "valid node quota with 6 nodes",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 6,
			},
			expectError: false,
		},
		{
			name: "valid node quota with 9 nodes",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      16,
					MemoryGi: 32,
				},
				NodeCount: 9,
			},
			expectError: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateTiKVGroupNodeQuota(tc.nodeQuota)
			if tc.expectError {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tc.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestValidateTiKVGroupStorages(t *testing.T) {
	testCases := []struct {
		name          string
		storages      []*entity.Storage
		provider      string
		expectError   bool
		expectedError string
	}{
		{
			name:          "empty storages",
			storages:      []*entity.Storage{},
			provider:      "aws",
			expectError:   true,
			expectedError: "storages are required",
		},
		{
			name:          "nil storages",
			storages:      nil,
			provider:      "aws",
			expectError:   true,
			expectedError: "storages are required",
		},
		{
			name: "missing DATA storage",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "DATA storage is required",
		},
		{
			name: "missing RAFTLOG storage",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.GP3,
					SizeGi:        500,
					Iops:          10000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "RAFTLOG storage is required",
		},
		{
			name: "invalid storage name",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.GP3,
					SizeGi:        500,
					Iops:          10000,
					ThroughputMbs: 500,
				},
				{
					Name:   apiv1.Component_Storage_SLOWLOG,
					Type:   entity.GP3,
					SizeGi: 100,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "invalid storage name",
		},
		{
			name: "AWS gp3 missing size",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.GP3,
					SizeGi:        0,
					Iops:          10000,
					ThroughputMbs: 500,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "gp3 storage size must be greater than zero",
		},
		{
			name: "AWS gp3 missing iops",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.GP3,
					SizeGi:        500,
					Iops:          0,
					ThroughputMbs: 500,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "gp3 storage iops must be greater than zero",
		},
		{
			name: "AWS gp3 missing throughput",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.GP3,
					SizeGi:        500,
					Iops:          10000,
					ThroughputMbs: 0,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "gp3 storage throughput must be greater than zero",
		},
		{
			name: "AWS io2 with throughput",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.IO2,
					SizeGi:        500,
					Iops:          10000,
					ThroughputMbs: 500,
				},
				{
					Name:   apiv1.Component_Storage_RAFTLOG,
					Type:   entity.IO2,
					SizeGi: 100,
					Iops:   5000,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "io2 storage does not support setting throughput",
		},
		{
			name: "AWS invalid storage type",
			storages: []*entity.Storage{
				{
					Name:   apiv1.Component_Storage_DATA,
					Type:   "invalid-type",
					SizeGi: 500,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "invalid storage type 'invalid-type' for AWS",
		},
		{
			name: "AWS local-ssd with iops",
			storages: []*entity.Storage{
				{
					Name: apiv1.Component_Storage_DATA,
					Type: entity.LocalSSD,
					Iops: 10000,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "local-ssd storage does not support setting iops",
		},
		{
			name: "AliCloud cloud-auto missing size",
			storages: []*entity.Storage{
				{
					Name:   apiv1.Component_Storage_DATA,
					Type:   entity.CloudAuto,
					SizeGi: 0,
				},
				{
					Name:   apiv1.Component_Storage_RAFTLOG,
					Type:   entity.CloudAuto,
					SizeGi: 100,
				},
			},
			provider:      "alicloud",
			expectError:   true,
			expectedError: "cloud-auto storage size must be greater than zero",
		},
		{
			name: "AliCloud cloud-auto with throughput",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.CloudAuto,
					SizeGi:        500,
					ThroughputMbs: 100,
				},
				{
					Name:   apiv1.Component_Storage_RAFTLOG,
					Type:   entity.CloudAuto,
					SizeGi: 100,
				},
			},
			provider:      "alicloud",
			expectError:   true,
			expectedError: "cloud-auto storage does not support setting throughput",
		},
		{
			name: "valid AWS gp3 storages",
			storages: []*entity.Storage{
				{
					Name:          apiv1.Component_Storage_DATA,
					Type:          entity.GP3,
					SizeGi:        500,
					Iops:          10000,
					ThroughputMbs: 500,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:    "aws",
			expectError: false,
		},
		{
			name: "valid AWS io2 storages",
			storages: []*entity.Storage{
				{
					Name:   apiv1.Component_Storage_DATA,
					Type:   entity.IO2,
					SizeGi: 500,
					Iops:   10000,
				},
				{
					Name:   apiv1.Component_Storage_RAFTLOG,
					Type:   entity.IO2,
					SizeGi: 100,
					Iops:   5000,
				},
			},
			provider:    "aws",
			expectError: false,
		},
		{
			name: "valid AWS local-ssd DATA with gp3 RAFTLOG",
			storages: []*entity.Storage{
				{
					Name: apiv1.Component_Storage_DATA,
					Type: entity.LocalSSD,
				},
				{
					Name:          apiv1.Component_Storage_RAFTLOG,
					Type:          entity.GP3,
					SizeGi:        100,
					Iops:          5000,
					ThroughputMbs: 500,
				},
			},
			provider:    "aws",
			expectError: false,
		},
		{
			name: "RAFTLOG cannot use local-ssd on AWS",
			storages: []*entity.Storage{
				{
					Name: apiv1.Component_Storage_DATA,
					Type: entity.LocalSSD,
				},
				{
					Name: apiv1.Component_Storage_RAFTLOG,
					Type: entity.LocalSSD,
				},
			},
			provider:      "aws",
			expectError:   true,
			expectedError: "RAFTLOG storage cannot use local-ssd type",
		},
		{
			name: "RAFTLOG cannot use local-ssd on AliCloud",
			storages: []*entity.Storage{
				{
					Name: apiv1.Component_Storage_DATA,
					Type: entity.LocalSSD,
				},
				{
					Name: apiv1.Component_Storage_RAFTLOG,
					Type: entity.LocalSSD,
				},
			},
			provider:      "alicloud",
			expectError:   true,
			expectedError: "RAFTLOG storage cannot use local-ssd type",
		},
		{
			name: "valid AliCloud cloud-auto storages",
			storages: []*entity.Storage{
				{
					Name:   apiv1.Component_Storage_DATA,
					Type:   entity.CloudAuto,
					SizeGi: 500,
				},
				{
					Name:   apiv1.Component_Storage_RAFTLOG,
					Type:   entity.CloudAuto,
					SizeGi: 100,
				},
			},
			provider:    "alicloud",
			expectError: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateTiKVGroupStorages(tc.storages, tc.provider)
			if tc.expectError {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tc.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestValidateTiKVGroupLabels(t *testing.T) {
	testCases := []struct {
		name          string
		labels        map[string]string
		expectError   bool
		expectedError string
	}{
		{
			name:        "nil labels",
			labels:      nil,
			expectError: false,
		},
		{
			name:        "empty labels",
			labels:      map[string]string{},
			expectError: false,
		},
		{
			name: "valid labels",
			labels: map[string]string{
				"role": "hot",
				"az":   "az1",
			},
			expectError: false,
		},
		{
			name: "reserved label key - keyspace",
			labels: map[string]string{
				tioperator.LabelKeyKeyspace: "my-keyspace",
			},
			expectError:   true,
			expectedError: "label key 'keyspace' is reserved",
		},
		{
			name: "reserved label key - cidr",
			labels: map[string]string{
				tioperator.LabelKeyCidr: "10.0.0.0/8",
			},
			expectError:   true,
			expectedError: "label key 'cidr' is reserved",
		},
		{
			name: "reserved label key - host",
			labels: map[string]string{
				"host": "myhost",
			},
			expectError:   true,
			expectedError: "label key 'host' is reserved",
		},
		{
			name: "reserved label key - region",
			labels: map[string]string{
				"region": "us-west-2",
			},
			expectError:   true,
			expectedError: "label key 'region' is reserved",
		},
		{
			name: "reserved label key - zone",
			labels: map[string]string{
				"zone": "us-west-2a",
			},
			expectError:   true,
			expectedError: "label key 'zone' is reserved",
		},
		{
			name: "too many labels",
			labels: func() map[string]string {
				labels := make(map[string]string)
				for i := 0; i <= MaxLabelsCount; i++ {
					labels[string(rune('a'+i))] = "value"
				}
				return labels
			}(),
			expectError:   true,
			expectedError: "too many labels",
		},
		{
			name: "invalid label key with uppercase",
			labels: map[string]string{
				"MyKey": "value",
			},
			expectError:   true,
			expectedError: "invalid label key",
		},
		{
			name: "empty label key",
			labels: map[string]string{
				"": "value",
			},
			expectError:   true,
			expectedError: "label key cannot be empty",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateTiKVGroupLabels(tc.labels)
			if tc.expectError {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tc.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestCheckSharedPoolCanAcceptNewComponent(t *testing.T) {
	testCases := []struct {
		name          string
		sharedPool    *entity.SharedPool
		expectError   bool
		expectedCode  codes.Code
		expectedError string
	}{
		{
			name: "available shared pool",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Phase: apiv1.SharedPool_AVAILABLE,
				},
			},
			expectError: false,
		},
		{
			name: "modifying shared pool",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Phase: apiv1.SharedPool_MODIFYING,
				},
			},
			expectError: false,
		},
		{
			name: "creating shared pool",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Phase: apiv1.SharedPool_CREATING,
				},
			},
			expectError:   true,
			expectedCode:  codes.FailedPrecondition,
			expectedError: "shared pool is creating",
		},
		{
			name: "pausing shared pool",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Phase: apiv1.SharedPool_PAUSING,
				},
			},
			expectError:   true,
			expectedCode:  codes.FailedPrecondition,
			expectedError: "shared pool is pausing",
		},
		{
			name: "unavailable shared pool",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Phase: apiv1.SharedPool_UNAVAILABLE,
				},
			},
			expectError:   true,
			expectedCode:  codes.FailedPrecondition,
			expectedError: "shared pool is unavailable",
		},
		{
			name: "deleting shared pool",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					Phase: apiv1.SharedPool_DELETING,
				},
			},
			expectError:   true,
			expectedCode:  codes.FailedPrecondition,
			expectedError: "shared pool is being deleted",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := checkSharedPoolCanAcceptNewComponent(tc.sharedPool)
			if tc.expectError {
				require.Error(t, err)
				st, ok := status.FromError(err)
				require.True(t, ok)
				assert.Equal(t, tc.expectedCode, st.Code())
				assert.Contains(t, st.Message(), tc.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func TestCheckTikvGroupIdUniqueness(t *testing.T) {
	testCases := []struct {
		name          string
		sharedPool    *entity.SharedPool
		tikvGroupId   string
		expectError   bool
		expectedError string
	}{
		{
			name: "unique tikv group id",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					ComponentGroups: []*apiv1.Component{
						{
							Id:            "existingtikv",
							ComponentType: apiv1.ComponentType_TIKV,
						},
					},
				},
			},
			tikvGroupId: "new-tikv",
			expectError: false,
		},
		{
			name: "duplicate tikv group id",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					ComponentGroups: []*apiv1.Component{
						{
							Id:            "existingtikv",
							ComponentType: apiv1.ComponentType_TIKV,
						},
					},
				},
			},
			tikvGroupId:   "existingtikv",
			expectError:   true,
			expectedError: "tikv group with id 'existingtikv' already exists",
		},
		{
			name: "same id but different component type",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					ComponentGroups: []*apiv1.Component{
						{
							Id:            "sameid",
							ComponentType: apiv1.ComponentType_TIDB,
						},
					},
				},
			},
			tikvGroupId: "sameid",
			expectError: false,
		},
		{
			name: "empty component groups",
			sharedPool: &entity.SharedPool{
				SharedPool: apiv1.SharedPool{
					ComponentGroups: []*apiv1.Component{},
				},
			},
			tikvGroupId: "new-tikv",
			expectError: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := checkTikvGroupIdUniqueness(tc.sharedPool, tc.tikvGroupId)
			if tc.expectError {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tc.expectedError)
			} else {
				require.NoError(t, err)
			}
		})
	}
}

func setupSuccessfulCreateTiKVGroup(phase apiv1.SharedPool_Phase) func(
	*repomocks.SharedPoolRepo,
	*repomocks.LockRepo,
	*dataplanemgrmocks.Client,
	*executormocks.RemoteExecutor,
	*bractionmocks.ScaleManager,
) {
	return func(
		sharedPoolRepo *repomocks.SharedPoolRepo,
		lockRepo *repomocks.LockRepo,
		dataplaneClient *dataplanemgrmocks.Client,
		dagExecutor *executormocks.RemoteExecutor,
		scaleManager *bractionmocks.ScaleManager,
	) {
		lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
			return lock.Key == "testpool"
		})).Return(func() error { return nil }, nil)

		sharedPoolRepo.On("Load", mock.Anything, "testpool", false).Return(&entity.SharedPool{
			SharedPool: apiv1.SharedPool{
				Id:            "testpool",
				Phase:         phase,
				DataplaneName: "dp-test",
				ComponentGroups: []*apiv1.Component{{
					Id:            "default",
					ComponentType: apiv1.ComponentType_TIKV,
					Phase:         apiv1.Component_AVAILABLE,
				}},
			},
		}, nil)

		dataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{
			NodePools: []*dplanemgrv1.NodePool{{
				CpuCore:               8,
				MemoryGi:              16,
				AvailableCpuMilliCore: 7800,
				AvailableMemoryMi:     15600,
				NodeType:              dplanemgrv1.NodeType_Dedicated,
			}},
		}, nil)

		dagExecutor.On("QueryRunnings", mock.Anything, mock.Anything, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil)
		scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil)
		sharedPoolRepo.On("Save", mock.Anything, mock.MatchedBy(func(sp *entity.SharedPool) bool {
			return sp.GetId() == "testpool" && sp.GetPhase() == apiv1.SharedPool_MODIFYING
		})).Return(nil)
		dagExecutor.On("Submit", mock.Anything, dag.PlanCreateComponent, mock.MatchedBy(func(args actions.ActionArgs) bool {
			exclusive, err := args.GetBool(workflow.TiKVPlacementExclusiveKey)
			return err == nil && exclusive
		}), mock.Anything).Return(nil)
	}
}

func TestCreateTiKVGroup(t *testing.T) {
	testCases := []struct {
		name         string
		sharedPoolId string
		tikvGroupId  string
		nodeQuota    *entity.NodeQuota
		storages     []*entity.Storage
		labels       map[string]string
		exclusive    bool
		force        bool
		setupMocks   func(
			*repomocks.SharedPoolRepo,
			*repomocks.LockRepo,
			*dataplanemgrmocks.Client,
			*executormocks.RemoteExecutor,
			*bractionmocks.ScaleManager,
		)
		expectError   bool
		expectedCode  codes.Code
		expectedError string
	}{
		{
			name:         "empty shared pool id",
			sharedPoolId: "",
			tikvGroupId:  "tikv1",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
			},
			setupMocks: func(*repomocks.SharedPoolRepo, *repomocks.LockRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) {
			},
			expectError:   true,
			expectedCode:  codes.InvalidArgument,
			expectedError: "shared_pool_id is required",
		},
		{
			name:         "empty tikv group id",
			sharedPoolId: "testpool",
			tikvGroupId:  "",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
			},
			setupMocks: func(*repomocks.SharedPoolRepo, *repomocks.LockRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) {
			},
			expectError:   true,
			expectedCode:  codes.InvalidArgument,
			expectedError: "tikv_group_id is required",
		},
		{
			name:         "invalid tikv group id with hyphen",
			sharedPoolId: "testpool",
			tikvGroupId:  "tikv-group",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
			},
			setupMocks: func(*repomocks.SharedPoolRepo, *repomocks.LockRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) {
			},
			expectError:   true,
			expectedCode:  codes.InvalidArgument,
			expectedError: "tikv_group_id 'tikv-group' is invalid",
		},
		{
			name:         "node count not multiple of 3",
			sharedPoolId: "testpool",
			tikvGroupId:  "tikv1",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 5,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
			},
			setupMocks: func(*repomocks.SharedPoolRepo, *repomocks.LockRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) {
			},
			expectError:   true,
			expectedCode:  codes.InvalidArgument,
			expectedError: "node_quota.node_count must be multiple of 3",
		},
		{
			name:         "missing DATA storage",
			sharedPoolId: "testpool",
			tikvGroupId:  "tikv1",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
			},
			setupMocks: func(*repomocks.SharedPoolRepo, *repomocks.LockRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) {
			},
			expectError:   true,
			expectedCode:  codes.InvalidArgument,
			expectedError: "DATA storage is required",
		},
		{
			name:         "reserved label key",
			sharedPoolId: "testpool",
			tikvGroupId:  "tikv1",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 3,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
			},
			labels: map[string]string{
				"zone": "az1",
			},
			setupMocks: func(*repomocks.SharedPoolRepo, *repomocks.LockRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) {
			},
			expectError:   true,
			expectedCode:  codes.InvalidArgument,
			expectedError: "label key 'zone' is reserved",
		},
		{
			name:         "successful create tikv group",
			sharedPoolId: "testpool",
			tikvGroupId:  "tikv2",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec:  &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
				NodeCount: 6,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.IO2, SizeGi: 100, Iops: 5000},
			},
			labels:      map[string]string{"role": "hot"},
			exclusive:   true,
			setupMocks:  setupSuccessfulCreateTiKVGroup(apiv1.SharedPool_AVAILABLE),
			expectError: false,
		},
		{
			name:         "force creates tikv group when shared pool is unavailable",
			sharedPoolId: "testpool",
			tikvGroupId:  "tikv2",
			nodeQuota: &apiv1.NodeQuota{
				NodeSpec: &apiv1.NodeSpec{
					Cpu:      8,
					MemoryGi: 16,
				},
				NodeCount: 6,
			},
			storages: []*entity.Storage{
				{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
				{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.IO2, SizeGi: 100, Iops: 5000},
			},
			labels: map[string]string{
				"role": "hot",
			},
			exclusive:   true,
			force:       true,
			setupMocks:  setupSuccessfulCreateTiKVGroup(apiv1.SharedPool_UNAVAILABLE),
			expectError: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Setup mocks
			sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
			lockRepo := repomocks.NewLockRepo(t)
			dataplaneClient := dataplanemgrmocks.NewClient(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			scaleManager := bractionmocks.NewScaleManager(t)

			tc.setupMocks(sharedPoolRepo, lockRepo, dataplaneClient, dagExecutor, scaleManager)

			// Create server
			server := &serverImpl{
				cloudProvider:   "aws",
				sharedPoolRepo:  sharedPoolRepo,
				lockRepo:        lockRepo,
				dataplaneClient: dataplaneClient,
				dagExecutor:     dagExecutor,
				scaleManager:    scaleManager,
				transactioner:   &mockCreateTiKVGroupTransactioner{},
			}

			// Execute
			result, err := server.CreateTiKVGroup(
				context.Background(),
				tc.sharedPoolId,
				tc.tikvGroupId,
				tc.nodeQuota,
				tc.storages,
				tc.labels,
				tc.exclusive,
				tc.force,
			)

			// Verify
			if tc.expectError {
				require.Error(t, err)
				st, ok := status.FromError(err)
				require.True(t, ok)
				assert.Equal(t, tc.expectedCode, st.Code())
				assert.Contains(t, st.Message(), tc.expectedError)
				assert.Nil(t, result)
			} else {
				require.NoError(t, err)
				require.NotNil(t, result)
				assert.Equal(t, tc.tikvGroupId, result.GetId())
				assert.Equal(t, apiv1.ComponentType_TIKV, result.GetComponentType())
				assert.Equal(t, apiv1.Component_CREATING, result.GetPhase())
			}
		})
	}
}

func TestCreateTiKVGroup_GuardBlocksOnRunningTiKVRemoteWorkerDAG(t *testing.T) {
	for _, runningPlan := range tikvRemoteWorkersDAGPlans() {
		t.Run(runningPlan, func(t *testing.T) {
			sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
			sharedPoolRepo, lockRepo, dagExecutor, scaleManager := newTKRWHandlerMocksWithRunningDag(t, sharedPool, runningPlan, runningDagStatus(tkrwKey(t, tkrwEnabledBoth(t))))
			server := &serverImpl{cloudProvider: "aws", sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, dagExecutor: dagExecutor, scaleManager: scaleManager, transactioner: &mockCreateTiKVGroupTransactioner{}}

			_, err := server.CreateTiKVGroup(context.Background(), testTKRWSharedPoolID, "tikvnew", &entity.NodeQuota{NodeSpec: &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16}, NodeCount: 3}, validCreateTiKVGroupStorage(), nil, false, true)
			require.Error(t, err)
			assert.Equal(t, codes.FailedPrecondition, status.Code(err))
			dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, dag.PlanCreateComponent, mock.Anything, mock.Anything)
		})
	}
}

func TestCreateTiKVGroup_GuardBlocksOnTerminalTiKVRemoteWorkerDrift(t *testing.T) {
	sharedPool := newTKRWSharedPool(tkrwEnabledBoth(t), 3)
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	lockRepo := repomocks.NewLockRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	scaleManager := bractionmocks.NewScaleManager(t)
	sharedPoolRepo.On("Load", mock.Anything, testTKRWSharedPoolID, false).Return(sharedPool, nil)
	lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil)
	setupAllDAGMocks(dagExecutor)
	scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe()

	server := &serverImpl{cloudProvider: "aws", sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, dagExecutor: dagExecutor, scaleManager: scaleManager, dataplaneManager: newTKRWDataplaneManager(t), transactioner: &mockCreateTiKVGroupTransactioner{}}
	_, err := server.CreateTiKVGroup(context.Background(), testTKRWSharedPoolID, "tikvnew", &entity.NodeQuota{NodeSpec: &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16}, NodeCount: 3}, validCreateTiKVGroupStorage(), nil, false, false)
	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	dagExecutor.AssertNotCalled(t, "Submit", mock.Anything, dag.PlanCreateComponent, mock.Anything, mock.Anything)
}

func validCreateTiKVGroupStorage() []*entity.Storage {
	return []*entity.Storage{
		{Name: apiv1.Component_Storage_DATA, Type: entity.GP3, SizeGi: 500, Iops: 10000, ThroughputMbs: 500},
		{Name: apiv1.Component_Storage_RAFTLOG, Type: entity.GP3, SizeGi: 100, Iops: 5000, ThroughputMbs: 500},
	}
}