package cluster

import (
	"context"
	"testing"
	"time"

	"github.com/samber/lo"
	. "github.com/smartystreets/goconvey/convey"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/adapter/account/mock_account"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	clusterrepo "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/cluster"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	featgatesvc "github.com/tidbcloud/tidb-management-service/utilities/featuregate"
)

type storageScaleOperationRepo struct {
	clusterrepo.OperationRepo

	recent    map[globaldedicatedv1.ComponentType][]*entity.ClusterOperation
	loadCalls int
}

func (r *storageScaleOperationRepo) LoadRecentScaleOperations(
	_ context.Context,
	_ uint64,
	_ []uint64,
) (map[globaldedicatedv1.ComponentType][]*entity.ClusterOperation, bizerr.BizError) {
	r.loadCalls++
	return r.recent, nil
}

func TestStorageScaleCheck(t *testing.T) {
	ctrl := gomock.NewController(t, gomock.WithOverridableExpectations())
	cfg := globalConf.Load()
	accountCli := mock_account.NewMockClient(ctrl)
	cfg.ServerConfig.FeatureGates = map[string]featgatesvc.FeatureGate{
		featgatesvc.EnableStorageScaleDown: {
			All: false,
		},
	}
	featGateSvc := featgatesvc.NewFeatureGate(
		cfg.ServerConfig.FeatureGates,
		account.BuildOrgCreateTimeFetcher(accountCli),
	)

	s := &serviceImpl{
		featGateCli: featGateSvc,
	}
	ctx := t.Context()

	Convey("TestScaleTiFlashReplicasToZeroThenScaleWithSmallerStorageSize", t, func() {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType: globaldedicatedv1.ComponentType_TIFLASH,
						Count:         lo.ToPtr(uint32(0)),
						Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
							globaldedicatedv1.Sku_DATA.String(): {
								VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_GP3),
								SizeGi:     200,
							},
						},
					},
				},
			},
		}

		desiredComponents := []*entity.DesiredClusterComponent{
			{
				ClusterComponent: &entity.ClusterComponent{
					ComponentType: globaldedicatedv1.ComponentType_TIFLASH,
					Count:         lo.ToPtr(uint32(1)),
					Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
						globaldedicatedv1.Sku_DATA.String(): {
							VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_GP3),
							SizeGi:     100,
						},
					},
				},
			},
		}
		err := s.checkStorageScale(ctx, cluster, desiredComponents)
		assert.Nil(t, err)
	})

	Convey("TestTiFlashDecreaseStorageSize", t, func() {
		cluster := &entity.Cluster{
			Cluster: &globaldedicatedv1.Cluster{
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
					{
						ComponentType: globaldedicatedv1.ComponentType_TIFLASH,
						Count:         lo.ToPtr(uint32(2)),
						Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
							globaldedicatedv1.Sku_DATA.String(): {
								VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_GP3),
								SizeGi:     200,
							},
						},
					},
				},
			},
		}

		desiredComponents := []*entity.DesiredClusterComponent{
			{
				ClusterComponent: &entity.ClusterComponent{
					ComponentType: globaldedicatedv1.ComponentType_TIFLASH,
					Count:         lo.ToPtr(uint32(2)),
					Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
						globaldedicatedv1.Sku_DATA.String(): {
							VolumeType: lo.ToPtr(globaldedicatedv1.StorageMu_GP3),
							SizeGi:     100,
						},
					},
				},
			},
		}
		err := s.checkStorageScale(ctx, cluster, desiredComponents)
		assert.NotNil(t, err)
		assert.Equal(t, err.ErrCode(), bizerr.ErrCodeInvalidClusterModification)
	})
}

func TestGCPHyperdiskStorageCooldown(t *testing.T) {
	featureGateSvc := featgatesvc.NewFeatureGate(
		map[string]featgatesvc.FeatureGate{
			featgatesvc.EnableStorageScaleDown: {All: true},
		},
		nil,
	)

	tests := []struct {
		name                 string
		provider             globalcommonv1.CloudProvider_Name
		volumeType           globaldedicatedv1.StorageMu_VolumeType
		modifyDesiredStorage func(*globaldedicatedv1.Cluster_ClusterComponent_Storage)
		recentAge            time.Duration
		wantLoadCalls        int
		wantCooldownMessage  string
	}{
		{
			name:       "GCP Hyperdisk size change is blocked during four-hour cooldown",
			provider:   globalcommonv1.CloudProvider_GCP,
			volumeType: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED,
			modifyDesiredStorage: func(storage *globaldedicatedv1.Cluster_ClusterComponent_Storage) {
				storage.SizeGi = 1000
			},
			recentAge:           -time.Hour,
			wantLoadCalls:       1,
			wantCooldownMessage: "Wait at least 4 hours",
		},
		{
			name:       "GCP Hyperdisk IOPS change is blocked during four-hour cooldown",
			provider:   globalcommonv1.CloudProvider_GCP,
			volumeType: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED,
			modifyDesiredStorage: func(storage *globaldedicatedv1.Cluster_ClusterComponent_Storage) {
				storage.Iops = 5000
			},
			recentAge:           -time.Hour,
			wantLoadCalls:       1,
			wantCooldownMessage: "Wait at least 4 hours",
		},
		{
			name:       "GCP Hyperdisk throughput change is blocked during four-hour cooldown",
			provider:   globalcommonv1.CloudProvider_GCP,
			volumeType: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED,
			modifyDesiredStorage: func(storage *globaldedicatedv1.Cluster_ClusterComponent_Storage) {
				storage.ThroughputMibps = 250
			},
			recentAge:           -time.Hour,
			wantLoadCalls:       1,
			wantCooldownMessage: "Wait at least 4 hours",
		},
		{
			name:       "GCP Hyperdisk storage change is allowed after four hours",
			provider:   globalcommonv1.CloudProvider_GCP,
			volumeType: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED,
			modifyDesiredStorage: func(storage *globaldedicatedv1.Cluster_ClusterComponent_Storage) {
				storage.SizeGi = 1000
			},
			recentAge:     -5 * time.Hour,
			wantLoadCalls: 1,
		},
		{
			name:       "GCP non-Hyperdisk storage does not use the cooldown",
			provider:   globalcommonv1.CloudProvider_GCP,
			volumeType: globaldedicatedv1.StorageMu_PD_SSD,
			modifyDesiredStorage: func(storage *globaldedicatedv1.Cluster_ClusterComponent_Storage) {
				storage.SizeGi = 1000
			},
			recentAge: -time.Hour,
		},
		{
			name:       "unchanged GCP Hyperdisk storage does not use the cooldown",
			provider:   globalcommonv1.CloudProvider_GCP,
			volumeType: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED,
			recentAge:  -time.Hour,
		},
		{
			name:       "AWS keeps the existing six-hour cooldown",
			provider:   globalcommonv1.CloudProvider_AWS,
			volumeType: globaldedicatedv1.StorageMu_GP3,
			modifyDesiredStorage: func(storage *globaldedicatedv1.Cluster_ClusterComponent_Storage) {
				storage.SizeGi = 1000
			},
			recentAge:           -5 * time.Hour,
			wantLoadCalls:       1,
			wantCooldownMessage: "Wait at least 6 hours",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			const componentID = uint64(11)
			newStorage := func() *globaldedicatedv1.Cluster_ClusterComponent_Storage {
				return &globaldedicatedv1.Cluster_ClusterComponent_Storage{
					VolumeType:      lo.ToPtr(tt.volumeType),
					SizeGi:          500,
					Iops:            3000,
					ThroughputMibps: 140,
				}
			}
			currentStorage := newStorage()
			desiredStorage := newStorage()
			if tt.modifyDesiredStorage != nil {
				tt.modifyDesiredStorage(desiredStorage)
			}
			currentComp := &globaldedicatedv1.Cluster_ClusterComponent{
				Id:            componentID,
				ClusterId:     1,
				ComponentType: globaldedicatedv1.ComponentType_TIKV,
				Count:         lo.ToPtr(uint32(1)),
				Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
					globaldedicatedv1.Sku_DATA.String(): currentStorage,
				},
			}
			desiredComp := &entity.DesiredClusterComponent{ClusterComponent: &entity.ClusterComponent{
				Id:            componentID,
				ComponentType: globaldedicatedv1.ComponentType_TIKV,
				Count:         lo.ToPtr(uint32(1)),
				Storages: map[string]*globaldedicatedv1.Cluster_ClusterComponent_Storage{
					globaldedicatedv1.Sku_DATA.String(): desiredStorage,
				},
			}}
			repo := &storageScaleOperationRepo{
				recent: map[globaldedicatedv1.ComponentType][]*entity.ClusterOperation{
					globaldedicatedv1.ComponentType_TIKV: {
						{
							CreatedAt: time.Now().Add(tt.recentAge),
							Operation: &entity.ScaleOperation{
								ComponentID:    componentID,
								ComponentType:  globaldedicatedv1.ComponentType_TIKV,
								StorageType:    globaldedicatedv1.Sku_DATA,
								StorageChanged: true,
							},
						},
					},
				},
			}
			service := &serviceImpl{featGateCli: featureGateSvc, operationRepo: repo}
			cluster := &entity.Cluster{Cluster: &globaldedicatedv1.Cluster{
				ClusterId:         1,
				CloudProviderId:   tt.provider,
				ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{currentComp},
			}}

			err := service.checkStorageScale(t.Context(), cluster, []*entity.DesiredClusterComponent{desiredComp})

			assert.Equal(t, tt.wantLoadCalls, repo.loadCalls)
			if tt.wantCooldownMessage != "" {
				require.NotNil(t, err)
				assert.Equal(t, bizerr.ErrCodeStorageScaleCoolDown, err.ErrCode())
				assert.Contains(t, err.Error(), tt.wantCooldownMessage)
			} else {
				assert.Nil(t, err)
			}
		})
	}
}