package common

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/utils/ptr"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/feature"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

type mockClusterComponent struct {
	Replicas         int
	InstanceType     string
	ComponentType    cloudv1.ComponentType
	StorageClassName string
	StorageSize      *resource.Quantity
}

var _ cloudv1.ClusterComponent = &mockClusterComponent{}

func (m mockClusterComponent) GetReplicas() int { return m.Replicas }

func (m mockClusterComponent) GetInstanceType() string { return m.InstanceType }

func (m mockClusterComponent) GetComponentType() cloudv1.ComponentType { return m.ComponentType }

func (m mockClusterComponent) GetStorageClassName() (string, bool) { return m.StorageClassName, true }

func (m mockClusterComponent) GetStorageSize() (*resource.Quantity, bool) { return m.StorageSize, true }

func (m mockClusterComponent) GetUseLocalStorage() bool { return false }

func (m mockClusterComponent) NeedSuspend(action *cloudv1.SuspendAction) bool { return false }

func (m mockClusterComponent) GetSuspendActions(action *cloudv1.SuspendAction) *cloudv1.SuspendAction {
	return nil
}

func (m mockClusterComponent) SetReplicas(replicas int) {}

func (m mockClusterComponent) SetInstanceType(instanceType string) {}

func (m mockClusterComponent) SetStorageClassName(className string) {}

func (m mockClusterComponent) SetStorageSize(size resource.Quantity) {}

func (m mockClusterComponent) SetUseLocalStorage(useLocalStorage bool) {}

func (c mockClusterComponent) GetAdditionalStorages() []cloudv1.ComponentStorageVolume {
	return nil
}

func (c mockClusterComponent) SetAdditionalStorages(volumes []cloudv1.ComponentStorageVolume) {}

func TestTiDBClusterStatusResolverV2_computeVolumeStatus(t *testing.T) {
	tests := []struct {
		name            string
		comp            cloudv1.ClusterComponent
		prev            *cloudv1.VolumeStatus
		replicas        int
		currentReplicas int
		updateReplicas  int
		storageSizeMap  map[string]int
		scNameMap       map[string]int
		vacNameMap      map[string]int

		enableCMEK bool

		wantVs        *cloudv1.VolumeStatus
		wantModifying bool
		wantErr       bool
	}{
		{
			name: "initial status",
			comp: &mockClusterComponent{
				Replicas:         3,
				ComponentType:    cloudv1.ComponentTiKV,
				StorageClassName: "sc-1",
				StorageSize:      ptr.To(resource.MustParse("10Gi")),
			},
			prev:            nil,
			replicas:        3,
			currentReplicas: 3,
			updateReplicas:  3,
			storageSizeMap:  map[string]int{"10Gi": 3},
			scNameMap:       map[string]int{"sc-1": 3},
			wantVs: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         3,
				ModifiedCount:        3,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("10Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-1",
				ResizedCount:         3,
				ResizedCapacity:      resource.MustParse("10Gi"),
			},
		},
		{
			name: "scale out",
			comp: &mockClusterComponent{
				Replicas:         6,
				ComponentType:    cloudv1.ComponentTiKV,
				StorageClassName: "sc-1",
				StorageSize:      ptr.To(resource.MustParse("10Gi")),
			},
			prev: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         3,
				ModifiedCount:        3,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("10Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-1",
				ResizedCount:         3,
				ResizedCapacity:      resource.MustParse("10Gi"),
			},
			replicas:        6,
			currentReplicas: 6,
			updateReplicas:  6,
			storageSizeMap:  map[string]int{"10Gi": 6},
			scNameMap:       map[string]int{"sc-1": 6},
			wantVs: &cloudv1.VolumeStatus{
				BoundCount:           6,
				CurrentCount:         6,
				ModifiedCount:        6,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("10Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-1",
				ResizedCount:         6,
				ResizedCapacity:      resource.MustParse("10Gi"),
			},
		},
		{
			name: "modify storage, progress 1/3",
			comp: &mockClusterComponent{
				Replicas:         3,
				ComponentType:    cloudv1.ComponentTiKV,
				StorageClassName: "sc-2",
				StorageSize:      ptr.To(resource.MustParse("30Gi")),
			},
			prev: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         3,
				ModifiedCount:        3,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("10Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-1",
				ResizedCount:         3,
				ResizedCapacity:      resource.MustParse("10Gi"),
			},
			replicas:        3,
			currentReplicas: 2,
			updateReplicas:  1,
			storageSizeMap:  map[string]int{"10Gi": 2, "30Gi": 1},
			scNameMap:       map[string]int{"sc-1": 2, "sc-2": 1},
			wantVs: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         2,
				ModifiedCount:        1,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("30Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-2",
				ResizedCount:         1,
				ResizedCapacity:      resource.MustParse("30Gi"),
			},
			wantModifying: true,
		},
		{
			name: "modify storage, progress 2/3",
			comp: &mockClusterComponent{
				Replicas:         3,
				ComponentType:    cloudv1.ComponentTiKV,
				StorageClassName: "sc-2",
				StorageSize:      ptr.To(resource.MustParse("30Gi")),
			},
			prev: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         2,
				ModifiedCount:        1,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("30Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-2",
				ResizedCount:         1,
				ResizedCapacity:      resource.MustParse("30Gi"),
			},
			replicas:        3,
			currentReplicas: 1,
			updateReplicas:  2,
			storageSizeMap:  map[string]int{"10Gi": 1, "30Gi": 2},
			scNameMap:       map[string]int{"sc-1": 1, "sc-2": 2},
			wantVs: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         1,
				ModifiedCount:        2,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("30Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-2",
				ResizedCount:         2,
				ResizedCapacity:      resource.MustParse("30Gi"),
			},
			wantModifying: true,
		},
		{
			name: "modify storage, progress 3/3",
			comp: &mockClusterComponent{
				Replicas:         3,
				ComponentType:    cloudv1.ComponentTiKV,
				StorageClassName: "sc-2",
				StorageSize:      ptr.To(resource.MustParse("30Gi")),
			},
			prev: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         1,
				ModifiedCount:        2,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("30Gi"),
				CurrentStorageClass:  "sc-1",
				ModifiedStorageClass: "sc-2",
				ResizedCount:         2,
				ResizedCapacity:      resource.MustParse("30Gi"),
			},
			replicas:        3,
			currentReplicas: 3,
			updateReplicas:  3,
			storageSizeMap:  map[string]int{"10Gi": 0, "30Gi": 3},
			scNameMap:       map[string]int{"sc-1": 0, "sc-2": 3},
			wantVs: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         3,
				ModifiedCount:        3,
				CurrentCapacity:      resource.MustParse("30Gi"),
				ModifiedCapacity:     resource.MustParse("30Gi"),
				CurrentStorageClass:  "sc-2",
				ModifiedStorageClass: "sc-2",
				ResizedCount:         3,
				ResizedCapacity:      resource.MustParse("30Gi"),
			},
		},
		{
			name: "enable cmek",
			comp: &mockClusterComponent{
				Replicas:         3,
				ComponentType:    cloudv1.ComponentTiKV,
				StorageClassName: "sc-1",
				StorageSize:      ptr.To(resource.MustParse("10Gi")),
			},
			prev:            nil,
			replicas:        3,
			currentReplicas: 3,
			updateReplicas:  3,
			storageSizeMap:  map[string]int{"10Gi": 3},
			scNameMap:       map[string]int{"sc-1-200683210682-key-ab2a9e23-1197-4c9d-851d-007b3ead8889": 3},
			enableCMEK:      true,
			wantVs: &cloudv1.VolumeStatus{
				BoundCount:           3,
				CurrentCount:         3,
				ModifiedCount:        3,
				CurrentCapacity:      resource.MustParse("10Gi"),
				ModifiedCapacity:     resource.MustParse("10Gi"),
				CurrentStorageClass:  "sc-1-200683210682-key-ab2a9e23-1197-4c9d-851d-007b3ead8889",
				ModifiedStorageClass: "sc-1-200683210682-key-ab2a9e23-1197-4c9d-851d-007b3ead8889",
				ResizedCount:         3,
				ResizedCapacity:      resource.MustParse("10Gi"),
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			r := &TiDBClusterStatusResolverV2{
				Gates:  feature.NewGates(),
				Mapper: cloudv1.NewNopMapper(),
			}
			if tt.enableCMEK {
				r.Gates = feature.NewGates(feature.UseCMEK)
				mapper, err := cloudv1.NewAWSStorageClassMapper("arn:aws:kms:us-west-2:200683210682:key/ab2a9e23-1197-4c9d-851d-007b3ead8889")
				assert.NoError(t, err, "kms key can be parsed")
				r.Mapper = mapper
			}
			gotVs, gotModifying, err := r.computeVolumeStatus(tt.comp, tt.prev, tt.replicas, tt.currentReplicas, tt.updateReplicas, tt.storageSizeMap, tt.scNameMap, tt.vacNameMap)
			if (err != nil) != tt.wantErr {
				t.Errorf("computeVolumeStatus() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			assert.Equalf(t, tt.wantVs, gotVs, "computeVolumeStatus(%v, %v, %v, %v, %v, %v)", tt.comp, tt.prev, tt.replicas, tt.currentReplicas, tt.updateReplicas, tt.storageSizeMap)
			assert.Equalf(t, tt.wantModifying, gotModifying, "computeVolumeStatus(%v, %v, %v, %v, %v, %v)", tt.comp, tt.prev, tt.replicas, tt.currentReplicas, tt.updateReplicas, tt.storageSizeMap)
		})
	}
}