package cluster

import (
	"testing"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"

	. "github.com/onsi/gomega"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/utils/pointer"
	ctrl "sigs.k8s.io/controller-runtime"
)

func TestScaleStorage(t *testing.T) {
	cases := map[string]struct {
		setup  func(*cloudv1.Cluster, *v1alpha1.TidbCluster)
		expect func(*WithT, *v1alpha1.TidbCluster, error)
	}{
		"NoScale": {
			setup: func(c *cloudv1.Cluster, tc *v1alpha1.TidbCluster) {
				desiredSize := resource.MustParse("10Gi")
				c.Spec.Cluster.PD.StorageSize = desiredSize
				c.Spec.Cluster.TiKV.StorageSize = desiredSize
				c.Spec.Cluster.TiFlash.StorageSize = desiredSize
				tc.Spec.PD.Requests[corev1.ResourceStorage] = desiredSize
				tc.Spec.TiKV.Requests[corev1.ResourceStorage] = desiredSize
				tc.Spec.TiFlash.StorageClaims[0].Resources.Requests[corev1.ResourceStorage] = desiredSize
			},
			expect: func(g *WithT, tc *v1alpha1.TidbCluster, err error) {
				desiredSize := resource.MustParse("10Gi")
				g.Expect(err).Should(Succeed())
				g.Expect(tc.Spec.PD.Requests[corev1.ResourceStorage]).Should(Equal(desiredSize))
				g.Expect(tc.Spec.TiKV.Requests[corev1.ResourceStorage]).Should(Equal(desiredSize))
				g.Expect(tc.Spec.TiFlash.StorageClaims[0].Resources.Requests[corev1.ResourceStorage]).Should(Equal(desiredSize))
			},
		},
		"Scaleup": {
			setup: func(c *cloudv1.Cluster, tc *v1alpha1.TidbCluster) {
				currentSize := resource.MustParse("5Gi")
				desiredSize := resource.MustParse("10Gi")
				c.Spec.Cluster.PD.StorageSize = desiredSize
				c.Spec.Cluster.TiKV.StorageSize = desiredSize
				c.Spec.Cluster.TiFlash.StorageSize = desiredSize
				tc.Spec.PD.Requests[corev1.ResourceStorage] = currentSize
				tc.Spec.TiKV.Requests[corev1.ResourceStorage] = currentSize
				tc.Spec.TiFlash.StorageClaims[0].Resources.Requests[corev1.ResourceStorage] = currentSize
			},
			expect: func(g *WithT, tc *v1alpha1.TidbCluster, err error) {
				desiredSize := resource.MustParse("10Gi")
				g.Expect(err).Should(Succeed())
				g.Expect(tc.Spec.PD.Requests[corev1.ResourceStorage]).Should(Equal(desiredSize))
				g.Expect(tc.Spec.TiKV.Requests[corev1.ResourceStorage]).Should(Equal(desiredSize))
				g.Expect(tc.Spec.TiFlash.StorageClaims[0].Resources.Requests[corev1.ResourceStorage]).Should(Equal(desiredSize))
			},
		},
		"Scaledown": {
			setup: func(c *cloudv1.Cluster, tc *v1alpha1.TidbCluster) {
				currentSize := resource.MustParse("10Gi")
				desiredSize := resource.MustParse("5Gi")
				c.Spec.Cluster.PD.StorageSize = desiredSize
				c.Spec.Cluster.TiKV.StorageSize = desiredSize
				c.Spec.Cluster.TiFlash.StorageSize = desiredSize
				tc.Spec.PD.Requests[corev1.ResourceStorage] = currentSize
				tc.Spec.TiKV.Requests[corev1.ResourceStorage] = currentSize
				tc.Spec.TiFlash.StorageClaims[0].Resources.Requests[corev1.ResourceStorage] = currentSize
			},
			expect: func(g *WithT, tc *v1alpha1.TidbCluster, err error) {
				g.Expect(err).Should(HaveOccurred())
				g.Expect(err.Error()).Should(ContainSubstring("not support to scale down"))
			},
		},
		"ScaledownButLocalStorage": {
			setup: func(c *cloudv1.Cluster, tc *v1alpha1.TidbCluster) {
				currentSize := resource.MustParse("10Gi")
				desiredSize := resource.MustParse("5Gi")
				c.Spec.Cluster.PD.StorageSize = desiredSize
				c.Spec.Cluster.TiKV.StorageSize = desiredSize
				c.Spec.Cluster.TiFlash.StorageSize = desiredSize
				tc.Spec.PD.Requests[corev1.ResourceStorage] = desiredSize
				tc.Spec.TiKV.Requests[corev1.ResourceStorage] = currentSize
				tc.Spec.TiFlash.StorageClaims[0].Resources.Requests[corev1.ResourceStorage] = currentSize
				c.Spec.Cluster.TiKV.UseLocalStorage = pointer.BoolPtr(true)
				c.Spec.Cluster.TiFlash.UseLocalStorage = pointer.BoolPtr(true)
			},
			expect: func(g *WithT, tc *v1alpha1.TidbCluster, err error) {
				g.Expect(err).Should(Succeed())
			},
		},
	}

	for name, tt := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			cluster := &cloudv1.Cluster{}
			cluster.Name = "scale-test"
			cluster.Namespace = "default"
			cluster.Spec.Cluster.PD.Replicas = 3
			cluster.Spec.Cluster.TiDB.Replicas = 3
			cluster.Spec.Cluster.TiKV.Replicas = 3
			cluster.Spec.Cluster.TiFlash.Replicas = 3
			tc := &v1alpha1.TidbCluster{}
			tc.Name = "scale-test-tc"
			tc.Namespace = "default"
			requirements := corev1.ResourceRequirements{
				Requests: map[corev1.ResourceName]resource.Quantity{},
			}
			tc.Spec.PD = &v1alpha1.PDSpec{Replicas: 3, ResourceRequirements: *requirements.DeepCopy()}
			tc.Spec.TiDB = &v1alpha1.TiDBSpec{Replicas: 3, ResourceRequirements: *requirements.DeepCopy()}
			tc.Spec.TiKV = &v1alpha1.TiKVSpec{Replicas: 3, ResourceRequirements: *requirements.DeepCopy()}
			tc.Spec.TiFlash = &v1alpha1.TiFlashSpec{Replicas: 3, StorageClaims: make([]v1alpha1.StorageClaim, 1)}
			tc.Spec.TiFlash.StorageClaims[0] = v1alpha1.StorageClaim{Resources: *requirements.DeepCopy()}
			ctx := &reconcileContext{
				ClusterContext: ClusterContext{
					Target: cluster,
				},
				Log: ctrl.Log.WithName(t.Name()),
			}
			tt.setup(cluster, tc)

			err := ctx.scaleStorage(tc)
			tt.expect(g, tc, err)
		})
	}
}