package utils

import (
	"reflect"
	"testing"

	. "github.com/onsi/gomega"
	storagev1 "k8s.io/api/storage/v1"
	storagev1beta1 "k8s.io/api/storage/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
)

func TestParseStorageClassFromName(t *testing.T) {
	g := NewGomegaWithT(t)

	type Case struct {
		name     string
		expected *storagev1.StorageClass
	}
	validSCNames := []Case{
		{
			"tikv-gp2",
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tikv-gp2"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          ebsProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"type":                      "gp2",
					"encrypted":                 "true",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"tikv-pd-ssd",
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tikv-pd-ssd"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          pdProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"type":                      "pd-ssd",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"tiflash-gp3-0-2",
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tiflash-gp3-0-2"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          ebsProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"encrypted":                 "true",
					"type":                      "gp3",
					"iops":                      "0",
					"throughput":                "2",
					"tagSpecification":          "Vendor=PingCAP",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"tidb-io2-20-whatever", // for io1 and io2, throughput is ignored
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tidb-io2-20-whatever"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          ebsProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"encrypted":                 "true",
					"type":                      "io2",
					"iops":                      "20",
					"tagSpecification":          "Vendor=PingCAP",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"mysql-io1-0-100", // for io1 and io2, throughput is ignored
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "mysql-io1-0-100"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          ebsProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"encrypted":                 "true",
					"type":                      "io1",
					"iops":                      "0",
					"tagSpecification":          "Vendor=PingCAP",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"tikv-premiumv2-1000-800",
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tikv-premiumv2-1000-800"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          azureProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"DiskIOPSReadWrite":         "1000",
					"DiskMBpsReadWrite":         "800",
					"skuName":                   "PremiumV2_LRS",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"tikv-ultrassd-1000-800",
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tikv-ultrassd-1000-800"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          azureProvisioner,
				Parameters: map[string]string{
					"kind":                      "managed",
					"csi.storage.k8s.io/fstype": "ext4",
					"DiskIOPSReadWrite":         "1000",
					"DiskMBpsReadWrite":         "800",
					"skuName":                   "UltraSSD_LRS",
					"cachingMode":               "None",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
		{
			"tikv-hyperdisk-balanced-1000-800",
			&storagev1.StorageClass{
				ObjectMeta:           metav1.ObjectMeta{Name: "tikv-hyperdisk-balanced-1000-800"},
				VolumeBindingMode:    &bindingMode,
				AllowVolumeExpansion: pointer.Bool(true),
				ReclaimPolicy:        &reclaimPolicy,
				Provisioner:          pdProvisioner,
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype":        "ext4",
					"type":                             "hyperdisk-balanced",
					"provisioned-throughput-on-create": "800Mi",
					"provisioned-iops-on-create":       "1000",
				},
				MountOptions: []string{"nodelalloc,noatime"},
			},
		},
	}
	for _, c := range validSCNames {
		sc, err := ParseStorageClassFromName(c.name)
		g.Expect(err).Should(BeNil())
		g.Expect(sc).To(Equal(c.expected))
	}
	invalidSCNames := []Case{
		{"tikv", nil},
		{"tikv-pd-ssd-3-4", nil},
		{"tikv-pd-ssd-3", nil},
		{"tikv-gp2-100-100", nil},
		{"tikv-gp3", nil},
		{"tikv-gp3-xx-100", nil},
		{"tikv-io2-xx-100", nil},
		{"tikv-gp4-100-100", nil},
	}
	for _, c := range invalidSCNames {
		sc, err := ParseStorageClassFromName(c.name)
		g.Expect(err).ShouldNot(BeNil())
		g.Expect(sc).To(Equal(c.expected))
	}
}

func TestGetVolumeAttributesClassFromFromStorageClass(t *testing.T) {
	tests := []struct {
		name string
		sc   *storagev1.StorageClass
		want *storagev1beta1.VolumeAttributesClass
	}{
		{
			name: "aws-gp3",
			sc: &storagev1.StorageClass{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-gp3-3000-125",
				},
				Provisioner: "ebs.csi.aws.com",
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"encrypted":                 "true",
					"iops":                      "3000",
					"tagSpecification":          "Vendor=PingCAP",
					"tagSpecification_1":        "project=67890",
					"tagSpecification_2":        "servicetype=dedicated",
					"tagSpecification_3":        "tenant=12345",
					"tagSpecification_4":        "usedby=customer-tidb",
					"throughput":                "125",
					"type":                      "gp3",
				},
			},
			want: &storagev1beta1.VolumeAttributesClass{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-gp3-3000-125",
				},
				DriverName: "ebs.csi.aws.com",
				Parameters: map[string]string{
					"iops":       "3000",
					"throughput": "125",
					"type":       "gp3",
				},
			},
		},
		{
			name: "does not have supported parameters",
			sc: &storagev1.StorageClass{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-gp3-3000-125",
				},
				Provisioner: "ebs.csi.aws.com",
				Parameters: map[string]string{
					"csi.storage.k8s.io/fstype": "ext4",
					"encrypted":                 "true",
					"tagSpecification":          "Vendor=PingCAP",
				},
			},
			want: nil,
		},
		{
			name: "gcp",
			sc: &storagev1.StorageClass{
				ObjectMeta: metav1.ObjectMeta{
					Name: "gcp-pd-ssd",
				},
				Provisioner: pdProvisioner,
				Parameters: map[string]string{
					"type":       "pd-ssd",
					"iops":       "3000",
					"throughput": "125",
				},
			},
			want: nil,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := GetVolumeAttributesClassFromFromStorageClass(tt.sc); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetVolumeAttributesClassFromFromStorageClass() = %v, want %v", got, tt.want)
			}
		})
	}
}