package cluster

import (
	"testing"
	"time"

	. "github.com/onsi/gomega"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
	evictpod "github.com/tidbcloud/shared-provider/pkg/utils/pod"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func timeToPtr(t metav1.Time) *metav1.Time {
	return &t
}

func TestHandleSafeToEvictAnno(t *testing.T) {
	cases := map[string]struct {
		pods                []*corev1.Pod
		podNum              int
		lastActiveTimestamp *metav1.Time
		safeToEvict         bool
	}{
		"DelSafeToEvictAnno": {
			pods: []*corev1.Pod{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tidb",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "pd",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "pd",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tikv",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tikv",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tiflash",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tiflash",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "discovery",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "discovery",
						},
					},
				},
			},
			podNum:              4,
			lastActiveTimestamp: timeToPtr(metav1.NewTime(time.Now().Add(-30 * time.Minute))),
			safeToEvict:         false,
		},
		"AddSafeToEvictAnno": {
			pods: []*corev1.Pod{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tidb",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "pd",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "pd",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tikv",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tikv",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tiflash",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tiflash",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "discovery",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "discovery",
						},
					},
				},
			},
			podNum:              4,
			lastActiveTimestamp: timeToPtr(metav1.Now()),
			safeToEvict:         true,
		},
		"NoUpdate_HaveSafeToEvict": {
			pods: []*corev1.Pod{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tidb",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "pd",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "pd",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tikv",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tikv",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tiflash",
						Namespace: "test-user-ns",
						Annotations: map[string]string{
							evictpod.SafeToEvictAnno: "false",
						},
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tiflash",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "discovery",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "discovery",
						},
					},
				},
			},
			podNum: 4,
			// lastActiveTimestamp is nil
			safeToEvict: true,
		},
		"NoUpdate_NoSafeToEvict": {
			pods: []*corev1.Pod{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tidb",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "pd",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "pd",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tikv",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tikv",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tiflash",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "tiflash",
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "discovery",
						Namespace: "test-user-ns",
						Labels: map[string]string{
							label.InstanceLabelKey:  "test",
							label.ComponentLabelKey: "discovery",
						},
					},
				},
			},
			podNum:              4,
			lastActiveTimestamp: timeToPtr(metav1.NewTime(time.Now().Add(-30 * time.Minute))),
			safeToEvict:         false,
		},
	}
	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			pods := make([]runtime.Object, 0, len(c.pods))
			for _, pod := range c.pods {
				pods = append(pods, pod)
			}
			cluster := &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{Name: "test", Namespace: defaultNamespace},
				Status:     cloudv1.ClusterStatus{LastActiveTimestamp: c.lastActiveTimestamp},
			}
			r := newFakeReconcileContext(
				cluster,
				newFakeReconciler(pods, []runtime.Object{cluster}))
			g.Expect(r.handleSafeToEvictAnno()).To(Succeed())

			podList := &corev1.PodList{}
			err := r.Client.List(r.ctx, podList, &client.ListOptions{Namespace: r.tcKey.Namespace})
			g.Expect(err).To(Succeed())
			podNum := len(podList.Items)
			g.Expect(podNum).NotTo(BeZero())
			num := 0
			for i := 0; i < podNum; i++ {
				pod := &podList.Items[i]
				if evictpod.EvictComponents.Find(pod.Name) {
					num++
					_, ok := pod.Annotations[evictpod.SafeToEvictAnno]
					g.Expect(ok).To(Equal(c.safeToEvict))
				} else {
					_, ok := pod.Annotations[evictpod.SafeToEvictAnno]
					g.Expect(ok).To(BeFalse())
				}
			}
			g.Expect(num).To(Equal(c.podNum))
		})
	}
}