package tidbpodrescheduler

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
)

func newFakeReconciler(objs []client.Object, created, deleted, evicted *[]string) *Reconciler {
	c := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithObjects(objs...).
		WithIndex(&corev1.Pod{}, "spec.nodeName", func(obj client.Object) []string {
			pod := obj.(*corev1.Pod)
			return []string{pod.Spec.NodeName}
		}).
		WithIndex(&corev1.Pod{}, "tidb.rescheduler.pending", func(obj client.Object) []string {
			pod := obj.(*corev1.Pod)
			if v := pod.Annotations[types.AnnoTidbReschedule]; v != "" && pod.Status.Phase == corev1.PodPending {
				return []string{v}
			}
			return nil
		}).
		WithInterceptorFuncs(interceptor.Funcs{
			Create: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.CreateOption) error {
				err := client.Create(ctx, obj, opts...)
				if err == nil {
					if pod, ok := obj.(*corev1.Pod); ok && pod.Labels[types.TidbServerRevision] != "" {
						*created = append(*created, pod.Name)
					}
				}
				return err
			},
			Delete: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.DeleteOption) error {
				err := client.Delete(ctx, obj, opts...)
				if err == nil {
					if pod, ok := obj.(*corev1.Pod); ok && pod.Labels[types.TidbServerRevision] != "" {
						*deleted = append(*deleted, pod.Name)
					}
				}
				return err
			},
			SubResourceCreate: func(ctx context.Context, client client.Client, subResourceName string,
				obj client.Object, subResource client.Object, opts ...client.SubResourceCreateOption,
			) error {
				err := client.SubResource(subResourceName).Create(ctx, obj, subResource, opts...)
				if err == nil && subResourceName == "eviction" {
					*evicted = append(*evicted, obj.GetName())
				}
				return err
			},
		}).
		Build()
	return &Reconciler{Client: c}
}

func readyNode(name string, unsched bool) *corev1.Node {
	return &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: map[string]string{labels.ServerlessComponentNodeLabelKey: "tidb"},
		},
		Spec:   corev1.NodeSpec{Unschedulable: unsched},
		Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{Type: corev1.NodeReady, Status: corev1.ConditionTrue}}},
	}
}

func notReadyNode(name string) *corev1.Node {
	return &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: map[string]string{labels.ServerlessComponentNodeLabelKey: "tidb"},
		},
		Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{Type: corev1.NodeReady, Status: corev1.ConditionFalse}}},
	}
}

//nolint:unparam // ignore
func tidbPod(name, ns, specNodeName, affinityNodeName, revision string, anno map[string]string, phase corev1.PodPhase) *corev1.Pod {
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:        name,
			Namespace:   ns,
			Labels:      map[string]string{types.TidbServerRevision: revision},
			Annotations: anno,
		},
		Spec:   corev1.PodSpec{NodeName: specNodeName},
		Status: corev1.PodStatus{Phase: phase},
	}
	if affinityNodeName != "" {
		pod.Spec.Affinity = &corev1.Affinity{NodeAffinity: &corev1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
				NodeSelectorTerms: []corev1.NodeSelectorTerm{{MatchFields: []corev1.NodeSelectorRequirement{{
					Key:      "metadata.name",
					Operator: corev1.NodeSelectorOpIn,
					Values:   []string{affinityNodeName},
				}}}},
			},
		}}
	}
	return pod
}

func TestReconciler_reschedule(t *testing.T) {
	cases := []struct {
		name      string
		objs      []client.Object
		tManager  *clientfake.FakeTidbManager
		nodeScore serverlessv1alpha1.NodeScoreStatus

		wantDeleted            []string
		wantEvicted            []string
		wantCreatedPodOn       []string
		wantRescheduleFailNode []string
	}{
		{
			name: "no to nodes for unsched",
			objs: []client.Object{
				readyNode("node1", true),
				readyNode("node2", false),
				tidbPod("pod-not-evict", "tidb-pool", "from-node", "from-node", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 3,
				GetTidbFromFreeCacheR: map[string][]string{
					"node1": nil,
					"node2": {"pod1", "pod2", "pod-not-evict"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"node1": 80, "node2": 20},
			},
		},
		{
			name: "no to nodes for not ready",
			objs: []client.Object{
				notReadyNode("node1"),
				readyNode("node2", false),
				tidbPod("pod-not-evict", "tidb-pool", "from-node", "from-node", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 3,
				GetTidbFromFreeCacheR: map[string][]string{
					"node1": nil,
					"node2": {"pod1", "pod2", "pod-not-evict"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"node1": 80, "node2": 20},
			},
		},
		{
			name: "no from nodes",
			objs: []client.Object{
				readyNode("node1", false),
				readyNode("node2", false),
				tidbPod("pod-not-evict", "tidb-pool", "from-node", "from-node", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 3,
				GetTidbFromFreeCacheR: map[string][]string{
					"node1": nil,
					"node2": {"pod1", "pod2", "pod-not-evict"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"node1": 95, "node2": 100},
			},
		},
		{
			name: "evict from-node and create",
			objs: []client.Object{
				readyNode("to-node", false),
				readyNode("middle-node", false),
				readyNode("from-node", false),
				tidbPod("pod-pending-no-delete", "tidb-pool", "from-node", "from-node", "rev",
					map[string]string{types.AnnoTidbReschedule: "true"}, corev1.PodPending),
				tidbPod("pod-evict", "tidb-pool", "from-node", "from-node", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 6,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node":     nil,
					"middle-node": nil,
					"from-node":   {"pod1", "pod2", "pod-evict", "pod-evict-skip"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					AbnormalPods: []corev1.Pod{
						*tidbPod("pod-pending-no-delete", "tidb-pool", "from-node", "from-node", "rev",
							map[string]string{types.AnnoTidbReschedule: "true"}, corev1.PodPending),
						*tidbPod("pod-incoming", "tidb-pool", "to-node", "to-node", "rev", nil, corev1.PodRunning),
					},
					Revision: "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node": 100, "middle-node": 70, "from-node": 60},
				NodeNeedEvict:      &serverlessv1alpha1.EvictType{CPU: map[string]float32{"from-node": 1}},
			},
			wantEvicted:      []string{"pod-evict"},
			wantCreatedPodOn: []string{"to-node"},
		},
		{
			name: "middle-node does not donate free pods",
			objs: []client.Object{
				readyNode("to-node", false),
				readyNode("middle-node", false),
				readyNode("from-node", false),
				tidbPod("pod-evict", "tidb-pool", "middle-node", "middle-node", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 1,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node":     nil,
					"middle-node": {"pod2", "pod-evict"},
					"from-node":   {"pod3"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node": 100, "middle-node": 60, "from-node": 20},
			},
		},
		{
			name: "evict on NodeNeedEvict node",
			objs: []client.Object{
				readyNode("to-node-1", false),
				readyNode("middle-node", false),
				readyNode("from-node-2", false),
				readyNode("from-node-3", false),
				tidbPod("pod-evict", "tidb-pool", "from-node-2", "from-node-2", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 1,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node-1":   nil,
					"middle-node": nil,
					"from-node-2": {"pod1", "pod-evict"},
					"from-node-3": nil,
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node-1": 100, "middle-node": 70, "from-node-3": -50},
				NodeNeedEvict:      &serverlessv1alpha1.EvictType{CPU: map[string]float32{"from-node-2": 1}},
			},
			wantEvicted:      []string{"pod-evict"},
			wantCreatedPodOn: []string{"to-node-1"},
		},
		{
			name: "no free pods to reschedule",
			objs: []client.Object{
				readyNode("to-node", false),
				readyNode("from-node", false),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 2,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node":   nil,
					"from-node": {"pod1"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node": 100, "from-node": 60},
			},
		},
		{
			name: "node is not reschedulable",
			objs: []client.Object{
				readyNode("to-node-1", false),
				readyNode("to-node-2", false),
				readyNode("middle-node", false),
				readyNode("from-node-3", false),
				tidbPod("pod-pending", "tidb-pool", "", "to-node-1", "rev",
					map[string]string{types.AnnoTidbReschedule: "true"}, corev1.PodPending),
				tidbPod("pod-evict", "tidb-pool", "from-node-3", "from-node-3", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 1,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node-1":   nil,
					"to-node-2":   nil,
					"middle-node": nil,
					"from-node-3": {"pod2", "pod-evict", "pod-evict-skip"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node-1": 100, "to-node-2": 90, "middle-node": 80, "from-node-3": 70},
			},
			wantDeleted:            []string{"pod-pending"},
			wantEvicted:            []string{"pod-evict"},
			wantCreatedPodOn:       []string{"to-node-2"},
			wantRescheduleFailNode: []string{"to-node-1"},
		},
		{
			name: "node has enough free pods",
			objs: []client.Object{
				readyNode("to-node-1", false),
				readyNode("to-node-2", false),
				readyNode("middle-node", false),
				readyNode("from-node-3", false),
				tidbPod("pod-evict", "tidb-pool", "from-node-3", "from-node-3", "rev", nil, corev1.PodRunning),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 1,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node-1":   {"pod1"},
					"to-node-2":   nil,
					"middle-node": nil,
					"from-node-3": {"pod2", "pod-evict", "pod-evict-skip"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node-1": 100, "to-node-2": 90, "middle-node": 80, "from-node-3": 70},
			},
			wantEvicted:      []string{"pod-evict"},
			wantCreatedPodOn: []string{"to-node-2"},
		},
		{
			name: "no more free pods to reschedule",
			objs: []client.Object{
				readyNode("to-node-1", false),
				readyNode("to-node-2", false),
				readyNode("middle-node", false),
				readyNode("from-node", false),
				tidbPod("pod-evict", "tidb-pool", "to-node", "to-node", "rev", nil, corev1.PodPending),
			},
			tManager: &clientfake.FakeTidbManager{
				FreeCacheSizeR: 6,
				GetTidbFromFreeCacheR: map[string][]string{
					"to-node-1":   nil,
					"to-node-2":   nil,
					"middle-node": nil,
					"from-node":   {"pod1", "pod2", "pod-evict"},
				},
				TidbPool: &clientfake.FakeTidbPool{
					TidbPod:       &corev1.Pod{},
					PoolNamespace: "tidb-pool",
					Revision:      "rev",
				},
			},
			nodeScore: serverlessv1alpha1.NodeScoreStatus{
				SchedulerNodeScore: map[string]int{"to-node-1": 100, "to-node-2": 90, "middle-node": 80, "from-node": 60},
			},
			wantEvicted:      []string{"pod-evict"},
			wantCreatedPodOn: []string{"to-node-1"},
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			var deleted, evicted, created []string
			r := newFakeReconciler(c.objs, &created, &deleted, &evicted)

			r.reschedule(context.Background(), c.tManager, c.nodeScore)

			g.Expect(deleted).To(ConsistOf(c.wantDeleted))
			g.Expect(evicted).To(ConsistOf(c.wantEvicted))
			g.Expect(len(created)).To(Equal(len(c.wantCreatedPodOn)))
			for i := range created {
				ns := c.tManager.GetTidbPool().Namespace()
				po := &corev1.Pod{}
				g.Expect(r.Client.Get(context.Background(), client.ObjectKey{
					Name:      created[i],
					Namespace: ns,
				}, po)).To(Succeed())
				g.Expect(po.Annotations).To(HaveKey(types.AnnoTidbReschedule))
				g.Expect(po.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.
					NodeSelectorTerms[0].MatchFields[0].Values[0]).To(Equal(c.wantCreatedPodOn[i]))
			}
			for _, nodeName := range c.wantRescheduleFailNode {
				node := &corev1.Node{}
				g.Expect(r.Client.Get(context.Background(), client.ObjectKey{Name: nodeName}, node)).To(Succeed())
				g.Expect(node.Annotations).To(HaveKey(types.AnnoNodeRescheduleFailTimestamp))
			}
		})
	}
}