package tidbpool

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/cmd/manager/app/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

var tidbServerPodTemplateContent = `
apiVersion: v1
kind: Pod
spec:
  containers:
  - name: tidb
`

var updateTidbServerPodTemplateContent = `
apiVersion: v1
kind: Pod
metadata:
  labels:
    go-test: test
spec:
  containers:
  - name: tidb
`

var tidbResourcesContent = `
    limits:
      cpu: "8"
      memory: 1Gi
    requests:
      cpu: 100m
      memory: 512Mi
`

var updateTidbResourcesContent = `
    limits:
      cpu: "4"
      memory: 2Gi
    requests:
      cpu: 200m
      memory: 1Gi
`

func Test_tidbPool_GetPool(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-pool")
	tests := []struct {
		name                    string
		initObjs                []runtime.Object
		updateObjs              []client.Object
		initWant                *TidbPool
		updateWant              *TidbPool
		updateNodeGroupSelector []string
		revisionChange          bool
		haveFakeTidbJob         bool
	}{
		{
			name: "tidbImageKey is empty",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey: "",
					},
				},
			},
		},
		{
			name: "MinSize < 0",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.MinPoolReplicasKey: "-1",
						types.MaxPoolReplicasKey: "1",
					},
				},
			},
		},
		{
			name: "MinSize >= MaxSize",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "1",
					},
				},
			},
		},
		{
			name: "init empty tidb pool",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "0",
						types.MaxPoolReplicasKey: "0",
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			initWant: &TidbPool{},
		},
		{
			name: "init tidb pool, have fakeTidbJob/nodeGroups",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.NodeGroupSelector:  "ng-a , ng-b",
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
				&batchv1.Job{
					ObjectMeta: metav1.ObjectMeta{Name: fakeTidbJobName, Namespace: "tidb-pool", UID: "fake-job-uid"},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			initWant:        &TidbPool{NodeGroups: []string{"ng-a", "ng-b"}, MinSize: 1, MaxSize: 2},
			haveFakeTidbJob: true,
		},
		{
			name: "init tidb pool, update MinSize/MaxSize",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			updateObjs: []client.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "3",
						types.MaxPoolReplicasKey: "5",
					},
				},
			},
			initWant:   &TidbPool{MinSize: 1, MaxSize: 2},
			updateWant: &TidbPool{MinSize: 3, MaxSize: 5},
		},
		{
			name: "init tidb pool, update TidbResourcesKey",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			updateObjs: []client.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   updateTidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
			},
			initWant:       &TidbPool{MinSize: 1, MaxSize: 2},
			updateWant:     &TidbPool{MinSize: 1, MaxSize: 2},
			revisionChange: true,
		},
		{
			name: "init tidb pool, update NodeGroupSelector",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			updateObjs: []client.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.NodeGroupSelector:  " ng-1, ng-2 	",
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
			},
			initWant:                &TidbPool{MinSize: 1, MaxSize: 2},
			updateWant:              &TidbPool{NodeGroups: []string{"ng-1", "ng-2"}, MinSize: 1, MaxSize: 2},
			updateNodeGroupSelector: []string{"ng-1", "ng-2"},
			revisionChange:          true,
		},
		{
			name: "init tidb pool, update tidbServerPodTemplate",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			updateObjs: []client.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: updateTidbServerPodTemplateContent,
					},
				},
			},
			initWant:       &TidbPool{MinSize: 1, MaxSize: 2},
			updateWant:     &TidbPool{MinSize: 1, MaxSize: 2},
			revisionChange: true,
		},
		{
			name: "init tidb pool, update serverlessTidbConfig",
			initObjs: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.ServerlessPoolConfig, Namespace: "tidb-pool"},
					Data: map[string]string{
						types.TidbImageKey:       "tidb-image",
						types.TidbResourcesKey:   tidbResourcesContent,
						types.MinPoolReplicasKey: "1",
						types.MaxPoolReplicasKey: "2",
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: types.TidbServerPodTemplateCm, Namespace: "tidb-pool"},
					Data: map[string]string{
						tidbServerPodTemplate: tidbServerPodTemplateContent,
					},
				},
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "test"},
				},
			},
			updateObjs: []client.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{Name: serverlessTidbConfig, Namespace: "tidb-pool"},
					Data:       map[string]string{"bootstrap-sql-file": "update-test"},
				},
			},
			initWant:       &TidbPool{MinSize: 1, MaxSize: 2},
			updateWant:     &TidbPool{MinSize: 1, MaxSize: 2},
			revisionChange: true,
		},
	}
	cfg := &config.Config{
		TiDBWorkerNS: "tidb-worker",
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			fakeClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(tt.initObjs...).Build()
			tp := New(cfg, "tidb-pool", fakeClient)
			initPool := tp.GetPool(context.Background())
			if tt.initWant == nil {
				g.Expect(initPool).To(BeNil())
			} else {
				g.Expect(initPool.TidbPod).NotTo(BeNil())
				g.Expect(initPool.NodeGroups).To(Equal(tt.initWant.NodeGroups))
				g.Expect(initPool.Revision).NotTo(BeEmpty())
				g.Expect(initPool.MinSize).To(Equal(tt.initWant.MinSize))
				g.Expect(initPool.MaxSize).To(Equal(tt.initWant.MaxSize))
				g.Expect(initPool.TidbPod.Name).To(BeEmpty())
				g.Expect(initPool.TidbPod.Namespace).To(Equal("tidb-pool"))
				g.Expect(initPool.TidbPod.GenerateName).To(
					Equal(types.TidbServerNamePrefix + initPool.Revision + "-"))
				g.Expect(initPool.TidbPod.Labels[types.TidbServerRevision]).To(Equal(initPool.Revision))
				g.Expect(initPool.TidbPod.Labels[labels.NameLabelKey]).To(Equal(labels.TidbPodNameLabelValue))
				g.Expect(initPool.TidbPod.Labels[labels.ComponentLabelKey]).To(Equal(labels.TidbPodComponentLabelValue))
				g.Expect(initPool.TidbPod.Spec.NodeSelector).To(HaveKeyWithValue(labels.ServerlessComponentNodeLabelKey, "tidb"))
				g.Expect(initPool.TidbPod.Spec.Tolerations).To(ContainElement(corev1.Toleration{
					Key:      labels.ServerlessComponentNodeLabelKey,
					Operator: corev1.TolerationOpEqual,
					Effect:   corev1.TaintEffectNoSchedule,
					Value:    "tidb",
				}))
				g.Expect(initPool.TidbPod.Spec.Containers[0].Image).To(Equal("tidb-image"))
				g.Expect(initPool.TidbPod.Spec.Containers[0].Env).To(Equal([]corev1.EnvVar{
					{
						Name:  managerNsEnvName,
						Value: "tidb-pool",
					},
					{
						Name:  tidbWorkerNsEnvName,
						Value: "tidb-worker",
					},
				}))
				g.Expect(initPool.TidbPod.Spec.Containers[0].Resources).NotTo(Equal(corev1.ResourceRequirements{}))
				if tt.haveFakeTidbJob {
					g.Expect(initPool.TidbPod.OwnerReferences).To(Equal([]metav1.OwnerReference{
						{
							APIVersion:         "batch/v1",
							Kind:               "Job",
							Name:               fakeTidbJobName,
							UID:                "fake-job-uid",
							Controller:         ptr.To(true),
							BlockOwnerDeletion: ptr.To(true),
						},
					}))
				} else {
					g.Expect(initPool.TidbPod.OwnerReferences).To(BeEmpty())
				}
			}
			if len(tt.updateObjs) > 0 {
				for _, obj := range tt.updateObjs {
					g.Expect(fakeClient.Update(context.Background(), obj)).To(Succeed())
				}

				updatePool := tp.GetPool(context.Background())
				if tt.updateWant == nil {
					g.Expect(updatePool).To(BeNil())
				} else {
					g.Expect(updatePool.TidbPod).NotTo(BeNil())
					g.Expect(updatePool.NodeGroups).To(Equal(tt.updateWant.NodeGroups))
					g.Expect(updatePool.Revision).NotTo(BeEmpty())
					g.Expect(updatePool.MinSize).To(Equal(tt.updateWant.MinSize))
					g.Expect(updatePool.MaxSize).To(Equal(tt.updateWant.MaxSize))
					g.Expect(updatePool.TidbPod.Name).To(BeEmpty())
					g.Expect(updatePool.TidbPod.Namespace).To(Equal("tidb-pool"))
					g.Expect(updatePool.TidbPod.GenerateName).To(
						Equal(types.TidbServerNamePrefix + updatePool.Revision + "-"))
					g.Expect(updatePool.TidbPod.Labels[types.TidbServerRevision]).To(Equal(updatePool.Revision))
					g.Expect(updatePool.TidbPod.Labels[labels.NameLabelKey]).To(Equal(labels.TidbPodNameLabelValue))
					g.Expect(updatePool.TidbPod.Labels[labels.ComponentLabelKey]).To(Equal(labels.TidbPodComponentLabelValue))
					g.Expect(updatePool.TidbPod.Spec.NodeSelector).To(HaveKeyWithValue(labels.ServerlessComponentNodeLabelKey, "tidb"))
					if len(tt.updateNodeGroupSelector) > 0 {
						g.Expect(updatePool.TidbPod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.
							NodeSelectorTerms).To(ContainElement(corev1.NodeSelectorTerm{MatchExpressions: []corev1.NodeSelectorRequirement{{
							Key:      labels.ServerlessNodeGroupLabelKey,
							Operator: corev1.NodeSelectorOpIn,
							Values:   tt.updateNodeGroupSelector,
						}}}))
					}
					g.Expect(updatePool.TidbPod.Spec.Tolerations).To(ContainElement(corev1.Toleration{
						Key:      labels.ServerlessComponentNodeLabelKey,
						Operator: corev1.TolerationOpEqual,
						Effect:   corev1.TaintEffectNoSchedule,
						Value:    "tidb",
					}))
					g.Expect(updatePool.TidbPod.Spec.Containers[0].Image).To(Equal("tidb-image"))
					g.Expect(updatePool.TidbPod.Spec.Containers[0].Env).To(Equal([]corev1.EnvVar{
						{
							Name:  managerNsEnvName,
							Value: "tidb-pool",
						},
						{
							Name:  tidbWorkerNsEnvName,
							Value: "tidb-worker",
						},
					}))
					g.Expect(updatePool.TidbPod.Spec.Containers[0].Resources).NotTo(Equal(corev1.ResourceRequirements{}))
					if tt.revisionChange {
						g.Expect(updatePool.Revision).NotTo(Equal(initPool.Revision))
					}
				}
			}
		})
	}
}