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_handleAssignCache(t *testing.T) {
	tests := []struct {
		name                           string
		tidbPool                       *tidbpool.TidbPool
		record                         upgradeRecord
		oldPodsInFreeCache             []client.Object
		oldPodsInAssignCache           []client.Object
		newPodsInFreeCache             []client.Object
		abnormalPodCount               int
		newRevisionFreeCacheCapability int
		wantRemoveSizeFreeCache        int
		wantRemoveSizeAssignCache      int
		expectedTryTickCount           int
		wantErr                        bool
	}{
		{
			name: "no need to update",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  5,
				MaxSize:  6,
			},
			record:               upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 2},
			expectedTryTickCount: 0,
		},
		{
			name: "oldPodInAssignCacheSize increases",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  5,
				MaxSize:  6,
			},
			record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 1},
			oldPodsInFreeCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-free-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-free-cache": ""},
				}},
			},
			newPodsInFreeCache: []client.Object{},
			oldPodsInAssignCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token1"},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-assign-pod2",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token2"},
				}},
			},
			wantRemoveSizeFreeCache:        0,
			wantRemoveSizeAssignCache:      0,
			newRevisionFreeCacheCapability: 0,
			expectedTryTickCount:           0,
		},
		{
			name: "oldPodInAssignCacheSize decreases",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  5,
				MaxSize:  6,
			},
			record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 2},
			oldPodsInFreeCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-free-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-free-cache": ""},
				}},
			},
			oldPodsInAssignCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token1"},
				}},
			},
			wantRemoveSizeFreeCache:        0,
			wantRemoveSizeAssignCache:      0,
			newRevisionFreeCacheCapability: 0,
			expectedTryTickCount:           0,
		},
		{
			name: "only delete old tidb pods in free cache",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  1,
				MaxSize:  6,
			},
			record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 2},
			oldPodsInFreeCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-free-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-free-cache": ""},
				}},
			},
			newPodsInFreeCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "123-free-cache-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-free-cache": ""},
				}},
			},
			oldPodsInAssignCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token1"},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-assign-pod2",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token2"},
				}},
			},
			wantRemoveSizeFreeCache:        1,
			wantRemoveSizeAssignCache:      1,
			newRevisionFreeCacheCapability: 2,
			expectedTryTickCount:           maxTryTickCount,
		},
		{
			name: "min size zero creates a new revision pod and waits before replacing assigned old pod",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  0,
				MaxSize:  1,
			},
			record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 1},
			oldPodsInAssignCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token1"},
				}},
			},
			newRevisionFreeCacheCapability: 1,
			wantRemoveSizeAssignCache:      0,
			expectedTryTickCount:           maxTryTickCount,
			wantErr:                        true,
		},
		{
			name: "oldPodInAssignCacheSize keep stable, new revision unready, skip delete old tidb pool in free cache",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  5,
				MaxSize:  6,
			},
			record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 2},
			oldPodsInFreeCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-free-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-free-cache": ""},
				}},
			},
			oldPodsInAssignCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token1"},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-assign-pod2",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token2"},
				}},
			},
			wantRemoveSizeFreeCache:        0,
			wantRemoveSizeAssignCache:      0,
			newRevisionFreeCacheCapability: 0,
			expectedTryTickCount:           maxTryTickCount,
			wantErr:                        true,
		},
		{
			name: "oldPodInAssignCacheSize keep stable, we can only delete pods in both free cache and assign cache",
			tidbPool: &tidbpool.TidbPool{
				TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					GenerateName: types.TidbServerNamePrefix + "123-",
					Namespace:    "tidb-pool",
					Labels:       map[string]string{"in-free-cache": ""},
				}},
				Revision: "123",
				MinSize:  1,
				MaxSize:  6,
			},
			record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 2},
			oldPodsInFreeCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-free-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-free-cache": ""},
				}},
			},
			oldPodsInAssignCache: []client.Object{
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token1"},
				}},
				&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
					Name:      "old-assign-pod2",
					Namespace: "tidb-pool",
					Labels:    map[string]string{"in-assign-cache": "token2"},
				}},
			},
			newPodsInFreeCache: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      types.TidbServerNamePrefix + "123-free-cache-pod1",
						Namespace: "tidb-pool",
						Labels:    map[string]string{"in-free-cache": ""},
					},
				},
			},
			wantRemoveSizeFreeCache:        1,
			wantRemoveSizeAssignCache:      1,
			newRevisionFreeCacheCapability: 2,
			expectedTryTickCount:           maxTryTickCount,
		},
		// {
		// 	name: "too much abnormal pods",
		// 	tidbPool: &tidbpool.TidbPool{
		// 		TidbPod: &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			GenerateName: types.TidbServerNamePrefix + "123-",
		// 			Namespace:    "tidb-pool",
		// 			Labels:       map[string]string{"in-free-cache": ""},
		// 		}},
		// 		Revision: "123",
		// 		MinSize:  1,
		// 		MaxSize:  6,
		// 	},
		// 	record: upgradeRecord{revision: "123", tryTickCount: maxTryTickCount, preOldPodInAssignCacheSize: 2},
		// 	oldPodsInFreeCache: []client.Object{
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name:      "old-free-pod1",
		// 			Namespace: "tidb-pool",
		// 			Labels:    map[string]string{"in-free-cache": ""},
		// 		}},
		// 	},
		// 	oldPodsInAssignCache: []client.Object{
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name:      types.TidbServerNamePrefix + "321-old-assign-pod1",
		// 			Namespace: "tidb-pool",
		// 			Labels:    map[string]string{"in-assign-cache": "token1"},
		// 		}},
		// 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 			Name:      "old-assign-pod2",
		// 			Namespace: "tidb-pool",
		// 			Labels:    map[string]string{"in-assign-cache": "token2"},
		// 		}},
		// 	},
		// 	newPodsInFreeCache: []client.Object{
		// 		&corev1.Pod{
		// 			ObjectMeta: metav1.ObjectMeta{
		// 				Namespace: "tidb-pool",
		// 				Name:      types.TidbServerNamePrefix + "123-new-free-pod1",
		// 				Labels:    map[string]string{"in-free-cache": ""},
		// 			},
		// 		},
		// 	},
		// 	abnormalPodCount:               types.MaxAbnormalTidbPods + 1,
		// 	wantRemoveSizeFreeCache:        1,
		// 	wantRemoveSizeAssignCache:      0,
		// 	newRevisionFreeCacheCapability: 1,
		// 	expectedTryTickCount:           maxTryTickCount,
		// },
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			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{}
			oldPodsFreeCache := make([]string, 0, len(tt.oldPodsInFreeCache))
			newPodsFreeCache := make([]string, 0, len(tt.newPodsInFreeCache))
			for _, po := range tt.oldPodsInFreeCache {
				oldPodsFreeCache = append(oldPodsFreeCache, po.GetName())
			}
			for _, po := range tt.newPodsInFreeCache {
				newPodsFreeCache = append(newPodsFreeCache, po.GetName())
			}
			oldPodsAssignCache := make([]clientfake.TokenPod, 0, len(tt.oldPodsInAssignCache))
			for _, po := range tt.oldPodsInAssignCache {
				oldPodsAssignCache = append(oldPodsAssignCache, clientfake.TokenPod{
					Token: po.GetLabels()["in-assign-cache"], Pod: po.GetName(),
				})
			}
			fakeTManager.GetOldTidbFromFreeCacheR = oldPodsFreeCache
			fakeTManager.GetNewTidbFromFreeCacheR = newPodsFreeCache
			fakeTManager.FreeCacheSizeR = len(tt.oldPodsInFreeCache) + len(tt.newPodsInFreeCache)
			fakeTManager.OldTidbFromAssignCache = oldPodsAssignCache
			kubeClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).
				WithObjects(tt.oldPodsInFreeCache...).WithObjects(tt.oldPodsInAssignCache...).WithObjects(tt.newPodsInFreeCache...).Build()
			fakeTManager.ReplaceTidbFn = func(ctx context.Context, podName string) error {
				return kubeClient.Delete(ctx, &corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: "tidb-pool",
						Name:      podName,
					},
				})
			}
			fakeTManager.TryCreatePodFn = func(ctx context.Context, pod *corev1.Pod) error {
				return kubeClient.Create(ctx, pod)
			}
			fakeTManager.TidbPool = fakeTidbPool

			tu := &TidbUpgrader{
				Logger:            logr.New(log.NullLogSink{}),
				client:            kubeClient,
				pool:              fakeTidbPool,
				tidbPool:          tt.tidbPool,
				tManager:          fakeTManager,
				record:            tt.record,
				freeCacheSize:     fakeTManager.FreeCacheSize(),
				oldPodInFreeCache: fakeTManager.GetOldTidbFromFreeCache(tt.tidbPool.Revision, false),
			}

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

			var oldPodsFreeCacheCur, newPodsFreeCacheCur []string
			allPodsFreeCacheList := &corev1.PodList{}
			g.Expect(tu.client.List(context.Background(), allPodsFreeCacheList,
				client.InNamespace("tidb-pool"), client.HasLabels{"in-free-cache"})).To(Succeed())
			for _, po := range allPodsFreeCacheList.Items {
				if strings.HasPrefix(po.Name, tt.tidbPool.TidbPod.GenerateName) {
					newPodsFreeCacheCur = append(newPodsFreeCacheCur, po.Name)
				} else {
					oldPodsFreeCacheCur = append(oldPodsFreeCacheCur, po.Name)
				}
			}
			g.Expect(len(newPodsFreeCacheCur)).To(Equal(tt.newRevisionFreeCacheCapability))
			g.Expect(len(tt.oldPodsInFreeCache) - len(oldPodsFreeCacheCur)).To(Equal(tt.wantRemoveSizeFreeCache))

			oldPodsAssignCacheList := &corev1.PodList{}
			g.Expect(tu.client.List(context.Background(), oldPodsAssignCacheList,
				client.InNamespace("tidb-pool"), client.HasLabels{"in-assign-cache"})).To(Succeed())
			g.Expect(len(tt.oldPodsInAssignCache) - len(oldPodsAssignCacheList.Items)).To(Equal(tt.wantRemoveSizeAssignCache))
			g.Expect(tt.expectedTryTickCount).To(Equal(tu.record.tryTickCount))
		})
	}
}