package nodegroup

import (
	"testing"

	"github.com/stretchr/testify/assert"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

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

func TestBuilder_buildNodeGroupsForSingleAZ(t *testing.T) {
	const (
		compType        = cloudv1.ComponentTiDB
		currentReplicas = 2
	)
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "test",
			Namespace:   "default",
			Annotations: map[string]string{},
		},
		Spec: cloudv1.ClusterSpec{
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     2, // equal to `currentReplicas`
					InstanceType: "r5.xlarge",
				},
			},
		},
	}
	ms := &cloudv1.MachineSpec{
		MachineType: "r5.xlarge",
		MachineArch: "X86",
	}
	type args struct {
		zone            string
		clusterModifier func(cluster *cloudv1.Cluster)
	}
	tests := []struct {
		name string
		args args
		// ng1 is the ng of us-west-2a, ng2 is for us-west-2b
		wantAssertion func(t assert.TestingT, ng1, ng2 cloudv1.NodegroupSpec)
		wantErr       bool
	}{
		{
			name: "scale in should return error",
			args: args{
				zone: "us-west-2a",
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 1
				},
			},
			wantErr: true,
		},
		{
			name: "scale out and without standby node, should increase the min/max of ng1 to 3",
			args: args{
				zone: "us-west-2a",
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 4
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 3, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 3, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
			},
			wantErr: false,
		},
		{
			name: "scale nothing, should keep current spread no change",
			args: args{
				zone: "us-west-2a",
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 1, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
			},
			wantErr: false,
		},
		{
			name: "scale nothing but has 2 standby nodes, should increase the min/max of ng1 to 3",
			args: args{
				zone: "us-west-2a",
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = currentReplicas
					cluster.Annotations[cloudv1.AnnoStandbyNode] = `{"tidb": {"us-west-2a": 2}}`
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 3, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 3, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
			},
			wantErr: false,
		},
		{
			name: "standby node's AZ mismatches with the specific zone, should keep min/max unchanged",
			args: args{
				zone: "us-west-2a",
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = currentReplicas
					cluster.Annotations[cloudv1.AnnoStandbyNode] = `{"tidb": {"us-west-2b": 2}}`
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 1, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			b := &Builder{}
			cluster := cluster.DeepCopy()
			ms := ms.DeepCopy()
			spec := &NodeTopology{
				Cluster: cluster,
				CurrentSpreads: map[cloudv1.ComponentType][]CountPerZone{
					cloudv1.ComponentTiDB: {
						{Zone: "us-west-2a", Count: 1},
						{Zone: "us-west-2b", Count: 1},
					},
				},
			}
			if tt.args.clusterModifier != nil {
				tt.args.clusterModifier(cluster)
			}

			got, err := b.buildNodeGroupsForSingleAZ(spec, tt.args.zone, &cluster.Spec.Cluster.TiDB, currentReplicas, ms)
			if (err != nil) != tt.wantErr {
				t.Errorf("buildNodeGroupsForSingleAZ() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if tt.wantAssertion != nil {
				ng1 := got[1] // us-west-2a
				ng2 := got[0] //us-west-2b
				tt.wantAssertion(t, ng1, ng2)
			}
		})
	}
}

func TestBuilder_buildNodeGroupsPerZone(t *testing.T) {
	const (
		compType        = cloudv1.ComponentTiDB
		currentReplicas = 2
	)
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "test",
			Namespace:   "default",
			Annotations: map[string]string{},
		},
		Spec: cloudv1.ClusterSpec{
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     2, // equal to `currentReplicas`
					InstanceType: "r5.xlarge",
				},
			},
		},
	}
	ms := &cloudv1.MachineSpec{
		MachineType: "r5.xlarge",
		MachineArch: "X86",
	}
	type args struct {
		clusterModifier func(cluster *cloudv1.Cluster)
	}
	tests := []struct {
		name string
		args args
		// ng1 is the ng of us-west-2a, ng2 is for us-west-2b
		wantAssertion func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec)
		wantErr       bool
	}{
		{
			name: "scale in should return error",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 1
				},
			},
			wantErr: true,
		},
		{
			name: "scale out to 3, should add new ng to us-west-2c",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 3
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 1, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng3.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng3.Maximum, "maximum")
			},
			wantErr: false,
		},
		{
			name: "scale out to 4, should add new node to us-west-2a and create new ng for us-west-2c",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 4
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 2, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 2, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng3.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng3.Maximum, "maximum")
			},
			wantErr: false,
		},
		{
			name: "scale out to 3 and have 1 standby node in us-west-2a, should add new ng to us-west-2c and us-west-2a",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 3
					cluster.Annotations[cloudv1.AnnoStandbyNode] = `{"tidb": {"us-west-2a": 1}}`
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 2, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 2, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng3.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng3.Maximum, "maximum")
			},
			wantErr: false,
		}, {
			name: "scale out to 5 and have 1 standby node in us-west-2b, should add new ng to us-west-2c, us-west-2b and us-west-2a",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiDB.Replicas = 5
					cluster.Annotations[cloudv1.AnnoStandbyNode] = `{"tidb": {"us-west-2b": 1}}`
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 2, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 2, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 3, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 3, ng2.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng3.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng3.Maximum, "maximum")
			},
			wantErr: false,
		},
		{
			name: "have 3 standby node in us-west-2b, should change min/max to 4 for ng of us-west-2b",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Annotations[cloudv1.AnnoStandbyNode] = `{"tidb": {"us-west-2b": 3}}`
				},
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 1, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 4, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 4, ng2.Maximum, "maximum")
				assert.EqualValuesf(t, 0, ng3.Minimum, "minimum")
				assert.EqualValuesf(t, 0, ng3.Maximum, "maximum")
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			b := &Builder{}
			cluster := cluster.DeepCopy()
			ms := ms.DeepCopy()
			spec := &NodeTopology{
				Cluster: cluster,
				CurrentSpreads: map[cloudv1.ComponentType][]CountPerZone{
					cloudv1.ComponentTiDB: {
						{Zone: "us-west-2a", Count: 1},
						{Zone: "us-west-2b", Count: 1},
					},
				},
				Zones: []string{"us-west-2a", "us-west-2b", "us-west-2c"},
			}
			if tt.args.clusterModifier != nil {
				tt.args.clusterModifier(cluster)
			}

			got, err := b.buildNodeGroupsPerZone(spec, false, false, &cluster.Spec.Cluster.TiDB, currentReplicas, ms)
			if (err != nil) != tt.wantErr {
				t.Errorf("buildNodeGroupsForSingleAZ() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if tt.wantAssertion != nil {
				ng1 := got[0] // us-west-2a
				ng2 := got[1] //us-west-2b
				var ng3 cloudv1.NodegroupSpec
				if len(got) == 3 {
					ng3 = got[2] // us-west-2c
				}
				tt.wantAssertion(t, ng1, ng2, ng3)
			}
		})
	}
}

func TestBuilder_buildNodeGroupsPerZoneForTiProxy(t *testing.T) {
	const (
		compType = cloudv1.ComponentTiProxy
	)
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "test",
			Namespace:   "default",
			Annotations: map[string]string{},
		},
		Spec: cloudv1.ClusterSpec{
			Cluster: cloudv1.ClusterTopology{
				TiProxy: cloudv1.TiProxy{
					Replicas:     0, // equal to `currentReplicas`
					InstanceType: "r5.xlarge",
				},
			},
		},
	}
	ms := &cloudv1.MachineSpec{
		MachineType: "r5.xlarge",
		MachineArch: "X86",
	}
	type args struct {
		clusterModifier                func(cluster *cloudv1.Cluster)
		tiproxyMixedDeployedWithPD     bool
		componentCurrentActualReplicas int
	}
	tests := []struct {
		name string
		args args
		// ng1 is the ng of us-west-2a, ng2 is for us-west-2b
		wantAssertion    func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec)
		wantErr          bool
		wantNilNodegroup bool
	}{
		{
			name: "tiproxy mixed deploy with PD, should skip nodegroup building",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiProxy.Replicas = 0
				},
				tiproxyMixedDeployedWithPD:     true,
				componentCurrentActualReplicas: 3,
			},
			wantAssertion:    func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {},
			wantErr:          false,
			wantNilNodegroup: true,
		},
		{
			name: "tiproxy deployment is standalone, should build nodegroups",
			args: args{
				clusterModifier: func(cluster *cloudv1.Cluster) {
					cluster.Spec.Cluster.TiProxy.Replicas = 3
				},
				tiproxyMixedDeployedWithPD:     false,
				componentCurrentActualReplicas: 2,
			},
			wantAssertion: func(t assert.TestingT, ng1, ng2, ng3 cloudv1.NodegroupSpec) {
				assert.EqualValuesf(t, 1, ng1.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng1.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng2.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng2.Maximum, "maximum")
				assert.EqualValuesf(t, 1, ng3.Minimum, "minimum")
				assert.EqualValuesf(t, 1, ng3.Maximum, "maximum")
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			b := &Builder{}
			cluster := cluster.DeepCopy()
			ms := ms.DeepCopy()
			spec := &NodeTopology{
				Cluster: cluster,
				CurrentSpreads: map[cloudv1.ComponentType][]CountPerZone{
					cloudv1.ComponentTiProxy: {
						{Zone: "us-west-2a", Count: 1},
						{Zone: "us-west-2b", Count: 1},
					},
				},
				Zones: []string{"us-west-2a", "us-west-2b", "us-west-2c"},
			}
			if tt.args.clusterModifier != nil {
				tt.args.clusterModifier(cluster)
			}

			got, err := b.buildNodeGroupsPerZone(spec, false, tt.args.tiproxyMixedDeployedWithPD, &cluster.Spec.Cluster.TiProxy, tt.args.componentCurrentActualReplicas, ms)
			if (err != nil) != tt.wantErr {
				t.Errorf("buildNodeGroupsPerZone() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if tt.wantNilNodegroup && got == nil {
				return
			}
			if tt.wantAssertion != nil {
				ng1 := got[0] // us-west-2a
				ng2 := got[1] // us-west-2b
				var ng3 cloudv1.NodegroupSpec
				if len(got) == 3 {
					ng3 = got[2] // us-west-2c
				}
				tt.wantAssertion(t, ng1, ng2, ng3)
			}
		})
	}
}