package tidbmanager

import (
	"context"
	"net/url"
	"os"
	"sync"
	"testing"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/prometheus/client_golang/prometheus/testutil"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	"go.etcd.io/etcd/api/v3/etcdserverpb"
	"go.etcd.io/etcd/api/v3/mvccpb"
	clientv3 "go.etcd.io/etcd/client/v3"
	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/log"
)

type funcActivator func(context.Context) error

func (f funcActivator) do(ctx context.Context, host, token, podName string, param *tidb.ActivateParameter) error {
	return f(ctx)
}

func freePod(podIP string) freePodInfo {
	return freePodInfo{podIP: podIP}
}

func activatingFreePod(podIP string) freePodInfo {
	return freePodInfo{podIP: podIP, state: freePodActivating}
}

func quarantinedFreePod(podIP string) freePodInfo {
	return freePodInfo{podIP: podIP, state: freePodQuarantined}
}

func sequentialActivator(errs ...error) funcActivator {
	var index int
	return func(context.Context) error {
		if index >= len(errs) {
			return nil
		}
		err := errs[index]
		index++
		return err
	}
}

func Test_tidbManager_AssignTidb(t1 *testing.T) {
	t1.Setenv("NAMESPACE", "tidb-admin")

	tidbPoolNamespace := "tidb-pool"
	tidbPodInSameNode := "same-node"
	tests := []struct {
		name           string
		token          string
		getResp        clientv3.GetResponse
		txnResp        clientv3.TxnResponse
		activator      funcActivator
		tidbPool       *clientfake.FakeTidbPool
		poolConfigData map[string]string
		// removeRunSlowAnno    bool
		// test case not put back to free cache
		nodeScoreMap          *serverlessv1alpha1.NodeScore
		freeCache             map[string]podInfo
		assignCache           map[string]tidbInfo
		wantFreeCache         map[string]map[string]freePodInfo
		wantEndpoints         []string
		wantCreatedByTManager bool
		wantErr               error
		wantStandbyChecks     int
		wantTerminatingCount  int
		setupManager          func(*tidbManager)
	}{
		{
			name:          "not found token",
			token:         "token_not_found",
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantErr:       utils.ErrNotFoundToken,
		},
		{
			name:  "tidb replicas is zero",
			token: "token1",
			assignCache: map[string]tidbInfo{
				"token1": {replicas: 0},
			},
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantErr:       utils.ErrZeroReplica,
		},
		{
			name:  "use in assignCache",
			token: "token1",
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}, replicas: 2},
			},
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantEndpoints: []string{"pod1,ip1", "pod2,ip2"},
		},
		{
			name:  "assign cache size limit exceeded",
			token: "token1",
			poolConfigData: map[string]string{
				types.AssignCacheSizeLimitKey: "1",
			},
			assignCache: map[string]tidbInfo{
				"token1": {replicas: 1},
				"token2": {endpoints: []string{"pod2,ip2"}, replicas: 1},
			},
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantErr:       utils.ErrAssignCacheSizeLimitExceeded,
		},
		{
			name:    "not enough pod for assign, but use before cache",
			token:   "token1",
			getResp: clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("pod1,ip1")}}},
			tidbPool: &clientfake.FakeTidbPool{
				TidbPod:       &corev1.Pod{},
				PoolNamespace: tidbPoolNamespace,
				Revision:      "rev",
				MinSize:       1,
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1"}, replicas: 2},
			},
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantEndpoints: []string{"pod1,ip1"},
		},
		{
			name:    "AssignCacheSizeLimitKey missing is unlimited",
			token:   "token1",
			getResp: clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			txnResp: clientv3.TxnResponse{Succeeded: true},
			activator: func(ctx context.Context) error {
				return nil
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
			},
			assignCache: map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{
				"node1": {"pod1": activatingFreePod("ip1")},
			},
			wantEndpoints: []string{"pod1,ip1"},
		},
		{
			name:    "not enough pod for assign, but create a new pod",
			token:   "token1",
			getResp: clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			tidbPool: &clientfake.FakeTidbPool{
				TidbPod:       &corev1.Pod{},
				PoolNamespace: tidbPoolNamespace,
				Revision:      "rev",
				MinSize:       0,
			},
			assignCache:           map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache:         map[string]map[string]freePodInfo{},
			wantCreatedByTManager: true,
			wantErr:               utils.ErrNotEnoughPod,
		},
		{
			name:    "not enough pod for assign, but not create a new pod because of not ready",
			token:   "token1",
			getResp: clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			tidbPool: &clientfake.FakeTidbPool{
				TidbPod:       &corev1.Pod{},
				PoolNamespace: tidbPoolNamespace,
				AbnormalPods: []corev1.Pod{{ObjectMeta: metav1.ObjectMeta{
					Annotations: map[string]string{types.AnnoTidbCreatedByTManager: "token11"},
				}}},
				Revision: "rev",
				MinSize:  0,
			},
			assignCache:   map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantErr:       utils.ErrNotEnoughPod,
		},
		{
			name:      "ActivateByOtherTokenErr, but not enough pod for assign",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			activator: func(ctx context.Context) error { return tidb.ErrActivateByOtherToken },
			tidbPool: &clientfake.FakeTidbPool{
				TidbPod:       &corev1.Pod{},
				PoolNamespace: tidbPoolNamespace,
				Revision:      "rev",
				MinSize:       1,
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: tidbPodInSameNode},
				"pod2": {podIP: "ip2", nodeName: tidbPodInSameNode},
			},
			assignCache:   map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{},
			wantErr:       utils.ErrNotEnoughPod,
		},
		{
			name:      "terminating pod is quarantined and next pod is assigned",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			txnResp:   clientv3.TxnResponse{Succeeded: true},
			activator: sequentialActivator(tidb.ErrActivateTerminating, nil),
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{
						"node1": 80,
						"node2": 0,
					},
				},
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
				"pod2": {podIP: "ip2", nodeName: "node2"},
			},
			assignCache: map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{
				"node1": {"pod1": quarantinedFreePod("ip1")},
				"node2": {"pod2": activatingFreePod("ip2")},
			},
			wantEndpoints:        []string{"pod2,ip2"},
			wantStandbyChecks:    1,
			wantTerminatingCount: 1,
		},
		{
			name:    "terminating pod does not remove same-name replacement",
			token:   "token1",
			getResp: clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			txnResp: clientv3.TxnResponse{Succeeded: true},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
			},
			assignCache: map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{
				"node1": {"pod1": activatingFreePod("ip2")},
			},
			wantEndpoints:        []string{"pod1,ip2"},
			wantStandbyChecks:    1,
			wantTerminatingCount: 1,
			setupManager: func(manager *tidbManager) {
				var calls int
				manager.activator = funcActivator(func(context.Context) error {
					if calls == 0 {
						calls++
						manager.freeCache.setPod("node1", "pod1", "ip2")
						return tidb.ErrActivateTerminating
					}
					return nil
				})
			},
		},
		{
			name:      "Canceled error",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			activator: func(ctx context.Context) error { return &url.Error{Err: context.Canceled} },
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: tidbPodInSameNode},
			},
			assignCache: map[string]tidbInfo{"token1": {replicas: 1}},
			// test case not put back to free cache
			wantFreeCache: map[string]map[string]freePodInfo{
				tidbPodInSameNode: {"pod1": freePod("ip1")},
			},
			wantErr: &url.Error{Err: context.Canceled},
		},
		{
			name:      "assign one critical tidb",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			txnResp:   clientv3.TxnResponse{Succeeded: true},
			activator: func(ctx context.Context) error { return nil },
			poolConfigData: map[string]string{
				types.TidbCriticalKey: "true",
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
			},
			assignCache: map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{
				"node1": {"pod1": activatingFreePod("ip1")},
			},
			wantEndpoints: []string{"pod1,ip1"},
		},
		// {
		// 	name:              "assign one and removeRunSlowAnno",
		// 	token:             "token1",
		// 	getResp:           clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
		// 	txnResp:           clientv3.TxnResponse{Succeeded: true},
		// 	activator:         func(ctx context.Context) error { return nil },
		// 	removeRunSlowAnno: true,
		// 	freeCache: map[string]podInfo{
		// 		"pod1": {podIP: "ip1", nodeName: "node1"},
		// 	},
		// 	assignCache:   map[string]tidbInfo{"token1": {replicas: 1}},
		// 	wantFreeCache: map[string]map[string]string{},
		// 	wantEndpoints: []string{"pod1,ip1"},
		// },
		{
			name:      "assign one on highest score node",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("")}}},
			txnResp:   clientv3.TxnResponse{Succeeded: true},
			activator: func(ctx context.Context) error { return nil },
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{
						"node1": 80,
						"node2": 0,
						"node3": -50,
						"nodex": 90,
					},
				},
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
				"pod2": {podIP: "ip2", nodeName: "node2"},
				"pod3": {podIP: "ip3", nodeName: "node3"},
			},
			assignCache: map[string]tidbInfo{"token1": {replicas: 1}},
			wantFreeCache: map[string]map[string]freePodInfo{
				"node1": {"pod1": activatingFreePod("ip1")},
				"node2": {"pod2": freePod("ip2")},
				"node3": {"pod3": freePod("ip3")},
			},
			wantEndpoints: []string{"pod1,ip1"},
		},
		{
			name:      "assign one on other-node",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("pod1,ip1")}}},
			txnResp:   clientv3.TxnResponse{Succeeded: true},
			activator: func(ctx context.Context) error { return nil },
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{
						tidbPodInSameNode: 20,
						"other-node":      -5,
						"node4":           -40,
					},
				},
			},
			freeCache: map[string]podInfo{
				"pod2": {podIP: "ip2", nodeName: tidbPodInSameNode},
				"pod3": {podIP: "ip3", nodeName: "other-node"},
				"pod4": {podIP: "ip4", nodeName: "node4"},
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1"}, replicas: 2},
			},
			wantFreeCache: map[string]map[string]freePodInfo{
				tidbPodInSameNode: {"pod2": freePod("ip2")},
				"other-node":      {"pod3": activatingFreePod("ip3")},
				"node4":           {"pod4": freePod("ip4")},
			},
			wantEndpoints: []string{"pod1,ip1", "pod3,ip3"},
		},
		{
			name:      "assign one on same-node",
			token:     "token1",
			getResp:   clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("pod1,ip1;pod2,ip2")}}},
			txnResp:   clientv3.TxnResponse{Succeeded: true},
			activator: func(ctx context.Context) error { return nil },
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{
						tidbPodInSameNode: 20,
						"node4":           -90,
					},
				},
			},
			freeCache: map[string]podInfo{
				"pod3": {podIP: "ip3", nodeName: tidbPodInSameNode},
				"pod4": {podIP: "ip4", nodeName: "node4"},
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}, replicas: 3},
			},
			wantFreeCache: map[string]map[string]freePodInfo{
				tidbPodInSameNode: {"pod3": activatingFreePod("ip3")},
				"node4":           {"pod4": freePod("ip4")},
			},
			wantEndpoints: []string{"pod1,ip1", "pod2,ip2", "pod3,ip3"},
		},
		{
			name:    "assign one, txnResp fail",
			token:   "token1",
			getResp: clientv3.GetResponse{Kvs: []*mvccpb.KeyValue{{Value: []byte("pod1,ip1")}}},
			txnResp: clientv3.TxnResponse{
				Succeeded: false,
				Responses: []*etcdserverpb.ResponseOp{{
					Response: &etcdserverpb.ResponseOp_ResponseRange{ResponseRange: &etcdserverpb.RangeResponse{
						Kvs: []*mvccpb.KeyValue{{Value: []byte("pod2,ip2")}},
					}},
				}},
			},
			activator: func(ctx context.Context) error { return nil },
			freeCache: map[string]podInfo{
				"pod2": {podIP: "ip2", nodeName: tidbPodInSameNode},
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1"}, replicas: 2},
			},
			wantFreeCache: map[string]map[string]freePodInfo{
				tidbPodInSameNode: {"pod2": activatingFreePod("ip2")},
			},
			wantEndpoints: []string{"pod2,ip2"},
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			metrics.StandbyQuarantinedPodsGauge.WithLabelValues(tidbPoolNamespace).Set(0)
			terminatingCountBefore := testutil.ToFloat64(
				metrics.ActivateTerminatingCounter.WithLabelValues(tidbPoolNamespace),
			)
			t := &tidbManager{
				Logger:            logr.New(log.NullLogSink{}),
				ctx:               context.Background(),
				safeToEvict:       types.CASafeToEvict,
				tidbPoolNamespace: tidbPoolNamespace,
				etcdKV: clientfake.NewEtcdKV(clientfake.BuildArg{
					GetResponse: &tt.getResp,
					TxnResponse: &tt.txnResp,
				}),
				etcdLease:         clientfake.NewEtcdLease(),
				activator:         tt.activator,
				tidbPool:          tt.tidbPool,
				disLockerFunc:     func(token string) (sync.Locker, func(), error) { return &sync.Mutex{}, func() {}, nil },
				freeCache:         &freeCache{},
				assignCache:       &assignCache{},
				standbyCheckQueue: newStandbyCheckQueue(),
			}
			defer t.standbyCheckQueue.ShutDown()
			b := fake.NewClientBuilder().WithScheme(scheme.Scheme)
			if tt.nodeScoreMap == nil {
				tt.nodeScoreMap = &serverlessv1alpha1.NodeScore{ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				}}
			}
			b.WithRuntimeObjects(tt.nodeScoreMap)
			serverlessPoolConfig := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{
				Name:      types.ServerlessPoolConfig,
				Namespace: tidbPoolNamespace,
			}}
			if tt.poolConfigData != nil {
				serverlessPoolConfig.Data = tt.poolConfigData
			}
			b.WithRuntimeObjects(serverlessPoolConfig)
			for podName, info := range tt.freeCache {
				pod := &corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: podName, Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: info.nodeName},
					Status:     corev1.PodStatus{PodIP: info.podIP},
				}
				// if tt.removeRunSlowAnno {
				// 	if pod.Annotations == nil {
				// 		pod.Annotations = make(map[string]string)
				// 	}
				// 	pod.Annotations[types.AnnoTidbStatusRunSlow] = "true"
				// }
				b = b.WithRuntimeObjects(pod)
			}
			for _, tInfo := range tt.assignCache {
				for _, endpoint := range tInfo.endpoints {
					podName, podIP := ep.ParsePodFromEndpoint(endpoint)
					pod := &corev1.Pod{
						ObjectMeta: metav1.ObjectMeta{Name: podName, Namespace: tidbPoolNamespace},
						Spec:       corev1.PodSpec{NodeName: tidbPodInSameNode},
						Status:     corev1.PodStatus{PodIP: podIP},
					}
					b = b.WithRuntimeObjects(pod)
				}
			}
			t.client = b.Build()

			t.freeCache.init(tt.freeCache)
			t.assignCache.init(tt.assignCache)
			if tt.setupManager != nil {
				tt.setupManager(t)
			}

			gotEndpoints, err := t.AssignTidb(context.Background(), tt.token, &tidb.ActivateParameter{})
			if tt.wantErr == nil {
				g.Expect(err).To(Succeed())
				g.Expect(gotEndpoints).To(Equal(tt.wantEndpoints))
				if tt.poolConfigData[types.TidbCriticalKey] == "true" && len(tt.wantEndpoints) > 0 {
					podName, _ := ep.ParsePodFromEndpoint(tt.wantEndpoints[0])
					pod := &corev1.Pod{}
					g.Expect(t.client.Get(t.ctx, client.ObjectKey{
						Name:      podName,
						Namespace: tidbPoolNamespace,
					}, pod)).To(Succeed())
					g.Expect(pod.Annotations).To(HaveKeyWithValue(t.safeToEvict, "false"))
				}
				// if tt.removeRunSlowAnno && len(tt.wantEndpoints) > 0 {
				// 	for _, endpoint := range tt.wantEndpoints {
				// 		podName, _ := ep.ParsePodFromEndpoint(endpoint)
				// 		pod := &corev1.Pod{}
				// 		g.Expect(t.client.Get(t.ctx, client.ObjectKey{
				// 			Name:      podName,
				// 			Namespace: tidbPoolNamespace,
				// 		}, pod)).To(Succeed())
				// 		g.Expect(pod.Annotations).NotTo(HaveKey(types.AnnoTidbStatusRunSlow))
				// 	}
				// }
			} else {
				g.Expect(err).To(MatchError(tt.wantErr))
			}
			g.Expect(t.freeCache.nodePodMap).To(Equal(tt.wantFreeCache))
			g.Expect(t.standbyCheckQueue.Len()).To(Equal(tt.wantStandbyChecks))
			g.Expect(t.standbyChecks).To(HaveLen(tt.wantStandbyChecks))
			g.Expect(testutil.ToFloat64(
				metrics.StandbyQuarantinedPodsGauge.WithLabelValues(tidbPoolNamespace),
			)).To(Equal(float64(tt.wantStandbyChecks)))
			g.Expect(testutil.ToFloat64(
				metrics.ActivateTerminatingCounter.WithLabelValues(tidbPoolNamespace),
			)).To(Equal(terminatingCountBefore + float64(tt.wantTerminatingCount)))

			tidbPodList := &corev1.PodList{}
			g.Expect(t.client.List(t.ctx, tidbPodList, client.InNamespace(tidbPoolNamespace))).To(Succeed())
			if tt.wantCreatedByTManager {
				var found bool
				for _, pod := range tidbPodList.Items {
					if pod.Annotations[types.AnnoTidbCreatedByTManager] != "" {
						found = true
						break
					}
				}
				g.Expect(found).To(BeTrue(), "want at least one tidb pod has annotation %s", types.AnnoTidbCreatedByTManager)
			} else {
				for _, pod := range tidbPodList.Items {
					g.Expect(pod.Annotations).NotTo(HaveKey(types.AnnoTidbCreatedByTManager))
				}
			}
		})
	}
}