package tidbupgrader

import (
	"context"
	"strings"
	"testing"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbpool"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	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"
)

func TestTidbUpgrader_handleFreeCache(t *testing.T) {
	tests := []struct {
		name             string
		tidbPool         *tidbpool.TidbPool
		oldPods          []client.Object
		newPods          []client.Object
		abnormalPodCount int
		wantAddSize      int
		wantRemoveSize   int
		wantErr          bool
	}{
		{
			name: "freeCacheSize = 0",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  5,
				MaxSize:  6,
			},
		},
		// {
		// 	name: "< minOldRevPercent, do nothing",
		// 	tidbPool: &tidbpool.TidbPool{
		// 		TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
		// 		Revision: "123",
		// 		MinSize:  1,
		// 		MaxSize:  3,
		// 	},
		// 	oldPods: []client.Object{
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
		// 	},
		// 	newPods: []client.Object{
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod0", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod1", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod2", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod3", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod4", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod5", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod6", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod7", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod8", Namespace: "tidb-pool"}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name: types.TidbServerNamePrefix + "123-new-pod9", Namespace: "tidb-pool"}},
		// 	},
		// },
		{
			name: "freeCacheSize < tidbPool.MinSize",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  5,
				MaxSize:  6,
			},
			oldPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod1", Namespace: "tidb-pool"}},
			},
			newPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod1", Namespace: "tidb-pool"}},
			},
			wantAddSize: 3,
		},
		{
			name: "freeCacheSize in [tidbPool.MinSize, tidbPool.MaxSize], don't add new pod because reach min size",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  2,
				MaxSize:  4,
			},
			oldPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod1", Namespace: "tidb-pool"}},
			},
			newPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod1", Namespace: "tidb-pool"}},
			},
			wantAddSize: 0,
		},
		{
			name: "freeCacheSize > tidbPool.MaxSize and oldPodInFreeCacheSize <= tidbPool.MaxSize, don't add new pod because reach max size",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  2,
				MaxSize:  3,
			},
			oldPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod1", Namespace: "tidb-pool"}},
			},
			newPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod1", Namespace: "tidb-pool"}},
			},
			wantAddSize: 0,
		},
		{
			name: "freeCacheSize > tidbPool.MaxSize and oldPodInFreeCacheSize > tidbPool.MaxSize, remove old pod to reach max size",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  2,
				MaxSize:  3,
			},
			oldPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "456-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod1", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod2", Namespace: "tidb-pool"}},
			},
			newPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod1", Namespace: "tidb-pool"}},
			},
			wantAddSize:    0,
			wantRemoveSize: 1,
		},
		{
			name: "min size zero tolerates limited abnormal pods while upgrading old free cache",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  0,
				MaxSize:  1,
			},
			oldPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
			},
			abnormalPodCount: 1,
		},
		{
			name: "abnormal pods too more",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-", Namespace: "tidb-pool"}},
				Revision: "123",
				MinSize:  2,
				MaxSize:  3,
			},
			oldPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "321-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "456-old-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod1", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: "old-pod2", Namespace: "tidb-pool"}},
			},
			newPods: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod0", Namespace: "tidb-pool"}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name: types.TidbServerNamePrefix + "123-new-pod1", Namespace: "tidb-pool"}},
			},
			abnormalPodCount: types.MaxAbnormalTidbPods + 2,
			wantAddSize:      0,
			wantRemoveSize:   0,
			wantErr:          true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			kubeClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).
				WithObjects(tt.oldPods...).WithObjects(tt.newPods...).Build()
			fakeTidbPool := &clientfake.FakeTidbPool{
				TidbPod:       tt.tidbPool.TidbPod,
				PoolNamespace: "tidb-pool",
				Revision:      tt.tidbPool.Revision,
				MinSize:       tt.tidbPool.MinSize,
				MaxSize:       tt.tidbPool.MaxSize,
				AbnormalPods:  make([]corev1.Pod, tt.abnormalPodCount),
			}

			fakeTManager := &clientfake.FakeTidbManager{}
			oldPods := make([]string, 0, len(tt.oldPods))
			for _, po := range tt.oldPods {
				oldPods = append(oldPods, po.GetName())
			}
			fakeTManager.GetOldTidbFromFreeCacheR = oldPods

			newPods := make([]string, 0, len(tt.newPods))
			for _, po := range tt.newPods {
				newPods = append(newPods, po.GetName())
			}
			fakeTManager.GetNewTidbFromFreeCacheR = newPods
			fakeTManager.TryCreatePodFn = func(ctx context.Context, pod *corev1.Pod) error {
				return kubeClient.Create(ctx, pod)
			}
			fakeTManager.TidbPool = fakeTidbPool

			fakeTManager.FreeCacheSizeR = len(tt.oldPods) + len(tt.newPods)
			tu := &TidbUpgrader{
				Logger:   logr.New(log.NullLogSink{}),
				client:   kubeClient,
				pool:     fakeTidbPool,
				tidbPool: tt.tidbPool,
				tManager: fakeTManager,
			}

			err := tu.handleFreeCache()
			if tt.wantErr {
				g.Expect(err).To(HaveOccurred())
			} else {
				g.Expect(err).NotTo(HaveOccurred())
			}

			var oldPodsCur, newPodsCur []string
			allPodsList := &corev1.PodList{}
			g.Expect(tu.client.List(context.Background(), allPodsList, client.InNamespace("tidb-pool"))).To(Succeed())
			for _, po := range allPodsList.Items {
				if strings.HasPrefix(po.Name, tt.tidbPool.TidbPod.GenerateName) {
					newPodsCur = append(newPodsCur, po.Name)
				} else {
					oldPodsCur = append(oldPodsCur, po.Name)
				}
			}
			g.Expect(len(newPodsCur) - len(tt.newPods)).To(Equal(tt.wantAddSize))
			g.Expect(len(tt.oldPods) - len(oldPodsCur)).To(Equal(tt.wantRemoveSize))
		})
	}
}