package tidbmanager

import (
	"context"
	"net"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-logr/logr"
	perrors "github.com/pingcap/errors"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/cmd/manager/app/config"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"github.com/tidbcloud/aws-shared-provider/pkg/keyspacemigration"
	"github.com/tidbcloud/aws-shared-provider/pkg/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/notifier"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbpool"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/etcdlease"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/kubeapi"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/locker"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	infraclient "github.com/tidbcloud/tidb-management-service/api/serverless_infra/client"
	clientv3 "go.etcd.io/etcd/client/v3"
	"go.etcd.io/etcd/client/v3/concurrency"
	"golang.org/x/sync/errgroup"
	corev1 "k8s.io/api/core/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

var (
	_ TidbManager = &tidbManager{}
)

type tidbManager struct {
	logr.Logger

	ctx context.Context

	etcdKV      clientv3.KV
	etcdLease   clientv3.Lease
	etcdWatcher clientv3.Watcher

	client          client.Client
	infraClusterCli infraclient.Client
	infraApi        infra.InfraAPI

	safeToEvict          string
	tidbPoolNamespace    string
	etcdTokenPrefix      string
	etcdToFreeCacheAgain string

	rateLimiter    *rateLimiter
	activator      activator
	tidbPool       tidbpool.Pool
	tidbScaler     utils.Scaler
	tidbController tidb.Controller

	locLocker     locker.Locker
	disLockerFunc func(token string) (sync.Locker, func(), error)

	tokenRevision            int64
	toFreeCacheAgainRevision int64
	lockFreeCache            int64
	freeCache                *freeCache
	assignCache              *assignCache

	notifier *notifier.Notifier

	migrationLocker keyspacemigration.Locker

	removePodQueue    chan removePod
	standbyCheckQueue standbyCheckQueue
	standbyCheckMu    sync.Mutex
	standbyChecks     map[standbyCheck]struct{}
	elected           <-chan struct{}
	cfg               *config.Config
	creatingPods      sync.Map // key: podName, value: time.Time
}

type TidbManager interface {
	Start(ctx context.Context) error
	AssignTidb(ctx context.Context, token string, param *tidb.ActivateParameter) ([]string, error)
	FreeTidb(ctx context.Context, podName, podIP, comment string) error
	ReplaceTidb(ctx context.Context, podName string, deleteFn func(ctx context.Context, podName, namespace string) error) error
	CheckTidb(ctx context.Context, token, podIP string) error
	SyncCache(ctx context.Context, nodeName, podName, podIP string, podStatus utils.PodStatus) error
	GetTidbPool() tidbpool.Pool
	GetAssignedToken(podName, podIP string) string
	GetTidbFromFreeCache(node string, includeActivating bool) []string

	// GetTidbFromFreeCacheByRevision returns TiDB pods with the specified revision in FREE cache.
	GetTidbFromFreeCacheByRevision(revision string, includeActivating bool) []string
	// GetOldTidbFromFreeCache returns all the OLD revision TiDB pods expect the specified revision in FREE cache.
	GetOldTidbFromFreeCache(newRevision string, includeActivating bool) []string
	// GetOldTidbFromAssignCache returns all the OLD revision TiDB pods expect the specified revision in ASSIGN cache.
	GetOldTidbFromAssignCache(newRevision string, limit int) (map[string][]string, int)

	GetAssignedTidbByToken(token string) *ep.Endpoints
	WalkActiveTokenFromAssignCache(f func(token string, endpoints []string, replicas int) bool)
	LockFreeCache()
	UnlockFreeCache()
	FreeCacheSize() int
	AssignCacheSize() (token, endpoint int)
	FilterOrphanPods(inactivePods []string) []string
	SetTidbScaler(s utils.Scaler)
	GetTidbScaler() utils.Scaler
	DumpStatus() map[string][]any
	KillAll(ctx context.Context, token string, timeout time.Duration) error
	// TryCreatePod attempts to create a new TiDB pod by considering the current creating pods and free cache.
	// It's designed to be called by TidbScaler and TidbUpgrader.
	TryCreatePod(ctx context.Context) error
}

func New(
	mgr manager.Manager,
	safeToEvict string,
	tidbPoolNamespace string,
	etcdClient *clientv3.Client,
	etcdPrefix string,
	notifier *notifier.Notifier,
	tidbController tidb.Controller,
	infraClusterCli infraclient.Client,
	infraApi infra.InfraAPI,
	cfg *config.Config,
	migrationLocker keyspacemigration.Locker,
) (TidbManager, error) {
	log := ctrl.Log.WithName("tidbmanager").WithName(tidbPoolNamespace)
	tManager := &tidbManager{
		Logger: log,

		client:      mgr.GetClient(),
		etcdKV:      etcdClient.KV,
		etcdLease:   etcdClient.Lease,
		etcdWatcher: etcdClient.Watcher,

		infraClusterCli: infraClusterCli,
		infraApi:        infraApi,

		safeToEvict:          safeToEvict,
		tidbPoolNamespace:    tidbPoolNamespace,
		etcdTokenPrefix:      etcdPrefix + "token/",
		etcdToFreeCacheAgain: etcdPrefix + "to-freecache-again",

		activator:      newActivator(log, tidbPoolNamespace, tidbController),
		rateLimiter:    newRateLimiter(log, tidbPoolNamespace, mgr.GetClient()),
		tidbPool:       tidbpool.New(cfg, tidbPoolNamespace, mgr.GetClient()),
		tidbController: tidbController,

		locLocker: locker.New(),
		disLockerFunc: func(token string) (sync.Locker, func(), error) {
			s, err := concurrency.NewSession(etcdClient)
			if err != nil {
				return nil, nil, err
			}
			return concurrency.NewLocker(s, etcdPrefix+"lock/"+token), func() { _ = s.Close() }, nil
		},

		freeCache:   &freeCache{},
		assignCache: &assignCache{},

		notifier: notifier,

		migrationLocker: migrationLocker,

		removePodQueue:    make(chan removePod, 256),
		standbyCheckQueue: newStandbyCheckQueue(),
		standbyChecks:     make(map[standbyCheck]struct{}),
		elected:           mgr.Elected(),
		cfg:               cfg,
	}
	metrics.StandbyQuarantinedPodsGauge.WithLabelValues(tidbPoolNamespace).Set(0)
	return tManager, tManager.init()
}

func (t *tidbManager) IsLeader() bool {
	select {
	case <-t.elected:
		return true
	default:
		return false
	}
}

func (t *tidbManager) NeedLeaderElection() bool {
	return false
}

func (t *tidbManager) Start(ctx context.Context) error {
	g, gctx := errgroup.WithContext(ctx)
	t.ctx = gctx

	//nolint: gomnd // nolint
	t.processRemovePod(8)
	t.processStandbyCheck(1)

	g.Go(func() error { return t.startToFreeCacheAgainWatch() })
	g.Go(func() error { return t.startTokenWatch() })

	return g.Wait()
}

type podInfo struct {
	podIP    string
	nodeName string
}

func (t *tidbManager) fetchAllToken() (map[string]tidbInfo, int64, error) {
	resp, err := t.etcdKV.Get(context.TODO(), t.etcdTokenPrefix, clientv3.WithPrefix())
	if err != nil {
		return nil, 0, err
	}
	tokenMap := make(map[string]tidbInfo, resp.Count/2)
	for _, kv := range resp.Kvs {
		token, keyType := utils.ParseTokenFromKey(t.etcdTokenPrefix, string(kv.Key))
		if token == "" {
			t.Info("empty token", "key", kv.Key)
			continue
		}
		switch keyType {
		case utils.KeyEndpointsType:
			tInfo, ok := tokenMap[token]
			if ok {
				tInfo.setEndpoints(string(kv.Value), kv.ModRevision)
			} else {
				tInfo = newTidbInfo(string(kv.Value), 0, kv.ModRevision)
			}
			tokenMap[token] = tInfo
		case utils.KeyReplicasType:
			if tInfo, ok := tokenMap[token]; ok {
				replicas, err := strconv.Atoi(string(kv.Value))
				if err != nil {
					t.Error(err, "Atoi Replicas", "token", token, "replicas", kv.Value)
					continue
				}
				tInfo.replicas = replicas
				tInfo.replicasModRevision = kv.ModRevision
				tokenMap[token] = tInfo
			} else {
				replicas, err := strconv.Atoi(string(kv.Value))
				if err != nil {
					t.Error(err, "Atoi Replicas", "token", token, "replicas", kv.Value)
					continue
				}
				tInfo := newTidbInfo("", replicas, 0)
				tInfo.replicasModRevision = kv.ModRevision
				tokenMap[token] = tInfo
			}
		default:
			t.Info("unknown keyType", "keyType", keyType, "key", kv.Key)
		}
	}
	return tokenMap, resp.Header.GetRevision(), nil
}

func (t *tidbManager) init() error {
	tokenMap, revision, err := t.fetchAllToken()
	if err != nil {
		return err
	}

	t.freeCache.init(nil)
	t.assignCache.init(tokenMap)
	t.tokenRevision = revision
	t.toFreeCacheAgainRevision = revision

	return nil
}

func (t *tidbManager) syncCacheSizeMetric() {
	tokenSize, endpointSize := t.assignCache.getSize()
	metrics.CacheSizeGauge.WithLabelValues("assign_token", t.tidbPoolNamespace).Set(float64(tokenSize))
	metrics.CacheSizeGauge.WithLabelValues("assign_endpoint", t.tidbPoolNamespace).Set(float64(endpointSize))
	metrics.CacheSizeGauge.WithLabelValues("free", t.tidbPoolNamespace).Set(float64(t.FreeCacheSize()))
}

// SyncCache sync free/assign cache.
func (t *tidbManager) SyncCache(ctx context.Context, nodeName, podName, podIP string, podStatus utils.PodStatus) error {
	defer t.syncCacheSizeMetric()

	endpoint := ep.GenEndpointFromPod(podName, podIP)

	var err error
	t.standbyCheckMu.Lock()
	switch podStatus {
	case utils.PodReady:
		item := standbyCheck{podName: podName, podIP: podIP}
		if _, quarantined := t.standbyChecks[item]; quarantined {
			t.freeCache.setPodState(nodeName, podName, podIP, freePodQuarantined)
		} else if t.assignCache.findTokenByEndpoint(endpoint) == "" {
			t.freeCache.setPod(nodeName, podName, podIP)
		} else {
			t.freeCache.removePodIfIP(podName, podIP)
		}
	case utils.PodNotReady:
		t.freeCache.removePodIfIP(podName, podIP)
	case utils.PodTerminating:
		t.freeCache.removePodIfIP(podName, podIP)
	}
	t.standbyCheckMu.Unlock()

	if podStatus == utils.PodTerminating {
		err = t.removeEndpoint(ctx, endpoint, "pod terminating")
	}

	t.deleteCreatingPod(podName)

	return err
}

func (t *tidbManager) GetTidbPool() tidbpool.Pool {
	return t.tidbPool
}

func (t *tidbManager) LockFreeCache() {
	atomic.StoreInt64(&t.lockFreeCache, 1)
}

func (t *tidbManager) UnlockFreeCache() {
	atomic.StoreInt64(&t.lockFreeCache, 0)
}

func (t *tidbManager) FreeCacheSize() int {
	return t.freeCache.getSize()
}

func (t *tidbManager) AssignCacheSize() (token, endpoint int) {
	return t.assignCache.getSize()
}

func (t *tidbManager) removeEndpointByToken(ctx context.Context, token, endpoint, comment string) error {
	tInfo, ok := t.assignCache.getTidbInfo(token)
	if !ok {
		return nil
	}
	keyEndpoints := t.etcdTokenPrefix + token + "/" + utils.KeyEndpoints
	keyReplicas := t.etcdTokenPrefix + token + "/" + utils.KeyReplicas

	leaseID, err := etcdlease.GetLease30d(ctx, t.etcdLease)
	if err != nil {
		return err
	}

	for {
		pre := tInfo.getEndpoints()
		if !tInfo.removeEndpoint(endpoint) {
			return nil
		}
		tInfo.setComment(comment)
		now := tInfo.getEndpoints()

		var put []clientv3.Op
		if len(tInfo.endpoints) > 0 {
			// still remain lease
			put = append(put, clientv3.OpPut(keyEndpoints, now, clientv3.WithIgnoreLease()))
		} else {
			// no endpoint remain, reset new lease
			put = append(put, clientv3.OpPut(keyEndpoints, now, clientv3.WithLease(leaseID)))
			put = append(put, clientv3.OpPut(keyReplicas, "", clientv3.WithLease(leaseID), clientv3.WithIgnoreValue()))
		}

		resp, err := t.etcdKV.Txn(ctx).
			If(clientv3.Compare(clientv3.Value(keyEndpoints), "=", pre)).
			Then(put...).
			Else(clientv3.OpGet(keyEndpoints)).
			Commit()
		if err != nil {
			return err
		}
		if resp.Succeeded {
			return nil
		}

		kvs := resp.Responses[0].GetResponseRange().GetKvs()
		if len(kvs) == 0 {
			// token may deleted
			return nil
		}
		kv := kvs[0]
		tInfo.setEndpoints(string(kv.Value), kv.ModRevision)
	}
}

func (t *tidbManager) removeEndpoint(ctx context.Context, endpoint, comment string) error {
	token := t.assignCache.findTokenByEndpoint(endpoint)
	if token == "" {
		return t.trySetToFreeCacheAgain(ctx, endpoint, comment)
	}
	tInfo, ok := t.assignCache.getTidbInfo(token)
	if !ok {
		return t.trySetToFreeCacheAgain(ctx, endpoint, comment)
	}
	keyEndpoints := t.etcdTokenPrefix + token + "/" + utils.KeyEndpoints
	keyReplicas := t.etcdTokenPrefix + token + "/" + utils.KeyReplicas

	leaseID, err := etcdlease.GetLease30d(ctx, t.etcdLease)
	if err != nil {
		return err
	}

	for {
		pre := tInfo.getEndpoints()
		if !tInfo.removeEndpoint(endpoint) {
			return t.trySetToFreeCacheAgain(ctx, endpoint, comment)
		}
		tInfo.setComment(comment)
		now := tInfo.getEndpoints()

		var put []clientv3.Op
		if len(tInfo.endpoints) > 0 {
			// still remain lease
			put = append(put, clientv3.OpPut(keyEndpoints, now, clientv3.WithIgnoreLease()))
		} else {
			// no endpoint remain, reset new lease
			put = append(put, clientv3.OpPut(keyEndpoints, now, clientv3.WithLease(leaseID)))
			put = append(put, clientv3.OpPut(keyReplicas, "", clientv3.WithLease(leaseID), clientv3.WithIgnoreValue()))
		}

		resp, err := t.etcdKV.Txn(ctx).
			If(clientv3.Compare(clientv3.Value(keyEndpoints), "=", pre)).
			Then(put...).
			Else(clientv3.OpGet(keyEndpoints)).
			Commit()
		if err != nil {
			return err
		}
		if resp.Succeeded {
			return nil
		}

		kvs := resp.Responses[0].GetResponseRange().GetKvs()
		if len(kvs) == 0 {
			// token may deleted
			return t.trySetToFreeCacheAgain(ctx, endpoint, comment)
		}
		kv := kvs[0]
		tInfo.setEndpoints(string(kv.Value), kv.ModRevision)
	}
}

func (t *tidbManager) trySetToFreeCacheAgain(ctx context.Context, endpoint, comment string) error {
	if comment != "" {
		return nil
	}
	podName, _ := ep.ParsePodFromEndpoint(endpoint)
	if _, err := t.etcdKV.Put(ctx, t.etcdToFreeCacheAgain, podName); err != nil {
		return perrors.Annotatef(err, "failed to trySetToFreeCacheAgain %s", podName)
	}

	return nil
}

func (t *tidbManager) checkAndSetToFreeCache(ctx context.Context, podName string) {
	if atomic.LoadInt64(&t.lockFreeCache) > 0 {
		if t.IsLeader() {
			t.removePod(removePod{name: podName, force: true})
		}
		return
	}

	pod := &corev1.Pod{}
	if err := t.client.Get(ctx, client.ObjectKey{Name: podName, Namespace: t.tidbPoolNamespace}, pod); err != nil {
		t.Error(err, "checkAndSetToFreeCache get pod err", "pod name", podName)
		return
	}

	t.setPodToFreeCacheIfEligible(pod, pod.Status.PodIP)
}

func (t *tidbManager) setPodToFreeCacheIfEligible(pod *corev1.Pod, expectedIP string) bool {
	item := standbyCheck{podName: pod.Name, podIP: expectedIP}
	t.standbyCheckMu.Lock()
	defer t.standbyCheckMu.Unlock()

	if _, quarantined := t.standbyChecks[item]; quarantined {
		t.freeCache.setPodState(pod.Spec.NodeName, pod.Name, expectedIP, freePodQuarantined)
		return false
	}
	return t.setPodToFreeCacheIfEligibleLocked(pod, expectedIP)
}

// setPodToFreeCacheIfEligibleLocked requires standbyCheckMu to be held.
func (t *tidbManager) setPodToFreeCacheIfEligibleLocked(pod *corev1.Pod, expectedIP string) bool {
	if !t.podEligibleForFreeCacheLocked(pod, expectedIP) {
		return false
	}
	t.freeCache.setPod(pod.Spec.NodeName, pod.Name, expectedIP)
	return true
}

// restorePodToFreeCacheIfEligibleLocked requires standbyCheckMu to be held.
func (t *tidbManager) restorePodToFreeCacheIfEligibleLocked(pod *corev1.Pod, expectedIP string) bool {
	if !t.podEligibleForFreeCacheLocked(pod, expectedIP) {
		t.freeCache.removePodIfIP(pod.Name, expectedIP)
		return false
	}
	return t.freeCache.setPodState(pod.Spec.NodeName, pod.Name, expectedIP, freePodAvailable)
}

// podEligibleForFreeCacheLocked requires standbyCheckMu to be held.
func (t *tidbManager) podEligibleForFreeCacheLocked(pod *corev1.Pod, expectedIP string) bool {
	if expectedIP == "" || pod.Status.PodIP != expectedIP {
		t.Info("pod ip changed, skip set to free cache",
			"pod name", pod.Name, "expected pod ip", expectedIP, "pod ip", pod.Status.PodIP)
		return false
	}
	if pod.DeletionTimestamp != nil ||
		pod.Status.Phase == corev1.PodSucceeded ||
		pod.Status.Phase == corev1.PodFailed {
		t.Info("pod is terminating or terminal, skip set to free cache",
			"pod name", pod.Name, "pod ip", pod.Status.PodIP, "phase", pod.Status.Phase)
		return false
	}
	if !kubeapi.IsPodReady(pod) {
		t.Info("pod not ready, skip set to free cache", "pod name", pod.Name, "pod ip", pod.Status.PodIP)
		return false
	}
	if atomic.LoadInt64(&t.lockFreeCache) > 0 {
		if t.IsLeader() {
			t.removePod(removePod{name: pod.Name, force: true})
		}
		return false
	}
	if pod.Annotations[types.AnnoTidbRemoveAfterRun] == "true" {
		if t.IsLeader() {
			t.Info("pod AnnoTidbRemoveAfterRun is true", "pod name", pod.Name, "pod ip", pod.Status.PodIP)
			t.removePod(removePod{name: pod.Name, force: true})
		}
		return false
	}

	endpoint := ep.GenEndpointFromPod(pod.Name, expectedIP)
	if t.assignCache.findTokenByEndpoint(endpoint) != "" {
		return false
	}

	return true
}

func (t *tidbManager) modifyTidbInfo(ctx context.Context, token, endpoint string, tInfo *tidbInfo) error {
	keyEndpoints := t.etcdTokenPrefix + token + "/" + utils.KeyEndpoints
	keyReplicas := t.etcdTokenPrefix + token + "/" + utils.KeyReplicas

	for !tInfo.findEndpoint(endpoint) {
		pre := tInfo.getEndpoints()
		tInfo.appendEndpoint(endpoint)
		now := tInfo.getEndpoints()
		resp, err := t.etcdKV.Txn(ctx).
			If(clientv3.Compare(clientv3.Value(keyEndpoints), "=", pre)).
			Then(clientv3.OpPut(keyEndpoints, now), clientv3.OpPut(keyReplicas, "", clientv3.WithIgnoreValue())).
			Else(clientv3.OpGet(keyEndpoints)).
			Commit()
		if err != nil {
			return err
		}
		if resp.Succeeded {
			t.V(3).Info("put txn succeeded", "token", token, "endpoints", now)
			break
		}

		kvs := resp.Responses[0].GetResponseRange().GetKvs()
		if len(kvs) == 0 {
			return perrors.Annotatef(
				utils.ErrNotFoundToken, "failed to modifyTidbInfo for token %s, endpoint %s", token, endpoint)
		}
		kv := kvs[0]
		tInfo.setEndpoints(string(kv.Value), kv.ModRevision)
	}

	return nil
}

func (t *tidbManager) GetTidbFromFreeCache(node string, includeActivating bool) []string {
	tidbList := make([]string, 0)
	t.freeCache.walk(func(nodeName, podName string, info freePodInfo) bool {
		if info.isVisible(includeActivating) {
			if node == "" || nodeName == node {
				tidbList = append(tidbList, podName)
			}
		}
		return true
	})

	return tidbList
}

func (t *tidbManager) GetOldTidbFromFreeCache(newRevision string, includeActivating bool) []string {
	nodeScoreObj := &serverlessv1alpha1.NodeScore{}
	if err := t.client.Get(t.ctx, client.ObjectKey{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
		nodeScoreObj); err != nil {
		t.Error(err, "failed to get node score, skip sort")
	}

	type nodeTidb struct {
		nodeScore int
		tidbList  []string
	}
	oldNodeTidbList := make([]*nodeTidb, 0, len(nodeScoreObj.Status.SchedulerNodeScore))
	t.freeCache.walkNode(func(nodeName string, podMap map[string]freePodInfo) bool {
		nTidb := &nodeTidb{tidbList: make([]string, 0, len(podMap))}
		for podName, info := range podMap {
			if info.isVisible(includeActivating) {
				if !strings.HasPrefix(podName, types.TidbServerNamePrefix+newRevision) {
					nTidb.tidbList = append(nTidb.tidbList, podName)
				}
			}
		}
		if len(nTidb.tidbList) > 0 {
			if len(nodeScoreObj.Status.SchedulerNodeScore) > 0 {
				nTidb.nodeScore = nodeScoreObj.Status.SchedulerNodeScore[nodeName]
			}
			oldNodeTidbList = append(oldNodeTidbList, nTidb)
		}
		return true
	})

	sort.Slice(oldNodeTidbList, func(i, j int) bool {
		return oldNodeTidbList[i].nodeScore < oldNodeTidbList[j].nodeScore
	})

	tidbList := make([]string, 0)
	for _, nTidb := range oldNodeTidbList {
		tidbList = append(tidbList, nTidb.tidbList...)
	}
	return tidbList
}

func (t *tidbManager) GetTidbFromFreeCacheByRevision(revision string, includeActivating bool) []string {
	newTidbList := make([]string, 0)
	t.freeCache.walk(func(_, podName string, info freePodInfo) bool {
		if info.isVisible(includeActivating) {
			if strings.HasPrefix(podName, types.TidbServerNamePrefix+revision) {
				newTidbList = append(newTidbList, podName)
			}
		}
		return true
	})

	return newTidbList
}

func (t *tidbManager) GetOldTidbFromAssignCache(newRevision string, limit int) (map[string][]string, int) {
	var oldTidb map[string][]string
	var oldTidbLen int
	total := 0
	if limit < 0 {
		oldTidb = make(map[string][]string)
	} else {
		oldTidb = make(map[string][]string, limit)
	}

	t.assignCache.walkRevisionEndpointMap(func(revision string, endpoints map[string]string) bool {
		if revision == newRevision {
			return true
		}

		total += len(endpoints)

		if limit < 0 || oldTidbLen < limit {
			for endpoint, token := range endpoints {
				if name, _ := ep.ParsePodFromEndpoint(endpoint); name != "" {
					oldTidb[token] = append(oldTidb[token], name)
					oldTidbLen++
				}
				if limit >= 0 && oldTidbLen >= limit {
					break
				}
			}
		}
		return true
	})

	return oldTidb, total
}

func (t *tidbManager) WalkActiveTokenFromAssignCache(f func(token string, endpoints []string, replicas int) bool) {
	for token, tInfo := range t.assignCache.getActiveTokenInfo() {
		if !f(token, tInfo.endpoints, tInfo.replicas) {
			return
		}
	}
}

func (t *tidbManager) FilterOrphanPods(inactivePods []string) []string {
	var res []string
	for _, podName := range inactivePods {
		if t.freeCache.getPod(podName) != "" {
			continue
		}
		res = append(res, podName)
	}

	return res
}

func (t *tidbManager) SetTidbScaler(s utils.Scaler) {
	t.tidbScaler = s
}

func (t *tidbManager) GetTidbScaler() utils.Scaler {
	return t.tidbScaler
}

func (t *tidbManager) KillAll(ctx context.Context, token string, timeout time.Duration) error {
	endpoints := t.GetAssignedTidbByToken(token)
	if endpoints == nil || len(endpoints.Endpoints) == 0 {
		return nil
	}

	for _, endpoint := range endpoints.Endpoints {
		_, host := ep.ParsePodFromEndpoint(endpoint)
		err := t.tidbController.Kill(ctx, net.JoinHostPort(host, "10080"), utils.TrimToken(token))
		if err != nil {
			return err
		}
	}

	timer := time.NewTimer(timeout)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-timer.C:
		if endpoints = t.GetAssignedTidbByToken(token); endpoints == nil || len(endpoints.Endpoints) == 0 {
			return nil
		}
		return perrors.New("waiting all tidb to be killed")
	}
}

func (t *tidbManager) GetAssignedTidbByToken(token string) *ep.Endpoints {
	ti, ok := t.assignCache.getTidbInfo(token)
	if !ok {
		return nil
	}
	return &ep.Endpoints{
		Endpoints:   ti.endpoints,
		TidbPool:    t.tidbPoolNamespace,
		Comment:     ti.comment,
		ModRevision: ti.modRevision(),
		Replicas:    ti.replicas,
	}
}

func (t *tidbManager) GetAssignedToken(podName, podIP string) string {
	return t.assignCache.findTokenByEndpoint(ep.GenEndpointFromPod(podName, podIP))
}