package tidbpodevicter

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"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/query"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/record"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestReconciler_evict(t *testing.T) {
	var dataDeltaDesc, dataUsage string
	promServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Query().Get("query"), "delta") {
			fmt.Fprintln(w, dataDeltaDesc)
		} else {
			fmt.Fprintln(w, dataUsage)
		}
	}))
	defer promServer.Close()
	querier := query.NewPromQueryer(nil, promServer.URL)

	now := time.Now()

	tests := []struct {
		name              string
		promDataDeltaDesc string
		promDataUsage     string
		tidbManagers      map[string]tidbmanager.TidbManager
		nodes             map[string]float32
		objs              []runtime.Object

		wantPods  []string
		wantNodes []corev1.Node
	}{
		{
			name: "test evict",
			nodes: map[string]float32{
				"node_evict": 0.95,
			},
			promDataDeltaDesc: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"vip-tidb-pool","pod":"pod1"},"values":[[1667460382,"0.004"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.003"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod3"},"values":[[1667460382,"0.001"]]}]}}`,
			promDataUsage: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"vip-tidb-pool","pod":"pod1"},"values":[[1667460382,"0.03"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.05"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod3"},"values":[[1667460382,"0.04"]]}]}}`,
			objs: []runtime.Object{
				&corev1.Node{ObjectMeta: metav1.ObjectMeta{
					Name: "node_evict",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod1",
					Namespace: "vip-tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod2",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod3",
					Namespace: "tidb-pool",
				}},
			},
			wantPods: []string{"tidb-pool/pod3"},
			wantNodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_evict",
						Annotations: map[string]string{
							types.AnnoNodeEvictTimestampPrefix + "cpu":   "not-check-value",
							types.AnnoNodeEvictCountPeriodPrefix + "cpu": "1",
						},
					},
				},
			},
		},
		{
			name: "test evict more than maxEvictDuration",
			promDataDeltaDesc: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.003"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.001"]]}]}}`,
			promDataUsage: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.16"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.02"]]}]}}`,
			nodes: map[string]float32{
				"node_evict": 0.95,
			},
			objs: []runtime.Object{
				&corev1.Node{ObjectMeta: metav1.ObjectMeta{
					Name: "node_evict",
					Annotations: map[string]string{
						types.AnnoNodeEvictTimestampPrefix + "cpu":   now.Add(-maxEvictDuration).Format(time.RFC3339),
						types.AnnoNodeEvictCountPeriodPrefix + "cpu": strconv.Itoa(forceEvictCount),
					},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod1",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod2",
					Namespace: "tidb-pool",
				}},
			},
			wantPods: []string{"tidb-pool/pod2"},
			wantNodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_evict",
						Annotations: map[string]string{
							types.AnnoNodeEvictTimestampPrefix + "cpu":   "not-check-value",
							types.AnnoNodeEvictCountPeriodPrefix + "cpu": "1",
						},
					},
				},
			},
		},
		{
			name: "test force evict by evict-count",
			promDataDeltaDesc: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.003"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.001"]]}]}}`,
			promDataUsage: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.16"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.02"]]}]}}`,
			tidbManagers: map[string]tidbmanager.TidbManager{
				"vip-tidb-pool": &clientfake.FakeTidbManager{
					GetTidbFromFreeCacheR: map[string][]string{"node_force_evict": {"vip-pod-evict"}},
				},
				"tidb-pool": &clientfake.FakeTidbManager{
					GetTidbFromFreeCacheR: map[string][]string{"node_force_evict": {"pod-evict"}},
				},
			},
			nodes: map[string]float32{
				"node_force_evict": 0.95,
			},
			objs: []runtime.Object{
				&corev1.Node{ObjectMeta: metav1.ObjectMeta{
					Name: "node_force_evict",
					Annotations: map[string]string{
						types.AnnoNodeEvictTimestampPrefix + "cpu":   now.Format(time.RFC3339),
						types.AnnoNodeEvictCountPeriodPrefix + "cpu": strconv.Itoa(forceEvictCount),
					},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod1",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod2",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "vip-pod-evict",
					Namespace: "vip-tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod-evict",
					Namespace: "tidb-pool",
				}},
			},
			wantPods: []string{"tidb-pool/pod2"},
			wantNodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_force_evict",
						Annotations: map[string]string{
							types.AnnoNodeEvictTimestampPrefix + "cpu":   "not-check-value",
							types.AnnoNodeEvictCountPeriodPrefix + "cpu": strconv.Itoa(forceEvictCount + 1),
							types.AnnoNodeForceEvictTimestamp:            "not-check-value",
						},
					},
					Spec: corev1.NodeSpec{
						Unschedulable: true,
					},
				},
			},
		},
		{
			name: "test force evict again",
			promDataDeltaDesc: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.003"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.001"]]}]}}`,
			promDataUsage: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.16"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.02"]]}]}}`,
			nodes: map[string]float32{
				"node_force_evict_again": 0.95,
			},
			objs: []runtime.Object{
				&corev1.Node{ObjectMeta: metav1.ObjectMeta{
					Name: "node_force_evict_again",
					Annotations: map[string]string{
						types.AnnoNodeEvictTimestampPrefix + "cpu":   now.Format(time.RFC3339),
						types.AnnoNodeEvictCountPeriodPrefix + "cpu": strconv.Itoa(forceEvictCount),
						types.AnnoNodeForceEvictTimestamp:            now.Format(time.RFC3339),
					},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod1",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod2",
					Namespace: "tidb-pool",
				}},
			},
			wantPods: []string{"tidb-pool/pod2"},
			wantNodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_force_evict_again",
						Annotations: map[string]string{
							types.AnnoNodeEvictTimestampPrefix + "cpu":   "not-check-value",
							types.AnnoNodeEvictCountPeriodPrefix + "cpu": strconv.Itoa(forceEvictCount + 1),
							types.AnnoNodeForceEvictTimestamp:            "not-check-value",
						},
					},
				},
			},
		},
		{
			name: "test force evict by multi-nodes",
			promDataDeltaDesc: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.003"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.001"]]}]}}`,
			promDataUsage: `{"status":"success","data":{"resultType":"matrix","result":[` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod1"},"values":[[1667460382,"0.16"]]},` +
				`{"metric":{"namespace":"tidb-pool","pod":"pod2"},"values":[[1667460382,"0.02"]]}]}}`,
			tidbManagers: map[string]tidbmanager.TidbManager{
				"vip-tidb-pool": &clientfake.FakeTidbManager{
					GetTidbFromFreeCacheR: map[string][]string{"node_force_evict1": {"vip-pod-evict"}},
				},
				"tidb-pool": &clientfake.FakeTidbManager{
					GetTidbFromFreeCacheR: map[string][]string{"node_force_evict2": {"pod-evict"}},
				},
			},
			nodes: map[string]float32{
				"node_force_evict1": 0.95,
				"node_force_evict2": 0.98,
			},
			objs: []runtime.Object{
				&corev1.Node{ObjectMeta: metav1.ObjectMeta{
					Name: "node_force_evict1",
				}},
				&corev1.Node{ObjectMeta: metav1.ObjectMeta{
					Name: "node_force_evict2",
					Annotations: map[string]string{
						types.AnnoNodeEvictTimestampPrefix + "cpu":   now.Add(-maxEvictDuration).Format(time.RFC3339),
						types.AnnoNodeEvictCountPeriodPrefix + "cpu": strconv.Itoa(forceEvictCount),
					},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod1",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod2",
					Namespace: "tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "vip-pod-evict",
					Namespace: "vip-tidb-pool",
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "pod-evict",
					Namespace: "tidb-pool",
				}},
			},
			wantPods: []string{"tidb-pool/pod2"},
			wantNodes: []corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_force_evict1",
						Annotations: map[string]string{
							types.AnnoNodeEvictTimestampPrefix + "cpu":   "not-check-value",
							types.AnnoNodeEvictCountPeriodPrefix + "cpu": "1",
							types.AnnoNodeForceEvictTimestamp:            "not-check-value",
						},
					},
					Spec: corev1.NodeSpec{
						Unschedulable: true,
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_force_evict2",
						Annotations: map[string]string{
							types.AnnoNodeEvictTimestampPrefix + "cpu":   "not-check-value",
							types.AnnoNodeEvictCountPeriodPrefix + "cpu": "1",
							types.AnnoNodeForceEvictTimestamp:            "not-check-value",
						},
					},
					Spec: corev1.NodeSpec{
						Unschedulable: true,
					},
				},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			dataDeltaDesc, dataUsage = tt.promDataDeltaDesc, tt.promDataUsage
			r := newFakeReconciler(tt.objs, querier, tt.tidbManagers)
			r.evict(context.Background(), "cpu", 0.9, tt.nodes)

			podList := &corev1.PodList{}
			g.Expect(r.Client.List(context.Background(), podList)).To(Succeed())
			pods := make([]string, 0, len(podList.Items))
			for _, pod := range podList.Items {
				pods = append(pods, pod.Namespace+"/"+pod.Name)
			}
			g.Expect(pods).To(ConsistOf(tt.wantPods))

			nodeList := &corev1.NodeList{}
			g.Expect(r.Client.List(context.Background(), nodeList)).To(Succeed())
			g.Expect(len(nodeList.Items)).To(Equal(len(tt.wantNodes)))
			for _, wantNode := range tt.wantNodes {
				node := findNode(nodeList.Items, wantNode.Name, wantNode.Namespace)
				g.Expect(node).NotTo(BeNil(), "wantNode: %v/%v", wantNode.Namespace, wantNode.Name)
				for k, v := range wantNode.Annotations {
					if v == "not-check-value" {
						g.Expect(node.Annotations[k]).NotTo(BeEmpty(), "node: %v", node.Name)
						node.Annotations[k] = v
					}
				}
				g.Expect(node.Annotations).To(Equal(wantNode.Annotations), "node: %v", node.Name)
				g.Expect(node.Spec).To(Equal(wantNode.Spec), "node: %v", node.Name)
			}
		})
	}
}

func findNode(nodes []corev1.Node, name, namespace string) *corev1.Node {
	for i := range nodes {
		if nodes[i].Name == name && nodes[i].Namespace == namespace {
			return &nodes[i]
		}
	}
	return nil
}

func newFakeReconciler(
	objs []runtime.Object, querier query.PromQuerier, tManager map[string]tidbmanager.TidbManager,
) *Reconciler {
	fakeClient := fake.NewClientBuilder().
		WithScheme(scheme.Scheme).
		WithRuntimeObjects(objs...).
		Build()

	return &Reconciler{
		Client:       fakeClient,
		Recorder:     record.NewFakeRecorder(4096),
		TidbManagers: tManager,
		o11yQuerier:  querier,
	}
}

func Test_calculatePodScoreAndSort(t *testing.T) {
	type args struct {
		podDeltaList []podDelta
		podUsageMap  map[string]float32
	}
	tests := []struct {
		name string
		args args
		want []podScore
	}{
		{
			name: "normal",
			args: args{
				podDeltaList: []podDelta{
					{"ns1", "pod11", 3},
					{"ns1", "pod12", 2},
					{"ns2", "pod21", 0},
					{"ns3", "pod31", -1},
					{"ns3", "pod32", -2},
					{"ns4", "not_found_pod", -2},
				},
				podUsageMap: map[string]float32{
					"ns1/pod11":     0.1,
					"ns1/pod12":     0.1,
					"nsx/pod12":     0.2,
					"ns2/pod21":     0.1,
					"ns3/pod31":     0.2,
					"ns3/pod32":     0.1,
					"nsx/other_pod": 0.1,
				},
			},
			want: []podScore{
				{podDelta{"ns1", "pod11", 3}, 0.1, 58},
				{podDelta{"ns1", "pod12", 2}, 0.1, 42},
				{podDelta{"ns2", "pod21", 0}, 0.1, 5},
				{podDelta{"ns3", "pod31", -1}, 0.2, -6},
				{podDelta{"ns3", "pod32", -2}, 0.1, -27},
				{podDelta{"ns4", "not_found_pod", -2}, 0, -32},
			},
		},
		{
			name: "notPositiveDeltaTotal is zero",
			args: args{
				podDeltaList: []podDelta{
					{"ns1", "pod11", 0},
				},
				podUsageMap: map[string]float32{"ns1/pod11": 0.5},
			},
			want: []podScore{
				{podDelta{"ns1", "pod11", 0}, 0.5, 20},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			g.Expect(calculatePodScoreAndSort(tt.args.podDeltaList, tt.args.podUsageMap)).To(Equal(tt.want))
		})
	}
}