package nodegroup

import (
	"testing"

	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/utils/pointer"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/common"
	gardencorev1beta1 "github.com/tidbcloud/infra-api/pkg/gardener/core/apis/v1beta1"
	"github.com/tidbcloud/infra-api/pkg/gitops"
)

var x86ImageVersion = "7.8.2003"

func TestConvertGCPVolumeCounts(t *testing.T) {
	g := NewGomegaWithT(t)
	g.Expect(convertGCPVolumeCounts(resource.MustParse("375Gi"))).To(Equal(2))
	g.Expect(convertGCPVolumeCounts(resource.MustParse("750Gi"))).To(Equal(2))
	g.Expect(convertGCPVolumeCounts(resource.MustParse("1125Gi"))).To(Equal(4))
	g.Expect(convertGCPVolumeCounts(resource.MustParse("1500Gi"))).To(Equal(4))
	g.Expect(convertGCPVolumeCounts(resource.MustParse("1510Gi"))).To(Equal(6))
	g.Expect(convertGCPVolumeCounts(resource.MustParse("2240Gi"))).To(Equal(6))
	g.Expect(convertGCPVolumeCounts(resource.MustParse("2250Gi"))).To(Equal(6))
}

func TestBuildNodeGroup(t *testing.T) {
	cases := map[string]struct {
		nodeSpec    *NodeTopologySpec
		component   cloudv1.ClusterComponent
		machineSpec *cloudv1.MachineSpec
		Zone        string
		expect      *gardencorev1beta1.Worker
	}{
		"AWSTiKVCloudDisk": {
			nodeSpec: &NodeTopologySpec{
				Provider: &cloudv1.Provider{
					Spec: cloudv1.ProviderSpec{
						Type: "aws",
					},
				},
				Tags: map[string]string{
					"cluster": "1",
					"tenant":  "1",
				},
				OwnerName: "tidb1",
			},
			component: &cloudv1.TiKV{
				UseLocalStorage: pointer.Bool(false),
			},
			machineSpec: &cloudv1.MachineSpec{
				MachineType:      "r5.2xlarge",
				MachineImage:     "centos-tidb-v20210420",
				SystemVolumeType: "gp2",
				SystemVolumeSize: "50Gi",
			},
			Zone: "us-west-2a",
			expect: &gardencorev1beta1.Worker{
				Name: "b0c10b54",
				Labels: map[string]string{
					"tenant":                     "1",
					"cluster":                    "1",
					"gardener.wg.name":           "tidb1-tikv-2a",
					"gardener.wg.prefix.name":    "tidb1-tikv",
					"gardener.node.az.info":      "us-west-2a",
					"gardener.node.machine.type": "r5.2xlarge",
				},
				Annotations: map[string]string{},
				Machine: gardencorev1beta1.Machine{
					Type: "r5.2xlarge",
					Image: &gardencorev1beta1.ShootMachineImage{
						Name: "centos-tidb-v20210420",
					},
				},
				Maximum: 1,
				Minimum: 1,
				Taints: []corev1.Taint{{
					Key:    "dedicated",
					Value:  "tidb-tikv",
					Effect: "NoSchedule",
				}},
				Volume: &gardencorev1beta1.Volume{
					Type:       pointer.String("gp2"),
					VolumeSize: "50Gi",
				},
				Zones: []string{"us-west-2a"},
			},
		},
		"cluster1311285575352848384-TiKV": {
			nodeSpec: &NodeTopologySpec{
				Provider: &cloudv1.Provider{
					Spec: cloudv1.ProviderSpec{
						Type: "aws",
					},
				},
				Tags: map[string]string{
					"cluster": "1311285575352848384",
					"tenant":  "1311265642040528896",
					"project": "1311265642040528896",
				},
				OwnerName: "tidb1311285575352848384",
			},
			component: &cloudv1.TiKV{
				UseLocalStorage: pointer.Bool(true),
			},
			machineSpec: &cloudv1.MachineSpec{
				MachineType:         "i3.4xlarge",
				MachineImage:        "centos-tidb-v20210420",
				MachineImageVersion: x86ImageVersion,
				SystemVolumeType:    "gp2",
				SystemVolumeSize:    "50Gi",
			},
			Zone: "ap-northeast-2a",
			expect: &gardencorev1beta1.Worker{
				Name: "2ddb6316-big",
				Labels: map[string]string{
					"gardener.wg.prefix.name":           "tidb1311285575352848384-tikv",
					"gardener.node.machine.type":        "i3.4xlarge",
					"cloud.pingcap.com/node-protection": "true",
					"tenant":                            "1311265642040528896",
					"project":                           "1311265642040528896",
					"gardener.node.az.info":             "ap-northeast-2a",
					"pingcap.com/aws-local-ssd":         "true",
					"cluster":                           "1311285575352848384",
					"gardener.wg.name":                  "tidb1311285575352848384-tikv-2a",
				},
				Annotations: map[string]string{},
				Machine: gardencorev1beta1.Machine{
					Type: "i3.4xlarge",
					Image: &gardencorev1beta1.ShootMachineImage{
						Name:    "centos-tidb-v20210420",
						Version: &x86ImageVersion,
					},
				},
				Maximum: 1,
				Minimum: 1,
				Taints: []corev1.Taint{{
					Key:    "dedicated",
					Value:  "tidb-tikv",
					Effect: "NoSchedule",
				}},
				Volume: &gardencorev1beta1.Volume{
					Type:       pointer.String("gp2"),
					VolumeSize: "50Gi",
				},
				Zones: []string{"ap-northeast-2a"},
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			result := ToGardenerNodeSet((&Builder{}).buildNodeGroup(tc.nodeSpec, tc.component, tc.machineSpec, 1, tc.Zone))
			g.Expect(tc.expect).To(Equal(&result))
		})
	}
}

func TestBuild_NodeGroupDistribution(t *testing.T) {
	cases := map[string]struct {
		desiredReplicas   int
		Zones             []string
		currentSpreads    []CountPerZone
		currentPVCSpreads []CountPerZone
		desiredSpreads    []CountPerZone
		err               bool
	}{
		"newCluster": {
			desiredReplicas: 6,
			Zones:           []string{"a", "b", "c"},
			currentSpreads:  nil,
			desiredSpreads: []CountPerZone{{
				Count: 2,
				Zone:  "a",
			}, {
				Count: 2,
				Zone:  "b",
			}, {
				Count: 2,
				Zone:  "c",
			}},
		},
		"unorderedSpread": {
			desiredReplicas: 1,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{{
				Count: 1,
				Zone:  "b",
			}},
			desiredSpreads: []CountPerZone{{
				Count: 1,
				Zone:  "b",
			}},
		},
		"unevenSpread": {
			desiredReplicas: 18,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{{
				Count: 7,
				Zone:  "a",
			}, {
				Count: 6,
				Zone:  "b",
			}, {
				Count: 5,
				Zone:  "c",
			}},
			desiredSpreads: []CountPerZone{{
				Count: 7,
				Zone:  "a",
			}, {
				Count: 6,
				Zone:  "b",
			}, {
				Count: 5,
				Zone:  "c",
			}},
		},
		"unevenSpreadScaleOutToEven": {
			desiredReplicas: 21,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{{
				Count: 7,
				Zone:  "a",
			}, {
				Count: 6,
				Zone:  "b",
			}, {
				Count: 5,
				Zone:  "c",
			}},
			desiredSpreads: []CountPerZone{{
				Count: 7,
				Zone:  "a",
			}, {
				Count: 7,
				Zone:  "b",
			}, {
				Count: 7,
				Zone:  "c",
			}},
		},
		"unevenSpreadScaleOut": {
			desiredReplicas: 6,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{{
				Count: 3,
				Zone:  "a",
			}},
			desiredSpreads: []CountPerZone{{
				Count: 3,
				Zone:  "a",
			}, {
				Count: 2,
				Zone:  "b",
			}, {
				Count: 1,
				Zone:  "c",
			}},
		},
		"scaleInNodesLessThanCurrentRequirements": {
			desiredReplicas: 3,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{{
				Count: 2,
				Zone:  "a",
			}, {
				Count: 2,
				Zone:  "b",
			}, {
				Count: 2,
				Zone:  "c",
			}},
			err: true,
		},
		"unevenPVCSpreadScaleOut": {
			desiredReplicas: 1,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{
				{
					Count: 0,
					Zone:  "a",
				}, {
					Count: 0,
					Zone:  "b",
				}, {
					Count: 0,
					Zone:  "c",
				},
			},
			currentPVCSpreads: []CountPerZone{
				{
					Count: 0,
					Zone:  "a",
				}, {
					Count: 1,
					Zone:  "b",
				}, {
					Count: 0,
					Zone:  "c",
				},
			},
			desiredSpreads: []CountPerZone{
				{
					Count: 1,
					Zone:  "b",
				},
			},
		},
		"evenSpreadScaleOutPreferPVC": {
			desiredReplicas: 1,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{
				{
					Count: 0,
					Zone:  "a",
				}, {
					Count: 0,
					Zone:  "b",
				}, {
					Count: 0,
					Zone:  "c",
				},
			},
			currentPVCSpreads: []CountPerZone{
				{
					Count: 1,
					Zone:  "b",
				},
			},
			desiredSpreads: []CountPerZone{
				{
					Count: 1,
					Zone:  "b",
				},
			},
		},
		"unevenSpreadScaleOutPreferSkew": {
			desiredReplicas: 4,
			Zones:           []string{"a", "b", "c"},
			currentSpreads: []CountPerZone{
				{
					Count: 0,
					Zone:  "a",
				},
				{
					Count: 2,
					Zone:  "b",
				}, {
					Count: 1,
					Zone:  "c",
				},
			},
			currentPVCSpreads: []CountPerZone{
				{
					Count: 2,
					Zone:  "b",
				}, {
					Count: 3, // pvc count in zone c is greater
					Zone:  "c",
				},
			},
			desiredSpreads: []CountPerZone{
				{
					Count: 1, // scale out in zone a not zone c
					Zone:  "a",
				},
				{
					Count: 2,
					Zone:  "b",
				},
				{
					Count: 1,
					Zone:  "c",
				},
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			spec := &NodeTopologySpec{
				Provider: &cloudv1.Provider{
					Spec: cloudv1.ProviderSpec{
						Type:    "aws",
						Regions: nil,
						InstanceTypes: []cloudv1.InstanceType{
							{
								Name: "test",
							},
						},
					},
				},
				Components: []cloudv1.ClusterComponent{
					&cloudv1.TiDB{
						Replicas:     tc.desiredReplicas,
						InstanceType: "test",
					},
				},
				Zones:     tc.Zones,
				OwnerName: "test",
				CurrentSpreads: map[cloudv1.ComponentType][]CountPerZone{
					cloudv1.ComponentTiDB: tc.currentSpreads,
				},
			}
			if tc.currentPVCSpreads != nil {
				spec.CurrentPVCSpreads = map[cloudv1.ComponentType][]CountPerZone{
					cloudv1.ComponentTiDB: tc.currentPVCSpreads,
				}
				spec.SchedulePreferPVC = true
			}
			builder := &Builder{
				ConfigProvider: gitops.NewHardcode(),
			}
			ngSet, err := builder.Build(spec, false)
			if tc.err {
				g.Expect(err).NotTo(Succeed())
			} else {
				g.Expect(err).To(Succeed())
				var finalSpreads []CountPerZone
				for _, ng := range ngSet.NodeGroups {
					finalSpreads = append(finalSpreads, CountPerZone{
						Count: int(ng.Maximum),
						Zone:  ng.Zones[0],
					})
				}
				g.Expect(finalSpreads).To(ConsistOf(tc.desiredSpreads))
			}
		})
	}
}

func TestBuild_HadResourceTags(t *testing.T) {
	g := NewGomegaWithT(t)
	pdReplica := 3
	tikvReplica := 6
	ticdcReplica := 5
	spec := &NodeTopologySpec{
		Provider: &cloudv1.Provider{
			Spec: cloudv1.ProviderSpec{
				Type:    "aws",
				Regions: nil,
				InstanceTypes: []cloudv1.InstanceType{
					{
						Name: "test",
					},
				},
			},
		},
		Components: []cloudv1.ClusterComponent{
			&cloudv1.TiKV{
				Replicas:     tikvReplica,
				InstanceType: "test",
			},
			&cloudv1.PD{
				Replicas:     pdReplica,
				InstanceType: "test",
			},
			&cloudv1.TiCDC{
				Spec: cloudv1.TiCDCSpec{
					Replicas:     ticdcReplica,
					InstanceType: "test",
				},
			},
		},
		Zones:     []string{"a", "b", "c", "d"},
		OwnerName: "test",
	}
	builder := &Builder{
		ConfigProvider: gitops.NewHardcode(),
	}
	tagsCount := map[string]int{}
	ngSet, err := builder.Build(spec, false)
	g.Expect(err).To(Succeed())
	for _, node := range ngSet.NodeGroups {
		labels := node.Labels
		//t.Log(labels[common.Component], labels[common.UsedBy])
		tagsCount[labels[common.Component]] += 1 //Counting components
		tagsCount[labels[common.UsedBy]] += 1    //Counting Usedby
	}
	zoneSize := len(spec.Zones)
	//t.Log(tagsCount)
	g.Expect(tagsCount["pd"] + tagsCount["ticdc"] + tagsCount["tikv"]).To(Equal(tagsCount["customer-tidb"]))
	g.Expect(tagsCount["pd"]).To(Equal(min(zoneSize, pdReplica)))
	g.Expect(tagsCount["tikv"]).To(Equal(min(zoneSize, tikvReplica)))
	g.Expect(tagsCount["ticdc"]).To(Equal(min(zoneSize, ticdcReplica)))

}
func min(a, b int) int {
	if a > b {
		return b
	}
	return a
}