package zoneredundantevicter

import (
	"context"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func TestReconciler_evictComponent(t *testing.T) {
	tests := []struct {
		name           string
		evictNodes     []*corev1.Node
		evictPodLabels map[string]string
		objs           []client.Object
		evictErr       error
		wantWait       bool
		wantRetry      bool
	}{
		{
			name: "nothing evict",
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "node1",
						Annotations: map[string]string{annoNodeShouldNoEvicterTaintKey: "1"},
					},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
			},
			objs: []client.Object{
				&corev1.PersistentVolumeClaim{ObjectMeta: metav1.ObjectMeta{
					Name: "pvc-not-evict",
					Annotations: map[string]string{
						types.AnnoVolumeSelectedNode: "node1",
					},
				}},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:   "pod_skip_by_ignored_label",
						Labels: map[string]string{labels.ZREvicterIgnoredKey: "true"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
						Volumes: []corev1.Volume{{
							VolumeSource: corev1.VolumeSource{PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
								ClaimName: "pvc-not-evict",
							}},
						}},
					},
				},
			},
		},
		{
			name: "evict pods, not wait/retry",
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node1"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node2"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node3"},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_ds_not_evict",
						OwnerReferences: []metav1.OwnerReference{{
							Controller: ptr.To(true),
							Kind:       "DaemonSet",
							UID:        "pod_ds_not_evict-uid",
						}},
					},
					Spec: corev1.PodSpec{
						NodeName: "node3",
					},
				},
				&corev1.PersistentVolumeClaim{ObjectMeta: metav1.ObjectMeta{
					Name: "pvc-evict",
					Annotations: map[string]string{
						types.AnnoVolumeSelectedNode: "node1",
						annoObjectShouldBeDeletedKey: "1",
					},
				}},
				&corev1.PersistentVolumeClaim{ObjectMeta: metav1.ObjectMeta{
					Name: "pvc-not-evict",
					Annotations: map[string]string{
						types.AnnoVolumeSelectedNode: "node1",
					},
				}},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "pod-evict-pvc",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
						Labels:      map[string]string{labels.ZREvicterPVCEvictableKey: "true"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
						Volumes: []corev1.Volume{{
							VolumeSource: corev1.VolumeSource{PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
								ClaimName: "pvc-evict",
							}},
						}},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "pod_delete",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
						Volumes: []corev1.Volume{{
							VolumeSource: corev1.VolumeSource{PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
								ClaimName: "pvc-not-evict",
							}},
						}},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "pod_evict",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node3",
					},
				},
			},
		},
		{
			name: "only evict matched labels pods, not wait/retry",
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node1"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node2"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_skip_by_evictPodLabels",
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "pod_delete",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
						Labels: map[string]string{
							"label_key":  "label_evict",
							"label_key2": "label_other",
						},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
			},
			evictPodLabels: map[string]string{"label_key": "label_evict"},
		},
		{
			name: "evict Job pods, not wait/retry",
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node1"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node2"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_job_terminating_not_delete",
						OwnerReferences: []metav1.OwnerReference{{
							Controller: ptr.To(true),
							Kind:       "Job",
							UID:        "pod_job_terminating_not_delete-uid",
						}},
						DeletionTimestamp: &metav1.Time{Time: time.Now()},
						Finalizers:        []string{types.TidbPoolFinalizer},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_job_evict",
						OwnerReferences: []metav1.OwnerReference{{
							Controller: ptr.To(true),
							Kind:       "Job",
							UID:        "pod_job_evict-uid",
						}},
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
					},
				},
			},
		},
		{
			name: "evict Deployment pods, wait, no retry",
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node1"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node2"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node3"},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_deployment_terminating_not_evict",
						OwnerReferences: []metav1.OwnerReference{{
							Controller: ptr.To(true),
							Kind:       "Deployment",
							UID:        "pod_deployment-uid",
						}},
						DeletionTimestamp: &metav1.Time{Time: time.Now()},
						Finalizers:        []string{types.TidbPoolFinalizer},
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_deployment_delete",
						OwnerReferences: []metav1.OwnerReference{{
							Controller: ptr.To(true),
							Kind:       "Deployment",
							UID:        "pod_deployment-uid",
						}},
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_deployment_not_evict",
						OwnerReferences: []metav1.OwnerReference{{
							Controller: ptr.To(true),
							Kind:       "Deployment",
							UID:        "pod_deployment-uid",
						}},
					},
					Spec: corev1.PodSpec{
						NodeName: "node3",
					},
				},
			},
			wantWait: true,
		},
		{
			name: "evict pods, no wait, retry",
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node-evict-first"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node-evict-second"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "pod_not_evict",
					},
					Spec: corev1.PodSpec{
						NodeName: "node-evict-first",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "pod_delete",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node-evict-second",
					},
				},
			},
			evictErr:  apierrors.NewTooManyRequestsError("inject to ut"),
			wantRetry: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			allObjs := make([]client.Object, 0, len(tt.objs)+len(tt.evictNodes))
			allObjs = append(allObjs, tt.objs...)
			for _, n := range tt.evictNodes {
				allObjs = append(allObjs, n)
			}
			r := newFakeReconciler(allObjs, nil, nil, tt.evictErr)
			gotWait, gotRetry, err := r.evictComponent(context.Background(), "test", tt.evictNodes, tt.evictPodLabels)
			g.Expect(err).To(Succeed())
			g.Expect(gotWait).To(Equal(tt.wantWait))
			g.Expect(gotRetry).To(Equal(tt.wantRetry))
			checkNodesAndObjects(g, r.Client, tt.evictNodes, tt.objs)
		})
	}
}