package nodescorer

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/go-logr/logr"
	. "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/labels"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/query"
	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"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func Test_calculateNodeScore(t *testing.T) {
	type args struct {
		nodeScoreMap *serverlessv1alpha1.NodeScore
		nodeCpuIdle  map[string]float32
		nodeMemIdle  map[string]float32
	}
	tests := []struct {
		name string
		args args
		want map[string]int
	}{
		{
			name: "empty nodeScoreMap",
			args: args{
				nodeScoreMap: &serverlessv1alpha1.NodeScore{},
				nodeCpuIdle:  map[string]float32{"node1": 0.9},
				nodeMemIdle:  map[string]float32{"node1": 0.9},
			},
			want: nil,
		},
		{
			name: "only have cpu",
			args: args{
				nodeScoreMap: &serverlessv1alpha1.NodeScore{Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4},
				}},
				nodeCpuIdle: map[string]float32{"node1": 0.9},
				nodeMemIdle: map[string]float32{"node1": 0.9},
			},
			want: map[string]int{"node1": 90},
		},
		{
			name: "only have mem",
			args: args{
				nodeScoreMap: &serverlessv1alpha1.NodeScore{Spec: serverlessv1alpha1.NodeScoreSpec{
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6},
				}},
				nodeCpuIdle: map[string]float32{"node1": 0.9},
				nodeMemIdle: map[string]float32{"node1": 0.9},
			},
			want: map[string]int{"node1": 89},
		},
		{
			name: "cpu have more nodes then mem",
			args: args{
				nodeScoreMap: &serverlessv1alpha1.NodeScore{Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6},
				}},
				nodeCpuIdle: map[string]float32{"node1": 0.9, "node2": 0.8},
				nodeMemIdle: map[string]float32{"node1": 0.9},
			},
			want: map[string]int{"node1": 90, "node2": 80},
		},
		{
			name: "mem have more nodes then cpu",
			args: args{
				nodeScoreMap: &serverlessv1alpha1.NodeScore{Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6},
				}},
				nodeCpuIdle: map[string]float32{"node1": 0.9},
				nodeMemIdle: map[string]float32{"node1": 0.9, "node2": 0.8},
			},
			want: map[string]int{"node1": 90, "node2": 80},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			got := calculateNodeScore(tt.args.nodeScoreMap, tt.args.nodeCpuIdle, tt.args.nodeMemIdle)
			if len(tt.want) == 0 {
				g.Expect(got).To(BeEmpty())
			} else {
				g.Expect(got).To(Equal(tt.want))
			}
		})
	}
}

func TestNodeScorer_do(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-admin")

	promData := `{"status":"success","data":{"resultType":"matrix","result":[` +
		`{"metric":{"instance":"node1"},"values":[[1667460382,"0.7"]]},` +
		`{"metric":{"instance":"node2"},"values":[[1667460382,"0.8"]]}]}}`
	promDataNode2NeedEvict := `{"status":"success","data":{"resultType":"matrix","result":[` +
		`{"metric":{"instance":"node1"},"values":[[1667460382,"0.7"]]},` +
		`{"metric":{"instance":"node2"},"values":[[1667460382,"0.05"]]}]}}`
	nodeList := []runtime.Object{
		&corev1.Node{
			ObjectMeta: metav1.ObjectMeta{
				Name:   "node1",
				Labels: map[string]string{labels.ServerlessComponentNodeLabelKey: "tidb"},
			},
			Status: corev1.NodeStatus{
				Phase:      corev1.NodeRunning,
				Conditions: []corev1.NodeCondition{{Type: corev1.NodeReady, Status: corev1.ConditionTrue}},
			},
		},
		&corev1.Node{
			ObjectMeta: metav1.ObjectMeta{
				Name:   "node2",
				Labels: map[string]string{labels.ServerlessComponentNodeLabelKey: "tidb"},
			},
			Status: corev1.NodeStatus{
				Phase:      corev1.NodeRunning,
				Conditions: []corev1.NodeCondition{{Type: corev1.NodeReady, Status: corev1.ConditionTrue}},
			},
		},
	}
	now := metav1.Now()

	tests := []struct {
		name     string
		promData string

		drainerOnDemandNode        bool
		rebalanceNode              bool
		notReadyNode               bool
		deletionCandidateOfCANodes bool

		nodeScoreMap *serverlessv1alpha1.NodeScore
		wantStatus   *serverlessv1alpha1.NodeScoreStatus
	}{
		{
			name:     "empty cpu/mem, disable nodescore",
			promData: promData,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{UpdateTime: &now},
		},
		{
			name:     "only cpu",
			promData: promData,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore:  map[string]int{"node1": 70, "node2": 80},
				SchedulerNodeScore: map[string]int{"node1": 70, "node2": 80},
				NodeNeedEvict:      &serverlessv1alpha1.EvictType{},
				UpdateTime:         &now,
			},
		},
		{
			name:     "only mem",
			promData: promData,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore:  map[string]int{"node1": 69, "node2": 80},
				SchedulerNodeScore: map[string]int{"node1": 69, "node2": 80},
				NodeNeedEvict:      &serverlessv1alpha1.EvictType{},
				UpdateTime:         &now,
			},
		},
		{
			promData: promDataNode2NeedEvict,
			name:     "have cpu/mem, but node2 need evict",
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore:  map[string]int{"node1": 70, "node2": 5},
				SchedulerNodeScore: map[string]int{"node1": 70, "node2": 5},
				NodeNeedEvict: &serverlessv1alpha1.EvictType{
					CPU: map[string]float32{"node2": 0.95},
					MEM: map[string]float32{"node2": 0.95},
				},
				UpdateTime: &now,
			},
		},
		{
			promData:            promDataNode2NeedEvict,
			name:                "have cpu/mem and drainerOnDemandNode, but node2 need evict",
			drainerOnDemandNode: true,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore: map[string]int{"node1": 70, "node2": 5},
				SchedulerNodeScore: map[string]int{
					"node1": 70 - DrainerOnDemandNodeScoreDeduct,
					"node2": 5 - DrainerOnDemandNodeScoreDeduct,
				},
				NodeNeedEvict: &serverlessv1alpha1.EvictType{
					CPU: map[string]float32{"node2": 0.95},
					MEM: map[string]float32{"node2": 0.95},
				},
				UpdateTime: &now,
			},
		},
		{
			promData:      promDataNode2NeedEvict,
			name:          "have cpu/mem and rebalanceNode, but node2 need evict",
			rebalanceNode: true,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore: map[string]int{"node1": 70, "node2": 5},
				SchedulerNodeScore: map[string]int{
					"node1": 70 - RebalanceNodeScoreDeduct,
					"node2": 5 - RebalanceNodeScoreDeduct,
				},
				NodeNeedEvict: &serverlessv1alpha1.EvictType{
					CPU: map[string]float32{"node2": 0.95},
					MEM: map[string]float32{"node2": 0.95},
				},
				UpdateTime: &now,
			},
		},
		{
			promData:     promDataNode2NeedEvict,
			name:         "have cpu/mem and notReadyNode, but node2 need evict",
			notReadyNode: true,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore: map[string]int{"node1": 70, "node2": 5},
				SchedulerNodeScore: map[string]int{
					"node1": 70 - NotReadyDeduct,
					"node2": 5 - NotReadyDeduct,
				},
				NodeNeedEvict: &serverlessv1alpha1.EvictType{
					CPU: map[string]float32{"node2": 0.95},
					MEM: map[string]float32{"node2": 0.95},
				},
				UpdateTime: &now,
			},
		},
		{
			promData:                   promDataNode2NeedEvict,
			name:                       "have cpu/mem and deletionCandidateOfCANodes, but node2 need evict",
			deletionCandidateOfCANodes: true,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore: map[string]int{"node1": 70, "node2": 5},
				SchedulerNodeScore: map[string]int{
					"node1": 70 - DeletionCandidateOfCADeduct,
					"node2": 5 - DeletionCandidateOfCADeduct,
				},
				NodeNeedEvict: &serverlessv1alpha1.EvictType{
					CPU: map[string]float32{"node2": 0.95},
					MEM: map[string]float32{"node2": 0.95},
				},
				UpdateTime: &now,
			},
		},
		{
			promData:                   promDataNode2NeedEvict,
			name:                       "have cpu/mem and deletionCandidateOfCANodes and notReadyNode, but node2 need evict",
			deletionCandidateOfCANodes: true,
			notReadyNode:               true,
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
				Spec: serverlessv1alpha1.NodeScoreSpec{
					CPU: &serverlessv1alpha1.CpuProfile{Weight: 4, EvictThreshold: 0.9},
					MEM: &serverlessv1alpha1.MemProfile{Weight: 6, EvictThreshold: 0.9},
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": 70},
					NodeNeedEvict: &serverlessv1alpha1.EvictType{
						CPU: map[string]float32{"node1": 0.95},
						MEM: map[string]float32{"node1": 0.96},
					},
				},
			},
			wantStatus: &serverlessv1alpha1.NodeScoreStatus{
				ResourceNodeScore: map[string]int{"node1": 70, "node2": 5},
				SchedulerNodeScore: map[string]int{
					"node1": 70 - DeletionCandidateOfCADeduct - NotReadyDeduct,
					"node2": 5 - DeletionCandidateOfCADeduct - NotReadyDeduct,
				},
				NodeNeedEvict: &serverlessv1alpha1.EvictType{
					CPU: map[string]float32{"node2": 0.95},
					MEM: map[string]float32{"node2": 0.95},
				},
				UpdateTime: &now,
			},
		},
	}
	var data string
	promServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, data)
	}))
	defer promServer.Close()
	querier := query.NewPromQueryer(nil, promServer.URL)

	ns := &NodeScorer{
		Logger:      logr.New(log.NullLogSink{}),
		ctx:         context.Background(),
		o11yQuerier: querier,
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			data = tt.promData
			copyNodeList := make([]runtime.Object, 0, len(nodeList))
			for _, node := range nodeList {
				copyNode := node.DeepCopyObject().(*corev1.Node)
				if tt.drainerOnDemandNode {
					copyNode.Labels[labels.NodeDrainerCapacityTypeLabelKey] = labels.NodeDrainerCapacityTypeOnDemand
				}
				if tt.rebalanceNode {
					copyNode.Spec.Taints = append(copyNode.Spec.Taints, corev1.Taint{
						Key:    types.RebalanceSignalTaintKey,
						Value:  fmt.Sprint(now.Unix()),
						Effect: corev1.TaintEffectNoSchedule,
					})
				}
				if tt.deletionCandidateOfCANodes {
					copyNode.Spec.Taints = append(copyNode.Spec.Taints, corev1.Taint{
						Key:    DeletionCandidateOfCATaint,
						Value:  fmt.Sprint(now.Unix()),
						Effect: corev1.TaintEffectPreferNoSchedule,
					})
				}
				if tt.notReadyNode {
					for i := range copyNode.Status.Conditions {
						if copyNode.Status.Conditions[i].Type == corev1.NodeReady {
							copyNode.Status.Conditions[i].Status = corev1.ConditionFalse
							break
						}
					}
				}
				copyNodeList = append(copyNodeList, copyNode)
			}
			ns.client = fake.NewClientBuilder().WithScheme(scheme.Scheme).
				WithRuntimeObjects(tt.nodeScoreMap).
				WithRuntimeObjects(copyNodeList...).
				WithStatusSubresource(&serverlessv1alpha1.NodeScore{}).
				Build()
			ns.do()
			gotNodeScoreMap := &serverlessv1alpha1.NodeScore{}
			g.Expect(ns.client.Get(ns.ctx, client.ObjectKeyFromObject(tt.nodeScoreMap), gotNodeScoreMap)).To(Succeed())
			if tt.wantStatus.UpdateTime != nil {
				g.Expect(gotNodeScoreMap.Status.UpdateTime).NotTo(BeNil())
			}
			tt.wantStatus.UpdateTime = nil
			gotNodeScoreMap.Status.UpdateTime = nil
			g.Expect(&gotNodeScoreMap.Status).To(Equal(tt.wantStatus))
		})
	}
}

func Test_filterEvictThreshold(t *testing.T) {
	type args struct {
		m              map[string]float32
		evictThreshold float32
	}
	tests := []struct {
		name string
		args args
		want map[string]float32
	}{
		{
			name: "empty m",
			args: args{
				evictThreshold: 0.9,
			},
			want: nil,
		},
		{
			name: "node2 filter to evict",
			args: args{
				m:              map[string]float32{"node1": 0.6, "node2": 0.025},
				evictThreshold: 0.9,
			},
			want: map[string]float32{"node2": 0.98},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			got := filterEvictThreshold(tt.args.m, tt.args.evictThreshold)
			if len(tt.want) == 0 {
				g.Expect(got).To(BeEmpty())
			} else {
				g.Expect(got).To(Equal(tt.want))
			}
		})
	}
}