package tidbmanager

import (
	"context"
	"testing"
	"time"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/notifier"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	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_tidbManager_endpointsHandler(t1 *testing.T) {
	tidbPoolNamespace := "tidb-pool"
	tests := []struct {
		name            string
		token           string
		endpoints       string
		lockFreeCache   int64
		pods            []runtime.Object
		freeCache       map[string]podInfo
		assignCache     map[string]tidbInfo
		wantFreeCache   map[string]map[string]freePodInfo
		wantAssignCache map[string]tidbInfo
		wantRemovePod   []string
	}{
		{
			name:      "endpointsHandler",
			token:     "token1",
			endpoints: "pod1,ip1;pod2,ip2",
			pods: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: "pod1", Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						PodIP:      "ip1",
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: "pod2", Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						PodIP:      "ip2",
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: "pod3", Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						PodIP:      "ip3",
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: "podn", Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						PodIP:      "ipn",
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}},
					},
				},
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
				"podn": {podIP: "ipn", nodeName: "node1"},
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod2,ip2", "pod3,ip3"}, replicas: 2},
			},
			wantFreeCache: map[string]map[string]freePodInfo{
				"node1": {
					"podn": freePod("ipn"),
					"pod3": freePod("ip3"),
				},
			},
			wantAssignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}, replicas: 2},
			},
		},
		{
			name:      "endpointsHandler, clean endpoints but pod2 not exist",
			token:     "token1",
			endpoints: "",
			freeCache: map[string]podInfo{},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod2,ip2"}},
			},
			wantFreeCache:   map[string]map[string]freePodInfo{},
			wantAssignCache: map[string]tidbInfo{"token1": {}},
		},
		{
			name:      "endpointsHandler, clean endpoints but pod2 is not ready",
			token:     "token1",
			endpoints: "",
			pods: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: "pod2", Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionFalse}},
					},
				},
			},
			freeCache: map[string]podInfo{},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod2,ip2"}},
			},
			wantFreeCache:   map[string]map[string]freePodInfo{},
			wantAssignCache: map[string]tidbInfo{"token1": {}},
		},
		{
			name:      "endpointsHandler, clean endpoints but pod2 is deleting",
			token:     "token1",
			endpoints: "",
			pods: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:              "pod2",
						Namespace:         tidbPoolNamespace,
						DeletionTimestamp: &metav1.Time{Time: time.Now()},
						Finalizers:        []string{types.TidbPoolFinalizer},
					},
					Spec: corev1.PodSpec{NodeName: "node1"},
				},
			},
			freeCache: map[string]podInfo{},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod2,ip2"}},
			},
			wantFreeCache:   map[string]map[string]freePodInfo{},
			wantAssignCache: map[string]tidbInfo{"token1": {}},
		},
		{
			name:          "endpointsHandler, clean endpoints will remove pod2 because of lockFreeCache",
			token:         "token1",
			endpoints:     "",
			lockFreeCache: 1,
			pods: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{Name: "pod2", Namespace: tidbPoolNamespace},
					Spec:       corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						PodIP:      "ip2",
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}},
					},
				},
			},
			freeCache: map[string]podInfo{},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod2,ip2"}},
			},
			wantFreeCache:   map[string]map[string]freePodInfo{},
			wantAssignCache: map[string]tidbInfo{"token1": {}},
			wantRemovePod:   []string{"pod2"},
		},
		{
			name:      "endpointsHandler, clean endpoints will remove pod2 because of AnnoTidbRemoveAfterRun",
			token:     "token1",
			endpoints: "",
			pods: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "pod2",
						Namespace:   tidbPoolNamespace,
						Annotations: map[string]string{types.AnnoTidbRemoveAfterRun: "true"},
					},
					Spec: corev1.PodSpec{NodeName: "node1"},
					Status: corev1.PodStatus{
						Phase:      corev1.PodRunning,
						PodIP:      "ip2",
						Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: corev1.ConditionTrue}},
					},
				},
			},
			freeCache: map[string]podInfo{},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod2,ip2"}},
			},
			wantFreeCache:   map[string]map[string]freePodInfo{},
			wantAssignCache: map[string]tidbInfo{"token1": {}},
			wantRemovePod:   []string{"pod2"},
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			elected := make(chan struct{})
			close(elected)
			t := &tidbManager{
				Logger:            logr.New(log.NullLogSink{}),
				ctx:               context.Background(),
				tidbPoolNamespace: tidbPoolNamespace,
				lockFreeCache:     tt.lockFreeCache,
				client:            fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(tt.pods...).Build(),
				removePodQueue:    make(chan removePod, 10),
				freeCache:         &freeCache{},
				assignCache:       &assignCache{},
				notifier:          notifier.New(),
				elected:           elected,
			}
			t.freeCache.init(tt.freeCache)
			t.assignCache.init(tt.assignCache)
			ch := t.notifier.Register(tt.token, "test")
			wantNotify := newTidbInfo(tt.endpoints, tt.assignCache[tt.token].replicas, 0)

			t.endpointsHandler(tt.token, tt.endpoints, 0)
			g.Expect(t.freeCache.nodePodMap).To(Equal(tt.wantFreeCache))
			g.Expect(t.assignCache.tokenMap).To(Equal(tt.wantAssignCache))
			msg := <-ch
			endpoints := msg.(*ep.Endpoints)
			g.Expect(endpoints.Endpoints).To(Equal(wantNotify.endpoints))
			g.Expect(endpoints.TidbPool).To(Equal(tidbPoolNamespace))
			g.Expect(endpoints.Comment).To(Equal(wantNotify.comment))
			g.Expect(endpoints.ModRevision).To(Equal(int64(0)))
			g.Expect(endpoints.Replicas).To(Equal(wantNotify.replicas))

			if len(tt.wantRemovePod) > 0 {
				t.processRemovePod(1)
				time.Sleep(1 * time.Second) // wait processRemovePod start worker

				for _, name := range tt.wantRemovePod {
					err := t.client.Get(context.Background(), client.ObjectKey{
						Name:      name,
						Namespace: tidbPoolNamespace,
					}, &corev1.Pod{})
					g.Expect(apierrors.IsNotFound(err)).To(BeTrue())
				}
			}
		})
	}
}

func Test_tidbManager_replicasHandler(t1 *testing.T) {
	tidbPoolNamespace := "tidb-pool"
	tests := []struct {
		name            string
		token           string
		replicas        int
		modRevision     int64
		assignCache     map[string]tidbInfo
		wantAssignCache map[string]tidbInfo
	}{
		{
			name:        "replicasHandler",
			token:       "token1",
			replicas:    2,
			modRevision: 10,
			assignCache: map[string]tidbInfo{
				"token1": {
					endpoints:            []string{"pod1,ip1"},
					replicas:             1,
					comment:              "assigned",
					endpointsModRevision: 5,
				},
			},
			wantAssignCache: map[string]tidbInfo{
				"token1": {
					endpoints:            []string{"pod1,ip1"},
					replicas:             2,
					comment:              "assigned",
					endpointsModRevision: 5,
					replicasModRevision:  10,
				},
			},
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			t := &tidbManager{
				tidbPoolNamespace: tidbPoolNamespace,
				assignCache:       &assignCache{},
				notifier:          notifier.New(),
			}
			t.assignCache.init(tt.assignCache)
			ch := t.notifier.Register(tt.token, "test")

			t.replicasHandler(tt.token, tt.replicas, tt.modRevision)

			g.Expect(t.assignCache.tokenMap).To(Equal(tt.wantAssignCache))
			msg := <-ch
			want := tt.wantAssignCache[tt.token]
			g.Expect(msg).To(Equal(&ep.Endpoints{
				Endpoints:   want.endpoints,
				TidbPool:    tidbPoolNamespace,
				Comment:     want.comment,
				ModRevision: want.modRevision(),
				Replicas:    want.replicas,
			}))
		})
	}
}