package tidbmanager

import (
	"context"
	"os"
	"testing"
	"time"

	"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/tidb"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

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

	tests := []struct {
		name         string
		newRevision  string
		nodeScoreMap *serverlessv1alpha1.NodeScore
		freeCache    map[string]podInfo
		noOrder      bool
		want         []string
	}{
		{
			name:        "all old revisions in free cache",
			newRevision: "1234567",
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": -20, "node2": 60},
				},
			},
			freeCache: map[string]podInfo{
				"pod1": {podIP: "ip1", nodeName: "node1"},
				types.TidbServerNamePrefix + "7654321-pod1": {podIP: "ip1", nodeName: "node2"},
			},
			want: []string{"pod1", types.TidbServerNamePrefix + "7654321-pod1"},
		},
		{
			name:        "no old revisions in free cache",
			newRevision: "1234567",
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{},
			},
			freeCache: map[string]podInfo{
				types.TidbServerNamePrefix + "1234567-pod1": {podIP: "ip1", nodeName: "node1"},
			},
			want: []string{},
		},
		{
			name:        "got revisions in free cache",
			newRevision: "1234567",
			nodeScoreMap: &serverlessv1alpha1.NodeScore{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.NodeScoreMap,
					Namespace: os.Getenv("NAMESPACE"),
				},
				Status: serverlessv1alpha1.NodeScoreStatus{
					SchedulerNodeScore: map[string]int{"node1": -20, "node2": 60},
				},
			},
			freeCache: map[string]podInfo{
				types.TidbServerNamePrefix + "1234567-pod1": {podIP: "ip1", nodeName: "node1"},
				types.TidbServerNamePrefix + "7654321-pod2": {podIP: "ip2", nodeName: "node1"},
				"pod3": {podIP: "ip3", nodeName: "node1"},
			},
			noOrder: true,
			want:    []string{"pod3", types.TidbServerNamePrefix + "7654321-pod2"},
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			t := &tidbManager{
				Logger:    logr.New(log.NullLogSink{}),
				freeCache: &freeCache{},
				client:    fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(tt.nodeScoreMap).Build(),
			}
			t.freeCache.init(tt.freeCache)

			got := t.GetOldTidbFromFreeCache(tt.newRevision, false)
			if tt.noOrder {
				g.Expect(got).To(ConsistOf(tt.want))
			} else {
				g.Expect(got).To(Equal(tt.want))
			}
		})
	}
}

func Test_tidbManager_GetFreeCachePodsIncludeActivating(t1 *testing.T) {
	t1.Setenv("NAMESPACE", "tidb-admin")
	g := NewGomegaWithT(t1)

	newRevision := "1234567"
	currentPod := types.TidbServerNamePrefix + newRevision + "-pod1"
	oldActivatingPod := types.TidbServerNamePrefix + "7654321-pod2"
	oldAvailablePod := "pod3"
	quarantinedPod := "pod4"
	nodeScoreMap := &serverlessv1alpha1.NodeScore{
		ObjectMeta: metav1.ObjectMeta{
			Name:      types.NodeScoreMap,
			Namespace: os.Getenv("NAMESPACE"),
		},
		Status: serverlessv1alpha1.NodeScoreStatus{
			SchedulerNodeScore: map[string]int{"node1": 10, "node2": 20},
		},
	}
	t := &tidbManager{
		Logger:    logr.New(log.NullLogSink{}),
		freeCache: &freeCache{},
		client:    fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(nodeScoreMap).Build(),
	}
	t.freeCache.init(map[string]podInfo{
		currentPod:       {podIP: "ip1", nodeName: "node1"},
		oldActivatingPod: {podIP: "ip2", nodeName: "node1"},
		oldAvailablePod:  {podIP: "ip3", nodeName: "node2"},
		quarantinedPod:   {podIP: "ip4", nodeName: "node2"},
	})
	t.freeCache.nodePodMap["node1"][currentPod] = activatingFreePod("ip1")
	t.freeCache.nodePodMap["node1"][oldActivatingPod] = activatingFreePod("ip2")
	t.freeCache.nodePodMap["node2"][quarantinedPod] = quarantinedFreePod("ip4")

	g.Expect(t.GetTidbFromFreeCache("", false)).To(ConsistOf(oldAvailablePod))
	g.Expect(t.GetTidbFromFreeCache("", true)).To(ConsistOf(currentPod, oldActivatingPod, oldAvailablePod))
	g.Expect(t.GetTidbFromFreeCacheByRevision(newRevision, false)).To(BeEmpty())
	g.Expect(t.GetTidbFromFreeCacheByRevision(newRevision, true)).To(ConsistOf(currentPod))
	g.Expect(t.GetOldTidbFromFreeCache(newRevision, false)).To(ConsistOf(oldAvailablePod))
	g.Expect(t.GetOldTidbFromFreeCache(newRevision, true)).To(ConsistOf(oldActivatingPod, oldAvailablePod))
	g.Expect(t.FreeCacheSize()).To(Equal(3))
	g.Expect(t.FilterOrphanPods([]string{quarantinedPod, "orphan"})).To(Equal([]string{"orphan"}))
}

func Test_tidbManager_GetOldTidbFromAssignCache(t1 *testing.T) {
	type args struct {
		newRevision string
		limit       int
	}
	tests := []struct {
		name          string
		args          args
		assignCache   map[string]tidbInfo
		wantOldPods   map[string][]string
		wantOldPodsIn map[string][]string
		wantTotal     int
	}{
		{
			name: "limit = 0",
			args: args{newRevision: "1234567", limit: 0},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", types.TidbServerNamePrefix + "1234567-pod2"}},
				"token2": {endpoints: []string{"pod3,ip3"}},
				"token3": {endpoints: []string{types.TidbServerNamePrefix + "7654321-pod4,ip4"}},
			},
			wantOldPods: map[string][]string{},
			wantTotal:   3,
		},
		{
			name: "limit = -1",
			args: args{newRevision: "1234567", limit: -1},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", types.TidbServerNamePrefix + "1234567-pod2"}},
				"token2": {endpoints: []string{"pod3,ip3"}},
				"token3": {endpoints: []string{types.TidbServerNamePrefix + "7654321-pod4,ip4"}},
			},
			wantOldPods: map[string][]string{
				"token1": {"pod1"},
				"token2": {"pod3"},
				"token3": {types.TidbServerNamePrefix + "7654321-pod4"},
			},
			wantTotal: 3,
		},
		{
			name: "limit = 2, limit less than total",
			args: args{newRevision: "1234567", limit: 2},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", types.TidbServerNamePrefix + "1234567-pod2"}},
				"token2": {endpoints: []string{"pod3,ip3"}},
				"token3": {endpoints: []string{types.TidbServerNamePrefix + "7654321-pod4,ip4"}},
			},
			wantOldPodsIn: map[string][]string{
				"token1": {"pod1"},
				"token2": {"pod3"},
				"token3": {types.TidbServerNamePrefix + "7654321-pod4"},
			},
			wantTotal: 3,
		},
		{
			name: "limit = 3, limit less than total, two pods in same token",
			args: args{newRevision: "1234567", limit: 3},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", types.TidbServerNamePrefix + "1234567-pod2"}},
				"token2": {endpoints: []string{"pod3,ip3", "pod4,ip4"}},
				"token3": {endpoints: []string{types.TidbServerNamePrefix + "7654321-pod5,ip5"}},
			},
			wantOldPodsIn: map[string][]string{
				"token1": {"pod1"},
				"token2": {"pod3", "pod4"},
				"token3": {types.TidbServerNamePrefix + "7654321-pod5"},
			},
			wantTotal: 4,
		},
		{
			name: "limit = 5, limit greater than total",
			args: args{newRevision: "1234567", limit: 5},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", types.TidbServerNamePrefix + "1234567-pod2"}},
				"token2": {endpoints: []string{"pod3,ip3"}},
				"token3": {endpoints: []string{types.TidbServerNamePrefix + "7654321-pod4,ip4"}},
			},
			wantOldPods: map[string][]string{
				"token1": {"pod1"},
				"token2": {"pod3"},
				"token3": {types.TidbServerNamePrefix + "7654321-pod4"},
			},
			wantTotal: 3,
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			t := &tidbManager{
				Logger:      logr.New(log.NullLogSink{}),
				assignCache: &assignCache{},
			}
			t.assignCache.init(tt.assignCache)

			oldPods, total := t.GetOldTidbFromAssignCache(tt.args.newRevision, tt.args.limit)
			if tt.wantOldPodsIn != nil {
				for token, pods := range oldPods {
					g.Expect(tt.wantOldPodsIn).Should(HaveKeyWithValue(token, ContainElements(pods)))
				}
			} else {
				g.Expect(oldPods).Should(Equal(tt.wantOldPods))
			}
			if tt.args.limit >= 0 && tt.args.limit <= len(tt.assignCache) {
				oldPodsLen := 0
				for _, pods := range oldPods {
					oldPodsLen += len(pods)
				}
				g.Expect(oldPodsLen).Should(Equal(tt.args.limit))
			}
			g.Expect(total).Should(Equal(tt.wantTotal))
		})
	}
}

type fakeTidbController struct {
	tidb.Controller
	cache *assignCache
}

func (tc *fakeTidbController) Kill(ctx context.Context, host string, keyspaceName string) error {
	// Mock the cache is clear
	if tc.cache != nil {
		tc.cache.removeToken(keyspaceName)
	}
	return nil
}

func Test_tidbManager_KillAll(t1 *testing.T) {
	tests := []struct {
		name           string
		assignCache    map[string]tidbInfo
		token          string
		timeout        time.Duration
		cleanCache     bool
		wantMinElapsed time.Duration
		wantErr        string
	}{
		{
			name: "kill token with pod ip endpoint",
			assignCache: map[string]tidbInfo{
				"token2": {endpoints: []string{"pod3,ip3"}},
			},
			token:      "token2",
			cleanCache: true,
		},
		{
			name: "kill token with mixed endpoint formats",
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", types.TidbServerNamePrefix + "1234567-pod2"}},
			},
			token:      "token1",
			cleanCache: true,
		},
		{
			name: "wait timeout before returning success",
			assignCache: map[string]tidbInfo{
				"token3": {endpoints: []string{types.TidbServerNamePrefix + "7654321-pod4,ip4"}},
			},
			token:          "token3",
			timeout:        50 * time.Millisecond,
			cleanCache:     true,
			wantMinElapsed: 50 * time.Millisecond,
		},
		{
			name: "return error when cache is not cleared",
			assignCache: map[string]tidbInfo{
				"token4": {endpoints: []string{"pod5,ip5"}},
			},
			token:   "token4",
			wantErr: "waiting all tidb to be killed",
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			cache := &assignCache{}
			cache.init(tt.assignCache)
			controller := &fakeTidbController{}
			if tt.cleanCache {
				controller.cache = cache
			}

			mgr := &tidbManager{
				Logger:         logr.New(log.NullLogSink{}),
				assignCache:    cache,
				tidbController: controller,
			}

			start := time.Now()
			err := mgr.KillAll(context.Background(), tt.token, tt.timeout)
			if tt.wantErr == "" {
				g.Expect(err).Should(Succeed())
			} else {
				g.Expect(err).Should(MatchError(tt.wantErr))
			}
			if tt.wantMinElapsed > 0 {
				g.Expect(time.Since(start)).Should(BeNumerically(">=", tt.wantMinElapsed))
			}
		})
	}
}

func Test_tidbManager_GetAssignedTidbByTokenReturnsLatestRevision(t1 *testing.T) {
	g := NewGomegaWithT(t1)
	mgr := &tidbManager{
		tidbPoolNamespace: "tidb-pool",
		assignCache:       &assignCache{},
	}
	mgr.assignCache.init(map[string]tidbInfo{
		"token1": {
			endpoints:            []string{"pod1,ip1"},
			replicas:             1,
			endpointsModRevision: 20,
			replicasModRevision:  10,
		},
		"token2": {
			endpoints:            []string{"pod2,ip2"},
			replicas:             2,
			endpointsModRevision: 10,
			replicasModRevision:  20,
		},
	})

	g.Expect(mgr.GetAssignedTidbByToken("token1").ModRevision).To(Equal(int64(20)))
	g.Expect(mgr.GetAssignedTidbByToken("token2").ModRevision).To(Equal(int64(20)))
}