package impl

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/suite"
	"github.com/tidbcloud/config-center/proto/config_center"
	"github.com/tidbcloud/dataflow-service/app/models/domain"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	"github.com/tidbcloud/dataflow-service/internal/config"
	grpcli_mock "github.com/tidbcloud/dataflow-service/pkg/grpc/grpcli/mock"
	globalv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	mgmtv1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
)

type bucketParamBuilderTestSuite struct {
	suite.Suite
}

func TestBucketParamBuilderSuite(t *testing.T) {
	suite.Run(t, new(bucketParamBuilderTestSuite))
}

func (s *bucketParamBuilderTestSuite) TestBucketParamBuilder() {
	snapshot := &entity.BucketSetting{
		Buckets: []entity.BackupSettingBucket{
			{Regions: []string{"us-east-1", "us-east-2"}, BucketId: 101, Master: true},
			{Regions: []string{"eu-central-1"}, BucketId: 202},
		},
	}

	tests := map[string]struct {
		backupRegion  string
		clusterRegion string
		backup        *entity.Backup
		wantBucketID  entity.ID
		wantErr       bool
	}{
		"explicit region - found in snapshot": {
			backupRegion:  "us-east-1",
			clusterRegion: "ap-southeast-1",
			backup:        &entity.Backup{BucketSettingSnapshot: snapshot, BucketID: 999},
			wantBucketID:  101,
		},
		"explicit region - second bucket": {
			backupRegion:  "eu-central-1",
			clusterRegion: "us-east-1",
			backup:        &entity.Backup{BucketSettingSnapshot: snapshot, BucketID: 999},
			wantBucketID:  202,
		},
		"explicit region - not found in snapshot": {
			backupRegion:  "ap-northeast-1",
			clusterRegion: "ap-northeast-1",
			backup:        &entity.Backup{BucketSettingSnapshot: snapshot, BucketID: 999},
			wantErr:       true,
		},
		"explicit region - snapshot nil": {
			backupRegion:  "us-east-1",
			clusterRegion: "us-east-1",
			backup:        &entity.Backup{BucketID: 999},
			wantErr:       true,
		},
		"no backup region - cluster region found in snapshot": {
			backupRegion:  "",
			clusterRegion: "eu-central-1",
			backup:        &entity.Backup{BucketSettingSnapshot: snapshot, BucketID: 999},
			wantBucketID:  202,
		},
		"no backup region - cluster region not found, fallback to default bucket": {
			backupRegion:  "",
			clusterRegion: "ap-northeast-1",
			backup:        &entity.Backup{BucketSettingSnapshot: snapshot, BucketID: 999},
			wantBucketID:  999,
		},
		"no backup region - snapshot nil, fallback to default bucket": {
			backupRegion:  "",
			clusterRegion: "us-east-1",
			backup:        &entity.Backup{BucketID: 777},
			wantBucketID:  777,
		},
		"no backup region - empty cluster region, fallback to default bucket": {
			backupRegion:  "",
			clusterRegion: "",
			backup:        &entity.Backup{BucketSettingSnapshot: snapshot, BucketID: 555},
			wantBucketID:  555,
		},
	}

	model := &baseRestoreModel{
		dedicatedGlobalClient: nil,
	}
	for name, tc := range tests {
		s.Run(name, func() {
			param := &createRestoreParam{
				createRestoreReq: &domain.CreateRestoreReqV2{
					BackupRegion:  tc.backupRegion,
					ClusterConfig: domain.ClusterConfig{Region: tc.clusterRegion},
				},
				backup: tc.backup,
			}

			err := model.bucketParamBuilder(context.Background(), param)

			if tc.wantErr {
				s.Error(err)
			} else {
				s.NoError(err)
				s.Equal(tc.wantBucketID, param.bucketID)
			}
		})
	}
}

// --- buildCreateClusterReqFor3LayerArch RestoreOverride integration tests ---

type buildCreateClusterReqTestSuite struct {
	suite.Suite
}

func TestBuildCreateClusterReqTestSuite(t *testing.T) {
	suite.Run(t, new(buildCreateClusterReqTestSuite))
}

func (s *buildCreateClusterReqTestSuite) baseParam() *createRestoreParam {
	return &createRestoreParam{
		backup: &entity.Backup{
			TenantID:  entity.ID(100),
			ProjectID: entity.ID(200),
			ClusterID: entity.ID(300),
		},
		createRestoreReq: &domain.CreateRestoreReqV2{
			UserID:    1,
			ProjectID: 200,
			ClusterConfig: domain.ClusterConfig{
				ClusterID: 1, // non-zero to skip ValidateOnly dry-run
				Region:    "us-west-2",
				Name:      "test-cluster",
				Port:      4000,
				Components: &domain.ClusterComponents{
					Tidb: &domain.TiDBComponent{NodeSize: "tidb.small", NodeQuantity: 2},
					Tikv: &domain.TiKVComponent{NodeSize: "tikv.small", NodeQuantity: 3, StorageSizeGib: 100, DataStoreIops: 2000, RaftIops: 500},
				},
			},
		},
		createClusterVersion: "v8.5.0",
	}
}

func (s *buildCreateClusterReqTestSuite) backupCluster() *mgmtv1.Cluster {
	return &mgmtv1.Cluster{
		OrganizationId: 999,
		CloudProvider:  "aws",
		Region:         "us-west-2",
		ClusterInfo: &mgmtv1.Cluster_DedicatedCluster{
			DedicatedCluster: &globalv1.Cluster{
				Version: "v8.5.0",
			},
		},
	}
}

func (s *buildCreateClusterReqTestSuite) TestBuildCreateClusterReqFor3LayerArch_RestoreOverride() {
	tests := map[string]struct {
		setupOverride        func()
		wantClusterConfig    bool
		wantAdditionalIops   int32
		wantClusterConfigLen int
	}{
		"no override matched - fields not set": {
			setupOverride: func() {
				config.Global().BackupConfig.RestoreOverrides = nil
			},
			wantClusterConfig:  false,
			wantAdditionalIops: 0,
		},
		"override with ComponentImageTags - ClusterConfiguration populated": {
			setupOverride: func() {
				config.Global().BackupConfig.RestoreOverrides = config.RestoreOverrides{
					{
						Gate: config.FeatureGate{
							EnabledTenants: []uint64{100},
						},
						ComponentImageTags: map[string]string{
							"tidb": "v8.5.0-custom",
							"tikv": "v8.5.0-custom",
						},
					},
				}
			},
			wantClusterConfig:    true,
			wantClusterConfigLen: 2,
			wantAdditionalIops:   0,
		},
		"override with TiKVDataAdditionalIops - TikvComp.AdditionalIops set": {
			setupOverride: func() {
				config.Global().BackupConfig.RestoreOverrides = config.RestoreOverrides{
					{
						Gate: config.FeatureGate{
							EnabledTenants: []uint64{100},
						},
						TiKVDataAdditionalIops: 3000,
					},
				}
			},
			wantClusterConfig:  false,
			wantAdditionalIops: 3000,
		},
		"override with both - both fields set": {
			setupOverride: func() {
				config.Global().BackupConfig.RestoreOverrides = config.RestoreOverrides{
					{
						Gate: config.FeatureGate{
							EnabledTenants: []uint64{100},
						},
						ComponentImageTags: map[string]string{
							"tidb": "v8.5.0-custom",
							"tikv": "v8.5.0-custom",
							"pd":   "v8.5.0-custom",
						},
						TiKVDataAdditionalIops: 5000,
					},
				}
			},
			wantClusterConfig:    true,
			wantClusterConfigLen: 3,
			wantAdditionalIops:   5000,
		},
		"override with zero TiKVDataAdditionalIops - field not set": {
			setupOverride: func() {
				config.Global().BackupConfig.RestoreOverrides = config.RestoreOverrides{
					{
						Gate: config.FeatureGate{
							EnabledTenants: []uint64{100},
						},
						TiKVDataAdditionalIops: 0,
					},
				}
			},
			wantClusterConfig:  false,
			wantAdditionalIops: 0,
		},
	}

	origOverrides := config.Global().BackupConfig.RestoreOverrides
	defer func() { config.Global().BackupConfig.RestoreOverrides = origOverrides }()

	for name, tc := range tests {
		s.Run(name, func() {
			tc.setupOverride()

			model := &baseRestoreModel{
				dedicatedGlobalClient: nil,
			}
			param := s.baseParam()
			backupCluster := s.backupCluster()

			req, err := model.buildCreateClusterReqFor3LayerArch(
				context.Background(), param, nil, backupCluster,
			)

			s.Require().NoError(err)
			s.Require().NotNil(req)

			if tc.wantClusterConfig {
				s.NotNil(req.ClusterConfiguration, "ClusterConfiguration should be set")
				s.Len(req.ClusterConfiguration.ComponentConfigurations, tc.wantClusterConfigLen)
			} else {
				s.Nil(req.ClusterConfiguration, "ClusterConfiguration should not be set")
			}

			s.Require().NotNil(req.TikvComp, "TikvComp should always be set")
			s.Equal(int32(2000), req.TikvComp.DataStoreIops)
			s.Equal(tc.wantAdditionalIops, req.TikvComp.AdditionalIops)
		})
	}
}

func (s *buildCreateClusterReqTestSuite) TestBuildCreateClusterReqFor3LayerArch_DataStoreIops() {
	originalOverrides := config.Global().BackupConfig.RestoreOverrides
	defer func() { config.Global().BackupConfig.RestoreOverrides = originalOverrides }()
	config.Global().BackupConfig.RestoreOverrides = nil

	tests := map[string]struct {
		dataStoreIops     int32
		wantDataStoreIops int32
	}{
		"provided TiKV data IOPS is forwarded": {
			dataStoreIops:     2000,
			wantDataStoreIops: 2000,
		},
		"omitted TiKV data IOPS preserves the default": {
			dataStoreIops:     0,
			wantDataStoreIops: 0,
		},
	}

	for name, tc := range tests {
		s.Run(name, func() {
			param := s.baseParam()
			param.createRestoreReq.ClusterConfig.Components.Tikv.DataStoreIops = tc.dataStoreIops

			req, err := (&baseRestoreModel{}).buildCreateClusterReqFor3LayerArch(
				context.Background(), param, nil, s.backupCluster(),
			)

			s.Require().NoError(err)
			s.Require().NotNil(req.TikvComp)
			s.Equal(tc.wantDataStoreIops, req.TikvComp.DataStoreIops)
		})
	}
}

func (s *bucketParamBuilderTestSuite) TestConvertRestoreProgressToDomain() {
	runningStatus := entity.RestoreKernelStatusRunning
	successStatus := entity.RestoreKernelStatusSuccess

	tests := map[string]struct {
		restore      *entity.Restore
		wantProgress float64
	}{
		"kernel metrics nil returns zero": {
			restore:      &entity.Restore{},
			wantProgress: 0,
		},
		"empty progresses returns zero": {
			restore: &entity.Restore{
				KernelMetrics: &entity.RestoreKernelMetrics{},
			},
			wantProgress: 0,
		},
		"snapshot mode rounds to two decimals": {
			restore: &entity.Restore{
				Mode: entity.RestoreModeSnapshot,
				KernelMetrics: &entity.RestoreKernelMetrics{
					Progresses: []entity.RestoreKernelProgress{
						{Progress: 33},
						{Progress: 67},
						{Progress: 100},
					},
				},
			},
			wantProgress: 66.67,
		},
		"pitr mode uses five steps divisor": {
			restore: &entity.Restore{
				Mode: entity.RestoreModePiTR,
				KernelMetrics: &entity.RestoreKernelMetrics{
					Progresses: []entity.RestoreKernelProgress{
						{Progress: 100},
						{Progress: 100},
						{Progress: 100},
					},
				},
			},
			wantProgress: 60,
		},
		"progress above hundred is capped to hundred": {
			restore: &entity.Restore{
				Mode: entity.RestoreModeSnapshot,
				KernelMetrics: &entity.RestoreKernelMetrics{
					Status: &successStatus,
					Progresses: []entity.RestoreKernelProgress{
						{Progress: 150},
						{Progress: 150},
						{Progress: 150},
					},
				},
			},
			wantProgress: 100,
		},
		"non-success status caps progress at ninety-nine": {
			restore: &entity.Restore{
				Mode: entity.RestoreModeSnapshot,
				KernelMetrics: &entity.RestoreKernelMetrics{
					Status: &runningStatus,
					Progresses: []entity.RestoreKernelProgress{
						{Progress: 100},
						{Progress: 100},
						{Progress: 100},
						{Progress: 100},
					},
				},
			},
			wantProgress: 99,
		},
		"status nil keeps full progress": {
			restore: &entity.Restore{
				Mode: entity.RestoreModeSnapshot,
				KernelMetrics: &entity.RestoreKernelMetrics{
					Progresses: []entity.RestoreKernelProgress{
						{Progress: 100},
						{Progress: 100},
						{Progress: 100},
					},
				},
			},
			wantProgress: 100,
		},
	}

	model := &baseRestoreModel{
		dedicatedGlobalClient: nil,
	}
	for name, tc := range tests {
		s.Run(name, func() {
			actual := model.convertRestoreProgressToDomain(tc.restore)
			s.InDelta(tc.wantProgress, actual, 0.001)
		})
	}
}

func TestBuildClusterConfiguration(t *testing.T) {
	tests := []struct {
		name     string
		tags     map[string]string
		expected *globalv1.ClusterConfiguration
	}{
		{
			name:     "empty map returns nil",
			tags:     map[string]string{},
			expected: nil,
		},
		{
			name: "single component",
			tags: map[string]string{"tidb": "v8.5.0-custom"},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
		{
			name: "multiple components",
			tags: map[string]string{
				"tidb": "v8.5.0-custom",
				"tikv": "v8.5.0-custom",
				"pd":   "v8.5.0-custom",
			},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIKV,
						OverwriteImageTag: "v8.5.0-custom",
					},
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
					{
						ComponentType:     globalv1.ComponentType_PD,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
		{
			name: "invalid key is ignored",
			tags: map[string]string{
				"tidb":    "v8.5.0-custom",
				"invalid": "v8.5.0-custom",
			},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := BuildClusterConfiguration(tt.tags)
			if tt.expected == nil {
				assert.Nil(t, result)
			} else {
				assert.NotNil(t, result)
				assert.Equal(t, len(tt.expected.ComponentConfigurations), len(result.ComponentConfigurations))
				for i, expected := range tt.expected.ComponentConfigurations {
					assert.Equal(t, expected.ComponentType, result.ComponentConfigurations[i].ComponentType)
					assert.Equal(t, expected.OverwriteImageTag, result.ComponentConfigurations[i].OverwriteImageTag)
				}
			}
		})
	}
}

func TestBuildClusterConfigurationWithFallback(t *testing.T) {
	tests := []struct {
		name                   string
		explicitOverrides      map[string]string
		useSameImageTags       []string
		sourceClusterImageTags map[string]string
		expected               *globalv1.ClusterConfiguration
	}{
		{
			name:                   "all empty returns nil",
			explicitOverrides:      map[string]string{},
			useSameImageTags:       []string{},
			sourceClusterImageTags: map[string]string{},
			expected:               nil,
		},
		{
			name:                   "only explicit overrides",
			explicitOverrides:      map[string]string{"tidb": "v8.5.0-custom"},
			useSameImageTags:       []string{},
			sourceClusterImageTags: map[string]string{"tidb": "v8.5.0-source"},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-custom",
					},
				},
			},
		},
		{
			name:                   "only useSameImageTags",
			explicitOverrides:      map[string]string{},
			useSameImageTags:       []string{"tidb", "tikv"},
			sourceClusterImageTags: map[string]string{"tidb": "v8.5.0-source", "tikv": "v8.5.0-source"},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIKV,
						OverwriteImageTag: "v8.5.0-source",
					},
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-source",
					},
				},
			},
		},
		{
			name:                   "explicit overrides take priority",
			explicitOverrides:      map[string]string{"tidb": "v8.5.0-override"},
			useSameImageTags:       []string{"tidb", "tikv"},
			sourceClusterImageTags: map[string]string{"tidb": "v8.5.0-source", "tikv": "v8.5.0-source"},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIKV,
						OverwriteImageTag: "v8.5.0-source",
					},
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-override",
					},
				},
			},
		},
		{
			name:                   "source cluster missing component",
			explicitOverrides:      map[string]string{},
			useSameImageTags:       []string{"tidb", "tikv"},
			sourceClusterImageTags: map[string]string{"tidb": "v8.5.0-source"},
			expected: &globalv1.ClusterConfiguration{
				ComponentConfigurations: []*globalv1.ClusterConfiguration_ComponentConfiguration{
					{
						ComponentType:     globalv1.ComponentType_TIDB,
						OverwriteImageTag: "v8.5.0-source",
					},
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := BuildClusterConfigurationWithFallback(
				tt.explicitOverrides,
				tt.useSameImageTags,
				tt.sourceClusterImageTags,
			)
			if tt.expected == nil {
				assert.Nil(t, result)
			} else {
				assert.NotNil(t, result)
				assert.Equal(t, len(tt.expected.ComponentConfigurations), len(result.ComponentConfigurations))
				for i, expected := range tt.expected.ComponentConfigurations {
					assert.Equal(t, expected.ComponentType, result.ComponentConfigurations[i].ComponentType)
					assert.Equal(t, expected.OverwriteImageTag, result.ComponentConfigurations[i].OverwriteImageTag)
				}
			}
		})
	}
}

// --- canCloneSourceImageTag tests ---

func TestCanCloneSourceImageTag(t *testing.T) {
	model := &baseRestoreModel{}

	tests := []struct {
		name      string
		backup    *entity.Backup
		cluster   *mgmtv1.Cluster
		override  config.RestoreOverride
		wantClone bool
		wantErr   bool
	}{
		{
			name:      "same version - always true",
			backup:    &entity.Backup{ClusterVersion: "v8.5.0"},
			cluster:   makeBackupCluster("v8.5.0"),
			override:  config.RestoreOverride{},
			wantClone: true,
		},
		{
			name:      "different version, ignore patch diff disabled",
			backup:    &entity.Backup{ClusterVersion: "v8.5.0"},
			cluster:   makeBackupCluster("v8.5.3"),
			override:  config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: false},
			wantClone: false,
		},
		{
			name:      "different patch version, ignore patch diff enabled",
			backup:    &entity.Backup{ClusterVersion: "v8.5.0"},
			cluster:   makeBackupCluster("v8.5.3"),
			override:  config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: true},
			wantClone: true,
		},
		{
			name:      "different minor version, ignore patch diff enabled",
			backup:    &entity.Backup{ClusterVersion: "v8.5.0"},
			cluster:   makeBackupCluster("v8.6.0"),
			override:  config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: true},
			wantClone: false,
		},
		{
			name:      "different major version, ignore patch diff enabled",
			backup:    &entity.Backup{ClusterVersion: "v8.5.0"},
			cluster:   makeBackupCluster("v9.5.0"),
			override:  config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: true},
			wantClone: false,
		},
		{
			name:     "version without v prefix - fails semver parse",
			backup:   &entity.Backup{ClusterVersion: "8.5.0"},
			cluster:  makeBackupCluster("8.5.3"),
			override: config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: true},
			wantErr:  true,
		},
		{
			name:     "invalid backup version",
			backup:   &entity.Backup{ClusterVersion: "invalid"},
			cluster:  makeBackupCluster("v8.5.0"),
			override: config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: true},
			wantErr:  true,
		},
		{
			name:     "invalid cluster version",
			backup:   &entity.Backup{ClusterVersion: "v8.5.0"},
			cluster:  makeBackupCluster("invalid"),
			override: config.RestoreOverride{CloneImageTagsIgnorePatchVersionDiff: true},
			wantErr:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := model.canCloneSourceImageTag(context.Background(), tt.backup, tt.cluster, tt.override)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			assert.NoError(t, err)
			assert.Equal(t, tt.wantClone, got)
		})
	}
}

// makeBackupCluster builds a minimal mgmtv1.Cluster with a DedicatedCluster of the given version.
func makeBackupCluster(version string) *mgmtv1.Cluster {
	return &mgmtv1.Cluster{
		ClusterInfo: &mgmtv1.Cluster_DedicatedCluster{
			DedicatedCluster: &globalv1.Cluster{
				Version: version,
			},
		},
	}
}

// --- overrideImages tests ---

type overrideImagesTestSuite struct {
	suite.Suite
}

func TestOverrideImagesSuite(t *testing.T) {
	suite.Run(t, new(overrideImagesTestSuite))
}

func (s *overrideImagesTestSuite) TestOverrideImages() {
	tests := map[string]struct {
		override          config.RestoreOverride
		backup            *entity.Backup
		backupCluster     *mgmtv1.Cluster
		setupMock         func(m *grpcli_mock.ConfigCenterGrpcClient)
		wantClusterConfig bool
		wantConfigLen     int
		wantErr           bool
	}{
		"no ComponentImageTags and no CloneImageTags - no-op": {
			override:      config.RestoreOverride{},
			backup:        &entity.Backup{ClusterVersion: "v8.5.0"},
			backupCluster: makeBackupCluster("v8.5.0"),
			setupMock:     func(m *grpcli_mock.ConfigCenterGrpcClient) {},
		},
		"ComponentImageTags only - sets ClusterConfiguration": {
			override: config.RestoreOverride{
				ComponentImageTags: map[string]string{"tidb": "v8.5.0-custom", "tikv": "v8.5.0-custom"},
			},
			backup:            &entity.Backup{ClusterVersion: "v8.5.0"},
			backupCluster:     makeBackupCluster("v8.5.0"),
			setupMock:         func(m *grpcli_mock.ConfigCenterGrpcClient) {},
			wantClusterConfig: true,
			wantConfigLen:     2,
		},
		"CloneImageTags only, can clone - calls config center and sets ClusterConfiguration": {
			override: config.RestoreOverride{
				CloneImageTags:                       []string{"tidb", "tikv"},
				CloneImageTagsIgnorePatchVersionDiff: true,
			},
			backup:        &entity.Backup{ClusterVersion: "v8.5.0"},
			backupCluster: makeBackupClusterWithID("v8.5.3", 12345),
			setupMock: func(m *grpcli_mock.ConfigCenterGrpcClient) {
				m.On("ListClusterComponents", mock.Anything, uint64(12345)).
					Return(map[string]*config_center.ClusterComponent{
						"tidb": {ComponentName: "tidb", EffectiveImageTag: "v8.5.3-tag"},
						"tikv": {ComponentName: "tikv", EffectiveImageTag: "v8.5.3-tag"},
					}, nil)
			},
			wantClusterConfig: true,
			wantConfigLen:     2,
		},
		"CloneImageTags only, cannot clone - no ClusterConfiguration": {
			override: config.RestoreOverride{
				CloneImageTags:                       []string{"tidb"},
				CloneImageTagsIgnorePatchVersionDiff: false,
			},
			backup:        &entity.Backup{ClusterVersion: "v8.5.0"},
			backupCluster: makeBackupCluster("v8.5.3"),
			setupMock:     func(m *grpcli_mock.ConfigCenterGrpcClient) {},
		},
		"ComponentImageTags overrides CloneImageTags": {
			override: config.RestoreOverride{
				ComponentImageTags:                   map[string]string{"tidb": "explicit-tag"},
				CloneImageTags:                       []string{"tidb", "tikv"},
				CloneImageTagsIgnorePatchVersionDiff: true,
			},
			backup:        &entity.Backup{ClusterVersion: "v8.5.0"},
			backupCluster: makeBackupClusterWithID("v8.5.3", 12345),
			setupMock: func(m *grpcli_mock.ConfigCenterGrpcClient) {
				m.On("ListClusterComponents", mock.Anything, uint64(12345)).
					Return(map[string]*config_center.ClusterComponent{
						"tidb": {ComponentName: "tidb", EffectiveImageTag: "v8.5.3-source"},
						"tikv": {ComponentName: "tikv", EffectiveImageTag: "v8.5.3-source"},
					}, nil)
			},
			wantClusterConfig: true,
			wantConfigLen:     2,
		},
	}

	for name, tc := range tests {
		s.Run(name, func() {
			mockCli := grpcli_mock.NewConfigCenterGrpcClient(s.T())
			tc.setupMock(mockCli)

			model := &baseRestoreModel{
				configCenterCli: mockCli,
			}
			req := &mgmtv1.RestoreDedicatedClusterRequest{
				TikvComp: &mgmtv1.RestoreDedicatedClusterRequest_StorageComp{},
			}

			err := model.overrideImages(context.Background(), tc.override, req, tc.backup, tc.backupCluster)

			if tc.wantErr {
				s.Error(err)
				return
			}
			s.NoError(err)

			if tc.wantClusterConfig {
				s.NotNil(req.ClusterConfiguration, "ClusterConfiguration should be set")
				s.Len(req.ClusterConfiguration.ComponentConfigurations, tc.wantConfigLen)
			} else {
				s.Nil(req.ClusterConfiguration, "ClusterConfiguration should not be set")
			}
		})
	}
}

func makeBackupClusterWithID(version string, clusterID uint64) *mgmtv1.Cluster {
	return &mgmtv1.Cluster{
		ClusterId: clusterID,
		ClusterInfo: &mgmtv1.Cluster_DedicatedCluster{
			DedicatedCluster: &globalv1.Cluster{
				Version: version,
			},
		},
	}
}

// --- overrideIops tests ---

func u32p(v uint32) *uint32 { return &v }

type overrideIopsTestSuite struct {
	suite.Suite
}

func TestOverrideIopsSuite(t *testing.T) {
	suite.Run(t, new(overrideIopsTestSuite))
}

func (s *overrideIopsTestSuite) TestOverrideIops() {
	tests := map[string]struct {
		override        config.RestoreOverride
		backupCluster   *mgmtv1.Cluster
		wantTikvIops    int32
		wantTiflashIops int32
		wantErr         bool
	}{
		"no override - no change": {
			override:      config.RestoreOverride{},
			backupCluster: makeBackupCluster("v8.5.0"),
			wantTikvIops:  0,
		},
		"TiKVDataAdditionalIops sets TikvComp.AdditionalIops": {
			override: config.RestoreOverride{
				TiKVDataAdditionalIops: 5000,
			},
			backupCluster: makeBackupCluster("v8.5.0"),
			wantTikvIops:  5000,
		},
		"CloneAdditionalIops tikv - clones from source cluster": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tikv"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIKV,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						globalv1.Sku_DATA.String(): {AdditionalIops: u32p(3000)},
					},
				},
			}),
			wantTikvIops: 3000,
		},
		"CloneAdditionalIops tiflash - clones from source cluster": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tiflash"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIFLASH,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						globalv1.Sku_DATA.String(): {AdditionalIops: u32p(2000)},
					},
				},
			}),
			wantTiflashIops: 2000,
		},
		"CloneAdditionalIops tikv+tiflash - clones both": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tikv", "tiflash"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIKV,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						globalv1.Sku_DATA.String(): {AdditionalIops: u32p(3000)},
					},
				},
				{
					ComponentType: globalv1.ComponentType_TIFLASH,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						globalv1.Sku_DATA.String(): {AdditionalIops: u32p(2000)},
					},
				},
			}),
			wantTikvIops:    3000,
			wantTiflashIops: 2000,
		},
		"TiKVDataAdditionalIops overrides CloneAdditionalIops for tikv": {
			override: config.RestoreOverride{
				TiKVDataAdditionalIops: 5000,
				CloneAdditionalIops:    []string{"tikv"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIKV,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						globalv1.Sku_DATA.String(): {AdditionalIops: u32p(3000)},
					},
				},
			}),
			wantTikvIops: 5000,
		},
		"CloneAdditionalIops with zero additional IOPS - no change": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tikv"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIKV,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						globalv1.Sku_DATA.String(): {AdditionalIops: u32p(0)},
					},
				},
			}),
			wantTikvIops: 0,
		},
		"CloneAdditionalIops with non-DATA storage type - ignored": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tikv"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIKV,
					Storages: map[string]*globalv1.Cluster_ClusterComponent_Storage{
						"OTHER": {AdditionalIops: u32p(3000)},
					},
				},
			}),
			wantTikvIops: 0,
		},
		"CloneAdditionalIops for non-storage component (tidb) - ignored": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tidb"},
			},
			backupCluster: makeBackupClusterWithComponents("v8.5.0", []*globalv1.Cluster_ClusterComponent{
				{
					ComponentType: globalv1.ComponentType_TIDB,
				},
			}),
			wantTikvIops: 0,
		},
		"nil dedicated cluster - no panic": {
			override: config.RestoreOverride{
				CloneAdditionalIops: []string{"tikv"},
			},
			backupCluster: &mgmtv1.Cluster{},
			wantTikvIops:  0,
		},
	}

	for name, tc := range tests {
		s.Run(name, func() {
			model := &baseRestoreModel{}
			req := &mgmtv1.RestoreDedicatedClusterRequest{
				TikvComp:    &mgmtv1.RestoreDedicatedClusterRequest_StorageComp{},
				TiflashComp: &mgmtv1.RestoreDedicatedClusterRequest_StorageComp{},
			}

			err := model.overrideIops(context.Background(), tc.override, req, nil, tc.backupCluster)

			if tc.wantErr {
				s.Error(err)
				return
			}
			s.NoError(err)
			s.Equal(tc.wantTikvIops, req.TikvComp.AdditionalIops, "TikvComp.AdditionalIops mismatch")
			if tc.wantTiflashIops > 0 {
				s.Equal(tc.wantTiflashIops, req.TiflashComp.AdditionalIops, "TiflashComp.AdditionalIops mismatch")
			}
		})
	}
}

func makeBackupClusterWithComponents(version string, components []*globalv1.Cluster_ClusterComponent) *mgmtv1.Cluster {
	return &mgmtv1.Cluster{
		ClusterInfo: &mgmtv1.Cluster_DedicatedCluster{
			DedicatedCluster: &globalv1.Cluster{
				Version:           version,
				ClusterComponents: components,
			},
		},
	}
}