package tidbmanager

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/prometheus/client_golang/prometheus/testutil"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/notifier"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	k8stypes "k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

const (
	standbyTestNamespace = "tidb-pool"
	standbyTestPodName   = "pod1"
	standbyTestNodeName  = "node1"
)

type fakeStandbyStatusController struct {
	tidb.Controller
	status *tidb.Status
	err    error
}

func (c *fakeStandbyStatusController) GetStatus(context.Context, string) (*tidb.Status, error) {
	return c.status, c.err
}

type sequencePodClient struct {
	client.Client
	pods []*corev1.Pod
	gets int
}

func (c *sequencePodClient) Get(
	ctx context.Context,
	key client.ObjectKey,
	obj client.Object,
	opts ...client.GetOption,
) error {
	pod, ok := obj.(*corev1.Pod)
	if !ok || len(c.pods) == 0 {
		return c.Client.Get(ctx, key, obj, opts...)
	}

	index := c.gets
	if index >= len(c.pods) {
		index = len(c.pods) - 1
	}
	c.gets++
	c.pods[index].DeepCopyInto(pod)
	return nil
}

type blockingSecondGetClient struct {
	client.Client
	pod              *corev1.Pod
	gets             int
	secondGetStarted chan struct{}
	allowSecondGet   chan struct{}
}

func (c *blockingSecondGetClient) Get(
	ctx context.Context,
	key client.ObjectKey,
	obj client.Object,
	opts ...client.GetOption,
) error {
	pod, ok := obj.(*corev1.Pod)
	if !ok {
		return c.Client.Get(ctx, key, obj, opts...)
	}

	c.gets++
	if c.gets == 2 {
		close(c.secondGetStarted)
		select {
		case <-c.allowSecondGet:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	c.pod.DeepCopyInto(pod)
	return nil
}

func newStandbyTestPod(podIP string, uid k8stypes.UID, ready bool) *corev1.Pod {
	status := corev1.ConditionFalse
	if ready {
		status = corev1.ConditionTrue
	}
	return &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{Name: standbyTestPodName, Namespace: standbyTestNamespace, UID: uid},
		Spec:       corev1.PodSpec{NodeName: standbyTestNodeName},
		Status: corev1.PodStatus{
			Phase:      corev1.PodRunning,
			PodIP:      podIP,
			Conditions: []corev1.PodCondition{{Type: corev1.PodReady, Status: status}},
		},
	}
}

func TestCheckStandbyAndSetFreeCache(t1 *testing.T) {
	const (
		namespace = standbyTestNamespace
		podName   = standbyTestPodName
		podIP     = "ip1"
		nodeName  = standbyTestNodeName
	)

	readyPod := func(ip string, ready bool) *corev1.Pod {
		return newStandbyTestPod(ip, "uid1", ready)
	}

	tests := []struct {
		name          string
		pods          []runtime.Object
		status        *tidb.Status
		statusErr     error
		freeCache     map[string]podInfo
		assignCache   map[string]tidbInfo
		wantRetry     bool
		wantResult    string
		wantFreeCache map[string]map[string]freePodInfo
	}{
		{
			name:       "standby pod returns to free cache",
			pods:       []runtime.Object{readyPod(podIP, true)},
			status:     &tidb.Status{State: "standby"},
			wantRetry:  false,
			wantResult: standbyCheckResultRestored,
			wantFreeCache: map[string]map[string]freePodInfo{
				nodeName: {podName: freePod(podIP)},
			},
		},
		{
			name:      "terminating pod stays quarantined",
			pods:      []runtime.Object{readyPod(podIP, true)},
			status:    &tidb.Status{State: "terminating", KeyspaceName: "token1"},
			freeCache: map[string]podInfo{podName: {podIP: podIP, nodeName: nodeName}},
			wantRetry: true,
			wantFreeCache: map[string]map[string]freePodInfo{
				nodeName: {podName: quarantinedFreePod(podIP)},
			},
		},
		{
			name:          "status error is retried",
			pods:          []runtime.Object{readyPod(podIP, true)},
			statusErr:     errors.New("status unavailable"),
			wantRetry:     true,
			wantFreeCache: map[string]map[string]freePodInfo{},
		},
		{
			name:          "activated pod is not requeued",
			pods:          []runtime.Object{readyPod(podIP, true)},
			status:        &tidb.Status{State: "activated", KeyspaceName: "token1"},
			wantRetry:     false,
			wantResult:    standbyCheckResultActivated,
			wantFreeCache: map[string]map[string]freePodInfo{},
		},
		{
			name:          "not ready pod is retried",
			pods:          []runtime.Object{readyPod(podIP, false)},
			wantRetry:     true,
			wantFreeCache: map[string]map[string]freePodInfo{},
		},
		{
			name:       "old queue item preserves pod with a new ip",
			pods:       []runtime.Object{readyPod("ip2", true)},
			freeCache:  map[string]podInfo{podName: {podIP: "ip2", nodeName: nodeName}},
			wantRetry:  false,
			wantResult: standbyCheckResultReplaced,
			wantFreeCache: map[string]map[string]freePodInfo{
				nodeName: {podName: freePod("ip2")},
			},
		},
		{
			name: "failed pod is not retried",
			pods: []runtime.Object{func() *corev1.Pod {
				pod := readyPod(podIP, false)
				pod.Status.Phase = corev1.PodFailed
				return pod
			}()},
			wantRetry:     false,
			wantResult:    standbyCheckResultPodGone,
			wantFreeCache: map[string]map[string]freePodInfo{},
		},
		{
			name:          "assigned pod is not requeued",
			pods:          []runtime.Object{readyPod(podIP, true)},
			status:        &tidb.Status{State: "standby"},
			assignCache:   map[string]tidbInfo{"token1": {endpoints: []string{"pod1,ip1"}}},
			wantRetry:     false,
			wantResult:    standbyCheckResultAssigned,
			wantFreeCache: map[string]map[string]freePodInfo{},
		},
	}

	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			manager := &tidbManager{
				Logger:            logr.New(log.NullLogSink{}),
				ctx:               context.Background(),
				client:            fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(tt.pods...).Build(),
				tidbPoolNamespace: namespace,
				tidbController: &fakeStandbyStatusController{
					status: tt.status,
					err:    tt.statusErr,
				},
				freeCache:   &freeCache{},
				assignCache: &assignCache{},
				standbyChecks: map[standbyCheck]struct{}{
					{podName: podName, podIP: podIP}: {},
				},
			}
			manager.freeCache.init(tt.freeCache)
			manager.assignCache.init(tt.assignCache)
			manager.syncStandbyQuarantinedPodsMetricLocked()

			var completedBefore float64
			if tt.wantResult != "" {
				completedBefore = testutil.ToFloat64(
					metrics.StandbyCheckCompletedCounter.WithLabelValues(namespace, tt.wantResult),
				)
			}

			item := standbyCheck{podName: podName, podIP: podIP}
			retry := manager.checkStandbyAndSetFreeCache(
				context.Background(),
				item,
			)

			g.Expect(retry).To(Equal(tt.wantRetry))
			g.Expect(manager.freeCache.nodePodMap).To(Equal(tt.wantFreeCache))
			_, quarantined := manager.standbyChecks[item]
			g.Expect(quarantined).To(Equal(tt.wantRetry))
			if tt.wantRetry {
				g.Expect(manager.GetTidbFromFreeCache("", true)).NotTo(ContainElement(podName))
				g.Expect(manager.FreeCacheSize()).To(Equal(0))
			}
			g.Expect(testutil.ToFloat64(
				metrics.StandbyQuarantinedPodsGauge.WithLabelValues(namespace),
			)).To(Equal(float64(len(manager.standbyChecks))))
			if tt.wantResult != "" {
				g.Expect(testutil.ToFloat64(
					metrics.StandbyCheckCompletedCounter.WithLabelValues(namespace, tt.wantResult),
				)).To(Equal(completedBefore + 1))
			}
		})
	}
}

func TestFinishStandbyCheckMetricsAreIdempotent(t1 *testing.T) {
	g := NewGomegaWithT(t1)
	item := standbyCheck{podName: standbyTestPodName, podIP: "ip1"}
	manager := &tidbManager{
		tidbPoolNamespace: standbyTestNamespace,
		freeCache:         &freeCache{},
		standbyChecks:     map[standbyCheck]struct{}{item: {}},
	}
	manager.freeCache.init(nil)
	manager.syncStandbyQuarantinedPodsMetricLocked()
	completedBefore := testutil.ToFloat64(
		metrics.StandbyCheckCompletedCounter.WithLabelValues(
			standbyTestNamespace,
			standbyCheckResultPodGone,
		),
	)

	manager.finishStandbyCheck(item, standbyCheckResultPodGone)
	manager.finishStandbyCheck(item, standbyCheckResultPodGone)

	g.Expect(testutil.ToFloat64(
		metrics.StandbyQuarantinedPodsGauge.WithLabelValues(standbyTestNamespace),
	)).To(Equal(float64(0)))
	g.Expect(testutil.ToFloat64(
		metrics.StandbyCheckCompletedCounter.WithLabelValues(
			standbyTestNamespace,
			standbyCheckResultPodGone,
		),
	)).To(Equal(completedBefore + 1))
}

func TestEnqueueStandbyCheckKeepsPodQuarantined(t1 *testing.T) {
	g := NewGomegaWithT(t1)
	item := standbyCheck{podName: standbyTestPodName, podIP: "ip1"}
	manager := &tidbManager{
		tidbPoolNamespace: standbyTestNamespace,
		freeCache:         &freeCache{},
		standbyCheckQueue: newStandbyCheckQueue(),
		standbyChecks:     map[standbyCheck]struct{}{},
	}
	defer manager.standbyCheckQueue.ShutDown()
	manager.freeCache.init(map[string]podInfo{
		item.podName: {podIP: item.podIP, nodeName: standbyTestNodeName},
	})
	podName, podIP := manager.freeCache.assignOneOnNode(standbyTestNodeName)
	g.Expect(podName).To(Equal(item.podName))
	g.Expect(podIP).To(Equal(item.podIP))

	manager.enqueueStandbyCheck(item.podName, item.podIP)

	g.Expect(manager.freeCache.nodePodMap).To(Equal(map[string]map[string]freePodInfo{
		standbyTestNodeName: {standbyTestPodName: quarantinedFreePod(item.podIP)},
	}))
	g.Expect(manager.GetTidbFromFreeCache("", false)).To(BeEmpty())
	g.Expect(manager.GetTidbFromFreeCache("", true)).To(BeEmpty())
	g.Expect(manager.FreeCacheSize()).To(Equal(0))
	g.Expect(manager.FilterOrphanPods([]string{item.podName})).To(BeEmpty())
	g.Expect(manager.standbyChecks).To(HaveKey(item))
	g.Expect(manager.standbyCheckQueue.Len()).To(Equal(1))
}

func TestSyncCacheRespectsStandbyQuarantine(t1 *testing.T) {
	const (
		podName  = "pod1"
		oldPodIP = "ip1"
		nodeName = "node1"
	)

	tests := []struct {
		name          string
		podIP         string
		freeCache     map[string]podInfo
		wantFreeCache map[string]map[string]freePodInfo
	}{
		{
			name:      "exact quarantined pod remains quarantined",
			podIP:     oldPodIP,
			freeCache: map[string]podInfo{podName: {podIP: oldPodIP, nodeName: nodeName}},
			wantFreeCache: map[string]map[string]freePodInfo{
				nodeName: {podName: quarantinedFreePod(oldPodIP)},
			},
		},
		{
			name:  "same-name replacement is not blocked",
			podIP: "ip2",
			wantFreeCache: map[string]map[string]freePodInfo{
				nodeName: {podName: freePod("ip2")},
			},
		},
	}

	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			item := standbyCheck{podName: podName, podIP: oldPodIP}
			manager := &tidbManager{
				freeCache:     &freeCache{},
				assignCache:   &assignCache{},
				standbyChecks: map[standbyCheck]struct{}{item: {}},
			}
			manager.freeCache.init(tt.freeCache)
			manager.assignCache.init(nil)

			g.Expect(manager.SyncCache(
				context.Background(),
				nodeName,
				podName,
				tt.podIP,
				utils.PodReady,
			)).To(Succeed())

			g.Expect(manager.freeCache.nodePodMap).To(Equal(tt.wantFreeCache))
			g.Expect(manager.standbyChecks).To(HaveKey(item))
		})
	}
}

func TestStandbyCheckDoesNotAddReplacementPod(t1 *testing.T) {
	const (
		namespace = standbyTestNamespace
		podName   = standbyTestPodName
		nodeName  = standbyTestNodeName
		oldPodIP  = "ip1"
	)

	g := NewGomegaWithT(t1)
	oldPod := newStandbyTestPod(oldPodIP, "uid1", true)
	replacementPod := newStandbyTestPod("ip2", "uid2", true)
	item := standbyCheck{podName: podName, podIP: oldPodIP}
	manager := &tidbManager{
		Logger:            logr.New(log.NullLogSink{}),
		ctx:               context.Background(),
		client:            &sequencePodClient{pods: []*corev1.Pod{oldPod, replacementPod}},
		tidbPoolNamespace: namespace,
		tidbController: &fakeStandbyStatusController{
			status: &tidb.Status{State: "standby"},
		},
		freeCache:     &freeCache{},
		assignCache:   &assignCache{},
		standbyChecks: map[standbyCheck]struct{}{item: {}},
	}
	manager.freeCache.init(nil)
	manager.assignCache.init(nil)

	g.Expect(manager.checkStandbyAndSetFreeCache(context.Background(), item)).To(BeFalse())
	g.Expect(manager.freeCache.nodePodMap).To(BeEmpty())
	g.Expect(manager.standbyChecks).NotTo(HaveKey(item))
	g.Expect(manager.client.(*sequencePodClient).gets).To(Equal(2))
}

func TestStandbyCheckRetriesSameIPReplacement(t1 *testing.T) {
	const oldPodIP = "ip1"

	g := NewGomegaWithT(t1)
	oldPod := newStandbyTestPod(oldPodIP, "uid1", true)
	replacementPod := newStandbyTestPod(oldPodIP, "uid2", true)
	item := standbyCheck{podName: standbyTestPodName, podIP: oldPodIP}
	manager := &tidbManager{
		Logger:            logr.New(log.NullLogSink{}),
		ctx:               context.Background(),
		client:            &sequencePodClient{pods: []*corev1.Pod{oldPod, replacementPod}},
		tidbPoolNamespace: standbyTestNamespace,
		tidbController: &fakeStandbyStatusController{
			status: &tidb.Status{State: "standby"},
		},
		freeCache:     &freeCache{},
		assignCache:   &assignCache{},
		standbyChecks: map[standbyCheck]struct{}{item: {}},
	}
	manager.freeCache.init(nil)
	manager.assignCache.init(nil)

	g.Expect(manager.checkStandbyAndSetFreeCache(context.Background(), item)).To(BeTrue())
	g.Expect(manager.freeCache.nodePodMap).To(BeEmpty())
	g.Expect(manager.standbyChecks).To(HaveKey(item))

	g.Expect(manager.checkStandbyAndSetFreeCache(context.Background(), item)).To(BeFalse())
	g.Expect(manager.freeCache.nodePodMap).To(Equal(map[string]map[string]freePodInfo{
		standbyTestNodeName: {standbyTestPodName: freePod(oldPodIP)},
	}))
	g.Expect(manager.standbyChecks).NotTo(HaveKey(item))
	g.Expect(manager.client.(*sequencePodClient).gets).To(Equal(4))
}

func TestFinalStandbyCheckIsSerializedWithSyncCache(t1 *testing.T) {
	const (
		namespace = standbyTestNamespace
		podName   = standbyTestPodName
		nodeName  = standbyTestNodeName
		podIP     = "ip1"
	)

	for _, podStatus := range []utils.PodStatus{utils.PodNotReady, utils.PodTerminating} {
		t1.Run(podStatus.String(), func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			pod := newStandbyTestPod(podIP, "uid1", true)
			item := standbyCheck{podName: podName, podIP: podIP}
			blockingClient := &blockingSecondGetClient{
				pod:              pod,
				secondGetStarted: make(chan struct{}),
				allowSecondGet:   make(chan struct{}),
			}
			manager := &tidbManager{
				Logger:            logr.New(log.NullLogSink{}),
				ctx:               context.Background(),
				client:            blockingClient,
				tidbPoolNamespace: namespace,
				tidbController: &fakeStandbyStatusController{
					status: &tidb.Status{State: "standby"},
				},
				freeCache:     &freeCache{},
				assignCache:   &assignCache{},
				standbyChecks: map[standbyCheck]struct{}{item: {}},
			}
			manager.freeCache.init(nil)
			manager.assignCache.init(nil)

			standbyDone := make(chan bool, 1)
			go func() {
				standbyDone <- manager.checkStandbyAndSetFreeCache(context.Background(), item)
			}()
			g.Eventually(blockingClient.secondGetStarted, time.Second).Should(BeClosed())

			syncDone := make(chan error, 1)
			go func() {
				syncDone <- manager.SyncCache(context.Background(), nodeName, podName, podIP, podStatus)
			}()
			g.Consistently(syncDone, 100*time.Millisecond).ShouldNot(Receive())

			close(blockingClient.allowSecondGet)

			var retry bool
			g.Eventually(standbyDone, time.Second).Should(Receive(&retry))
			g.Expect(retry).To(BeFalse())
			var syncErr error
			g.Eventually(syncDone, time.Second).Should(Receive(&syncErr))
			g.Expect(syncErr).To(Succeed())
			g.Expect(manager.freeCache.nodePodMap).To(BeEmpty())
			g.Expect(manager.standbyChecks).NotTo(HaveKey(item))
		})
	}
}

func TestFinalStandbyCheckIsSerializedWithAssignmentWatch(t1 *testing.T) {
	const (
		namespace = standbyTestNamespace
		podName   = standbyTestPodName
		nodeName  = standbyTestNodeName
		podIP     = "ip1"
		token     = "token1"
	)

	g := NewGomegaWithT(t1)
	pod := newStandbyTestPod(podIP, "uid1", true)
	item := standbyCheck{podName: podName, podIP: podIP}
	blockingClient := &blockingSecondGetClient{
		pod:              pod,
		secondGetStarted: make(chan struct{}),
		allowSecondGet:   make(chan struct{}),
	}
	manager := &tidbManager{
		Logger:            logr.New(log.NullLogSink{}),
		ctx:               context.Background(),
		client:            blockingClient,
		tidbPoolNamespace: namespace,
		tidbController: &fakeStandbyStatusController{
			status: &tidb.Status{State: "standby"},
		},
		freeCache:     &freeCache{},
		assignCache:   &assignCache{},
		standbyChecks: map[standbyCheck]struct{}{item: {}},
		notifier:      notifier.New(),
	}
	manager.freeCache.init(map[string]podInfo{
		podName: {podIP: podIP, nodeName: nodeName},
	})
	manager.freeCache.quarantinePodIfIP(podName, podIP)
	manager.assignCache.init(map[string]tidbInfo{
		token: {replicas: 1},
	})

	standbyDone := make(chan bool, 1)
	go func() {
		standbyDone <- manager.checkStandbyAndSetFreeCache(context.Background(), item)
	}()
	g.Eventually(blockingClient.secondGetStarted, time.Second).Should(BeClosed())

	watchDone := make(chan struct{})
	go func() {
		manager.endpointsHandler(token, "pod1,ip1", 1)
		close(watchDone)
	}()
	g.Consistently(watchDone, 100*time.Millisecond).ShouldNot(BeClosed())

	close(blockingClient.allowSecondGet)

	var retry bool
	g.Eventually(standbyDone, time.Second).Should(Receive(&retry))
	g.Expect(retry).To(BeFalse())
	g.Eventually(watchDone, time.Second).Should(BeClosed())
	g.Expect(manager.GetAssignedToken(podName, podIP)).To(Equal(token))
	g.Expect(manager.freeCache.nodePodMap).To(BeEmpty())
	g.Expect(manager.standbyChecks).NotTo(HaveKey(item))
}