package cluster

import (
	"context"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/pingcap/errors"
	"github.com/pingcap/kvproto/pkg/keyspacepb"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/entity"
	infrautil "github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"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/pd"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	clientv3 "go.etcd.io/etcd/client/v3"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	defaultKillAllTimeout     = 2 * time.Second
	auditLogInplaceUpdateWait = 3 * time.Second
	retryTidbPoolChanging     = 40 * time.Second
	waitNewTidbPoolRunning    = 5 * time.Second
)

var (
	keyspaceEnabled  = keyspacepb.KeyspaceState_ENABLED.String()
	keyspaceDisabled = keyspacepb.KeyspaceState_DISABLED.String()
)

func (r *reconcileContext) updateKeyspaceConfig() error {
	minRcu := strconv.FormatInt(int64(r.cluster.Spec.TiDB.MinRcu), 10)
	maxRcu := strconv.FormatInt(int64(r.cluster.Spec.TiDB.MaxRcu), 10)

	cfg := &pd.KeyspaceConfig{
		TenantID:  r.cluster.TagsTenant,
		ProjectID: r.cluster.TagsProject,
		ClusterID: r.cluster.Name,
		MinRCU:    minRcu,
		MaxRCU:    maxRcu,
	}

	if r.cluster.Spec.RuVersion > 0 {
		cfg.RuVersion = strconv.FormatInt(int64(r.cluster.Spec.RuVersion), 10)
	}

	err := r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, cfg)
	if err != nil {
		return errors.Annotate(err, "failed to update keyspace config")
	}

	return nil
}

// setBootstrappedForRestore can only be used in create reconcile flow, do not use it after cluster is running
func (r *reconcileContext) setBootstrappedForRestore() error {
	if r.cluster.Annotations[types.CreateFromRestore] != "true" {
		return nil
	}
	cfg := &pd.KeyspaceConfig{
		IsBootstrappedForRestore: "False",
	}
	err := r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, cfg)
	if err != nil {
		return errors.Annotate(err, "failed to set bootstrapped for restore")
	}
	r.Info("set bootstrapped for restore to False succeed")
	return nil
}

func (r *reconcileContext) removeOldEtcdToken() error {
	keepPoolRG := map[string][]string{r.specTidbPool: {""}}
	for rgName, rgSpec := range r.specResourceGroups {
		pool := rgSpec.Pool
		if pool == "" {
			pool = r.specTidbPool
		}
		keepPoolRG[pool] = append(keepPoolRG[pool], rgName)
	}
	for pool, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
		var keeps []struct{ s, e string }
		prefix := etcdTokenPrefix + r.keyspaceToken + "/"
		for _, rg := range keepPoolRG[pool] {
			if rg == "" {
				keeps = append(keeps, struct{ s, e string }{s: prefix + utils.KeyReplicas, e: prefix + utils.KeyReplicas + "\x00"})
				keeps = append(keeps, struct{ s, e string }{s: prefix + utils.KeyEndpoints, e: prefix + utils.KeyEndpoints + "\x00"})
			} else {
				keeps = append(keeps, struct{ s, e string }{s: prefix + rg + "/", e: clientv3.GetPrefixRangeEnd(prefix + rg + "/")})
			}
		}
		sort.Slice(keeps, func(i, j int) bool { return keeps[i].s < keeps[j].s })

		start := prefix
		ops := make([]clientv3.Op, 0, len(keeps)+1)
		for _, k := range keeps {
			ops = append(ops, clientv3.OpDelete(start, clientv3.WithRange(k.s)))
			start = k.e
		}
		ops = append(ops, clientv3.OpDelete(start, clientv3.WithRange(clientv3.GetPrefixRangeEnd(prefix))))
		if _, err := r.EtcdKV.Txn(r.ctx).Then(ops...).Commit(); err != nil {
			return errors.Annotatef(err, "failed to delete etcd prefix %s", prefix)
		}
	}
	return nil
}

func (r *reconcileContext) removeEtcdToken(tidbPool, token string) error {
	etcdTokenPrefix := r.EtcdTokenPrefixMap[tidbPool]
	if etcdTokenPrefix == "" {
		return errors.Errorf("etcdTokenPrefix not found in tidbPool %s", tidbPool)
	}

	r.Info("delete etcd token", "etcdTokenPrefix", etcdTokenPrefix, "tidbPool", tidbPool, "token", token)
	if _, err := r.EtcdKV.Delete(r.ctx, etcdTokenPrefix+token+"/", clientv3.WithPrefix()); err != nil {
		return errors.Annotatef(err, "failed to delete etcd prefix %s", etcdTokenPrefix+r.keyspaceToken+"/")
	}

	return nil
}

func (r *reconcileContext) setEtcdToken() error {
	leaseID, err := etcdlease.GetLease30d(r.ctx, r.EtcdLease)
	if err != nil {
		return errors.Annotate(err, "failed to get etcd lease 30d")
	}

	if err := r.resetToken(r.EtcdTokenPrefixMap[r.specTidbPool], r.keyspaceToken, r.cluster.Spec.TiDB.Replicas, leaseID); err != nil {
		return err
	}
	for rgName, rgSpec := range r.specResourceGroups {
		pool := rgSpec.Pool
		if pool == "" {
			pool = r.specTidbPool
		}
		etcdTokenPrefix := r.EtcdTokenPrefixMap[pool]
		if etcdTokenPrefix == "" {
			return errors.Errorf("etcdTokenPrefix not found in pool %s", pool)
		}
		if err := r.resetToken(etcdTokenPrefix, r.keyspaceToken+"/"+rgName, rgSpec.Replicas, leaseID); err != nil {
			return err
		}
	}

	return nil
}

func (r *reconcileContext) setEtcdReplicas(tidbPool, token string, replicas int) error {
	etcdTokenPrefix := r.EtcdTokenPrefixMap[tidbPool]
	if etcdTokenPrefix == "" {
		return errors.Errorf("etcdTokenPrefix not found in tidbPool %s", tidbPool)
	}

	r.Info("put etcd replicas", "etcdTokenPrefix", etcdTokenPrefix, "tidbPool", tidbPool, "token", token, "replicas", replicas)
	keyEndpoints := etcdTokenPrefix + token + "/" + utils.KeyEndpoints
	resp, err := r.EtcdKV.Get(r.ctx, keyEndpoints)
	if err != nil {
		return errors.Annotatef(err, "failed to get etcd path %s", keyEndpoints)
	}
	if len(resp.Kvs) == 0 {
		return errors.Errorf("tidbPool %s not exist", tidbPool)
	}

	keyReplicas := etcdTokenPrefix + token + "/" + utils.KeyReplicas
	if _, err := r.EtcdKV.Put(r.ctx, keyReplicas, strconv.Itoa(replicas), clientv3.WithIgnoreLease()); err != nil {
		return errors.Annotatef(err, "failed to put etcd path %s", keyReplicas)
	}

	return nil
}

func (r *reconcileContext) reduceOneReplicaFromOldPools(pool, sg string) error {
	if pool == "" {
		return errors.New("pool is empty")
	}

	token := r.keyspaceToken
	if sg != "" {
		token = r.keyspaceToken + "/" + sg
	}
	for tidbPool, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
		if pool == tidbPool {
			continue
		}
		keyReplicas := etcdTokenPrefix + token + "/" + utils.KeyReplicas

		// Try to get replicas under the old pool.
		resp, err := r.EtcdKV.Get(r.ctx, keyReplicas)
		if err != nil {
			return errors.Annotatef(err, "failed to get etcd path %s", keyReplicas)
		}
		if len(resp.Kvs) == 0 {
			continue
		}

		v := string(resp.Kvs[0].Value)
		n, convErr := strconv.Atoi(v)
		if convErr != nil {
			return errors.Annotatef(convErr, "invalid replicas value at %s: %q", keyReplicas, v)
		}
		if err := r.setTidbPoolReplicas(tidbPool, sg, n-1); err != nil {
			return errors.Annotatef(err, "failed to set tidbPool replicas: %s, %d", tidbPool, n-1)
		}
		break // only reduce one replica
	}

	return nil
}

func (r *reconcileContext) resetEtcdToken() error {
	leaseID, err := etcdlease.GetLease30d(r.ctx, r.EtcdLease)
	if err != nil {
		return errors.Annotate(err, "failed to get etcd lease 30d")
	}

	if err := r.resetToken(r.EtcdTokenPrefixMap[r.specTidbPool], r.keyspaceToken, r.cluster.Spec.TiDB.Replicas, leaseID); err != nil {
		return err
	}
	for rgName, rgSpec := range r.specResourceGroups {
		pool := rgSpec.Pool
		if pool == "" {
			pool = r.specTidbPool
		}
		etcdTokenPrefix := r.EtcdTokenPrefixMap[pool]
		if etcdTokenPrefix == "" {
			return errors.Errorf("etcdTokenPrefix not found in pool %s", pool)
		}
		if err := r.resetToken(etcdTokenPrefix, r.keyspaceToken+"/"+rgName, rgSpec.Replicas, leaseID); err != nil {
			return err
		}
	}

	return r.removeOldEtcdToken()
}

func (r *reconcileContext) resetToken(etcdTokenPrefix, token string, replicas int, leaseID clientv3.LeaseID) error {
	keyEndpoints := etcdTokenPrefix + token + "/" + utils.KeyEndpoints
	r.Info("put etcd", "etcdTokenPrefix", etcdTokenPrefix, "token", token)
	isPut := false
	resp, err := r.EtcdKV.Get(r.ctx, keyEndpoints)
	if err != nil {
		return errors.Annotatef(err, "failed to get etcd path %s", keyEndpoints)
	}
	if len(resp.Kvs) == 0 {
		if _, err := r.EtcdKV.Put(r.ctx, keyEndpoints, "", clientv3.WithLease(leaseID)); err != nil {
			return errors.Annotatef(err, "failed to put etcd path %s", keyEndpoints)
		}
		isPut = true
	}

	keyReplicas := etcdTokenPrefix + token + "/" + utils.KeyReplicas
	op := clientv3.WithIgnoreLease()
	if isPut {
		op = clientv3.WithLease(leaseID)
	}
	if _, err := r.EtcdKV.Put(r.ctx, keyReplicas, strconv.Itoa(replicas), op); err != nil {
		return errors.Annotatef(err, "failed to put etcd path %s", keyReplicas)
	}

	return nil
}

func (r *reconcileContext) createServerlessTidbCluster() error {
	if err := r.assignKeyspace(); err != nil {
		return errors.Annotate(err, "failed to assign keyspace")
	}
	if err := r.updateKeyspaceConfig(); err != nil {
		return errors.Annotate(err, "failed to update keyspace")
	}

	if err := r.setBootstrappedForRestore(); err != nil {
		return errors.Annotate(err, "failed to set bootstrapped for restore")
	}

	if err := r.resetEtcdToken(); err != nil {
		return errors.Annotate(err, "failed to create etcd token")
	}

	return errors.Annotate(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.TiDB = infra.TiDBStatus{
			Version:         r.tidbVersion,
			UpdatedReplicas: r.cluster.Spec.TiDB.Replicas,
		}
		status.TidbPool = r.specTidbPool

		statusResourceGroups := make(map[string]infra.ResourceGroup)
		for rgName, rgSpec := range r.specResourceGroups {
			rgSpecCopy := rgSpec.DeepCopy()
			if rgSpecCopy.Pool == "" {
				rgSpecCopy.Pool = r.specTidbPool
			}
			statusResourceGroups[rgName] = *rgSpecCopy
		}
		status.ResourceGroups = statusResourceGroups
		return nil
	}), "failed to ensureStatus in createServerlessTidbCluster")
}

func (r *reconcileContext) deleteServerlessTidbCluster() error {
	r.Info("delete serverless tidb cluster",
		"keyspaceToken", r.keyspaceToken,
		"skipClusterDeleteEtcdAndPd", r.SkipClusterDeleteEtcdAndPd)
	if r.SkipClusterDeleteEtcdAndPd {
		r.Info("skip delete serverless tidb cluster", "keyspaceToken", r.keyspaceToken)
		return nil
	}
	if err := r.deleteEtcdToken(); err != nil {
		return errors.Annotate(err, "error delete etcd token")
	}
	if err := r.deleteKeyspace(); err != nil {
		return errors.Annotate(err, "error delete Keyspace")
	}
	if err := r.deleteResourceGroup(); err != nil {
		return errors.Annotate(err, "error delete resource group")
	}
	r.Info("delete serverless tidb cluster successful", "keyspaceToken", r.keyspaceToken)
	return nil
}

func (r *reconcileContext) deleteKeyspace() error {
	if r.keyspaceToken == "" {
		return nil
	}

	keyspaceList := &serverlessv1alpha1.KeyspaceList{}
	err := r.Client.List(r.ctx, keyspaceList,
		client.InNamespace(r.keyspaceNs), client.MatchingFields{"spec.keyspaceName": r.keyspaceToken})
	if err != nil {
		return errors.Annotate(err, "error list Keyspace")
	}

	if len(keyspaceList.Items) == 0 {
		// the cluster may be created before 202212, so we try to update pd keyspace directly
		keyspace, err := r.PdClient.GetKeyspaceByName(r.ctx, r.keyspaceToken)
		if err != nil {
			if pd.IsKeyspaceNotExists(err) {
				return nil
			}
			return err
		}
		if keyspace != nil {
			switch keyspace.State {
			case keyspaceEnabled:
				if err := r.PdClient.UpdateKeyspaceState(r.ctx, r.keyspaceToken, keyspacepb.KeyspaceState_DISABLED); err != nil {
					return err
				}
				return errors.New("should requeue after disabling keyspace")
			case keyspaceDisabled:
				if err := r.PdClient.UpdateKeyspaceState(r.ctx, r.keyspaceToken, keyspacepb.KeyspaceState_ARCHIVED); err != nil {
					return err
				}
			}
		}
		return nil
	}

	return client.IgnoreNotFound(r.Client.Delete(r.ctx, &keyspaceList.Items[0]))
}

func (r *reconcileContext) deleteEtcdToken() error {
	if r.keyspaceToken == "" {
		return nil
	}

	for _, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
		_, err := r.EtcdKV.Delete(r.ctx, etcdTokenPrefix+r.keyspaceToken+"/", clientv3.WithPrefix())
		if err != nil {
			return errors.Annotatef(err, "failed to delete etcd prefix %s",
				etcdTokenPrefix+r.keyspaceToken+"/")
		}
	}
	return nil
}

func (r *reconcileContext) reconcileResumeEtcdToken() error {
	if r.cluster.Annotations[types.AnnoKeyResumeEtcdToken] == "" {
		return nil
	}

	r.Info("resume etcd token", "specTidbPool", r.specTidbPool)

	if err := r.setEtcdToken(); err != nil {
		return errors.Annotate(err, "failed to resume etcd token")
	}

	delete(r.cluster.Annotations, types.AnnoKeyResumeEtcdToken)
	return r.InfraClusterCli.Update(r.ctx, r.cluster)
}

func (r *reconcileContext) allOldTidbPoolPods(tidbPool, sg string) []client.ObjectKey {
	token := r.keyspaceToken
	if sg != "" {
		token = r.keyspaceToken + "/" + sg
	}
	oldPods := make([]client.ObjectKey, 0)
	for pool, tManager := range r.TidbManagers {
		if pool != tidbPool {
			if endpoints := tManager.GetAssignedTidbByToken(token); endpoints != nil {
				for _, endpoint := range endpoints.Endpoints {
					podName, _ := ep.ParsePodFromEndpoint(endpoint)
					oldPods = append(oldPods, client.ObjectKey{Namespace: pool, Name: podName})
				}
			}
		}
	}

	return oldPods
}

func (r *reconcileContext) newTidbPoolPods(tidbPool, sg string) []client.ObjectKey {
	token := r.keyspaceToken
	if sg != "" {
		token = r.keyspaceToken + "/" + sg
	}
	newPods := make([]client.ObjectKey, 0)
	if tManager := r.TidbManagers[tidbPool]; tManager != nil {
		if endpoints := tManager.GetAssignedTidbByToken(token); endpoints != nil {
			for _, endpoint := range endpoints.Endpoints {
				podName, _ := ep.ParsePodFromEndpoint(endpoint)
				newPods = append(newPods, client.ObjectKey{Namespace: tidbPool, Name: podName})
			}
		}
	}

	return newPods
}

func (r *reconcileContext) needSetEtcdToken() bool {
	if r.specTidbPool != r.cluster.Status.TidbPool {
		r.Info("tidb pool have changed", "old", r.cluster.Status.TidbPool, "now", r.specTidbPool)
		return true
	}
	for rgName, rgSpec := range r.specResourceGroups {
		pool := rgSpec.Pool
		if pool == "" {
			pool = r.specTidbPool
		}
		rgStatus := r.cluster.Status.ResourceGroups[rgName]
		if pool != rgStatus.Pool {
			r.Info("tidb resource group pool have changed", "rgName", rgName, "old", rgStatus.Pool, "now", pool)
			return true
		}
	}

	return false
}

func (r *reconcileContext) reconcileTidbPool() error {
	if r.needSetEtcdToken() {
		r.cluster.Annotations = utils.UpdateMap(r.cluster.Annotations, types.AnnoKeyTidbPoolChanging, time.Now().Format(time.RFC3339))
		if err := r.InfraClusterCli.Update(r.ctx, r.cluster); err != nil {
			return errors.Annotate(err, "failed to set infra cluster AnnoKeyTidbPoolChanging")
		}

		// for _, objKey := range r.allOldTidbPoolPods() {
		// 	data := fmt.Sprintf(`{"metadata":{"annotations":{"%s":"%s"}}}`, types.AnnoTidbStatusRunSlow, "true")
		// 	if err := r.Client.Patch(r.ctx, &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		// 		Name:      objKey.Name,
		// 		Namespace: objKey.Namespace,
		// 	}}, client.RawPatch(apitypes.MergePatchType, []byte(data))); err != nil {
		// 		return errors.Annotatef(err, "failed to patch tidb pod AnnoTidbStatusRunSlow %v", objKey)
		// 	}
		// }

		if err := r.setEtcdToken(); err != nil {
			return errors.Annotate(err, "failed to set tidb pool in etcd")
		}

		// for _, objKey := range r.newTidbPoolPods() {
		// 	newTidbPod := &corev1.Pod{}
		// 	if err := r.Client.Get(r.ctx, objKey, newTidbPod); err != nil {
		// 		return errors.Annotatef(err, "failed to get tidb pod %v", objKey)
		// 	}
		// 	if newTidbPod.Annotations[types.AnnoTidbStatusRunSlow] != "" {
		// 		data := fmt.Sprintf(`{"metadata":{"annotations":{"%s":null}}}`, types.AnnoTidbStatusRunSlow)
		// 		err := r.Client.Patch(r.ctx, newTidbPod, client.RawPatch(apitypes.StrategicMergePatchType, []byte(data)))
		// 		if err != nil {
		// 			return errors.Annotatef(err, "failed to remove tidb pod AnnoTidbStatusRunSlow %v", objKey)
		// 		}
		// 	}
		// }

		if err := r.ensureStatus(func(status *infra.ClusterStatus) error {
			status.TidbPool = r.specTidbPool
			for rgName, rgSpec := range r.specResourceGroups {
				rgSpecCopy := rgSpec.DeepCopy()
				if rgSpecCopy.Pool == "" {
					rgSpecCopy.Pool = r.specTidbPool
				}
				r.cluster.Status.ResourceGroups[rgName] = *rgSpecCopy
			}
			return nil
		}); err != nil {
			return errors.Annotate(err, "failed to ensureStatus in reconcileTidbPool")
		}
	}

	if tidbPoolChanging := r.cluster.Annotations[types.AnnoKeyTidbPoolChanging]; tidbPoolChanging != "" {
		if err := r.handleTidbPoolChanging(tidbPoolChanging, r.specTidbPool, "", r.cluster.Spec.TiDB.Replicas); err != nil {
			return err
		}
		for rgName, rgStatus := range r.cluster.Status.ResourceGroups {
			if err := r.handleTidbPoolChanging(tidbPoolChanging, rgStatus.Pool, rgName, rgStatus.Replicas); err != nil {
				return err
			}
		}

		// make sure all old tidb pool paths are removed
		if err := r.removeOldEtcdToken(); err != nil {
			return errors.Annotate(err, "failed to remove old tidb pools in etcd")
		}

		// remove AnnoKeyTidbPoolChanging when tidb pool changing is done
		delete(r.cluster.Annotations, types.AnnoKeyTidbPoolChanging)
		if err := r.InfraClusterCli.Update(r.ctx, r.cluster); err != nil {
			return errors.Annotate(err, "failed to remove infra cluster AnnoKeyTidbPoolChanging")
		}
		r.Info("tidb pool changing end")
	}

	return nil
}

func (r *reconcileContext) handleTidbPoolChanging(tidbPoolChanging, tidbPool, sg string, replicas int) error {
	if tidbPool == "" {
		return errors.New("tidbPool is empty")
	}

	allOldTidbPoolPodLen := len(r.allOldTidbPoolPods(tidbPool, sg))
	if allOldTidbPoolPodLen == 0 {
		return nil
	}

	tidbPoolChangingTimestamp, err := time.Parse(time.RFC3339, tidbPoolChanging)
	if err != nil {
		return errors.Annotatef(err, "failed to parse AnnoKeyTidbPoolChanging")
	}
	if d := time.Since(tidbPoolChangingTimestamp); d < retryTidbPoolChanging {
		return &utils.RequeueAfterError{RequeueAfter: retryTidbPoolChanging - d}
	}

	if allOldTidbPoolPodLen+len(r.newTidbPoolPods(tidbPool, sg)) > replicas {
		if err := r.reduceOneReplicaFromOldPools(tidbPool, sg); err != nil {
			return errors.Annotatef(err, "failed to reduce one replica from old pods")
		}
		r.Info("reduce one replica in old tidb pool")
		r.cluster.Annotations[types.AnnoKeyTidbPoolChanging] = time.Now().Format(time.RFC3339)
		if err := r.InfraClusterCli.Update(r.ctx, r.cluster); err != nil {
			// in order to change stability as much as possible, we just record error.
			r.Error(err, "failed to update infra cluster AnnoKeyTidbPoolChanging")
		}
		return &utils.RequeueAfterError{RequeueAfter: retryTidbPoolChanging}
	}
	return &utils.RequeueAfterError{
		Cause:        errors.New("waiting new tidb pool pods to running"),
		RequeueAfter: waitNewTidbPoolRunning,
	}
}

func (r *reconcileContext) reconcileTidbReplicas() error {
	if r.cluster.Spec.TiDB.Replicas != r.cluster.Status.TiDB.UpdatedReplicas {
		r.Info("tidb replicas have changed", "old", r.cluster.Status.TiDB.UpdatedReplicas, "now", r.cluster.Spec.TiDB.Replicas)
		if err := r.setTidbPoolReplicas(r.specTidbPool, "", r.cluster.Spec.TiDB.Replicas); err != nil {
			return errors.Annotatef(err, "failed to set tidbPool replicas: %s, %d", r.specTidbPool, r.cluster.Spec.TiDB.Replicas)
		}
	}
	for rgName, rgSpec := range r.specResourceGroups {
		statusRG, ok := r.cluster.Status.ResourceGroups[rgName]
		if !ok {
			r.Info("resource group not found in status, skip", "rgName", rgName)
			continue
		}
		if statusRG.Replicas != rgSpec.Replicas {
			r.Info("tidb rg replicas have changed", "resourceGroup", rgName, "old", statusRG.Replicas, "now", rgSpec.Replicas)
			if err := r.setTidbPoolReplicas(statusRG.Pool, rgName, rgSpec.Replicas); err != nil {
				return errors.Annotatef(err, "failed to set rg pool replicas: %s, %d", statusRG.Pool, rgSpec.Replicas)
			}
		}
	}

	return errors.Annotatef(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.TiDB.UpdatedReplicas = r.cluster.Spec.TiDB.Replicas
		for rgName, rgSpec := range r.specResourceGroups {
			if statusRG, ok := status.ResourceGroups[rgName]; ok {
				statusRG.Replicas = rgSpec.Replicas
				status.ResourceGroups[rgName] = statusRG
			}
		}
		return nil
	}), "failed to ensureStatus in reconcileTidbReplicas")
}

func (r *reconcileContext) ensureZoneRedundant() error {
	if r.cluster.Status.Phase != infra.ClusterPhaseRunning &&
		r.cluster.Status.Phase != infra.ClusterPhaseCreating {
		return nil
	}

	if r.cluster.Status.ZoneRedundantEnabled == r.cluster.Spec.EnableZoneRedundant {
		return nil
	}

	keyspaceZR := "false"
	if r.cluster.Spec.EnableZoneRedundant == "yes" {
		keyspaceZR = "true"
	}
	if err := r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, &pd.KeyspaceConfig{
		ZoneRedundant: keyspaceZR,
	}); err != nil {
		return errors.Annotatef(err, "failed to update keyspace ZoneRedundant to %v", keyspaceZR)
	}
	r.Info("enableZoneRedundant is updated", "enableZoneRedundant", r.cluster.Spec.EnableZoneRedundant)

	return errors.Annotate(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.ZoneRedundantEnabled = r.cluster.Spec.EnableZoneRedundant
		return nil
	}), "failed to ensureStatus in zoneRedundantEnabled")
}

func (r *reconcileContext) ensureRcuConfig() error {
	if r.cluster.Status.Phase != infra.ClusterPhaseRunning &&
		r.cluster.Status.Phase != infra.ClusterPhaseCreating {
		return nil
	}

	minRcu := strconv.FormatInt(int64(r.cluster.Spec.TiDB.MinRcu), 10)
	maxRcu := strconv.FormatInt(int64(r.cluster.Spec.TiDB.MaxRcu), 10)

	if r.cluster.Annotations[types.AnnoKeyPrevMinRCU] == minRcu &&
		r.cluster.Annotations[types.AnnoKeyPrevMaxRCU] == maxRcu {
		return nil
	}

	if err := r.updateKeyspaceConfig(); err != nil {
		return errors.Annotate(err, "failed to update keyspace config")
	}

	if err := r.updateResourceGroup(); err != nil {
		return errors.Annotate(err, "failed to update resource group")
	}

	if err := r.InfraClusterCli.Patch(r.ctx, map[string]any{
		fmt.Sprintf("annotations.'%s'", types.AnnoKeyPrevMaxRCU): maxRcu,
		fmt.Sprintf("annotations.'%s'", types.AnnoKeyPrevMinRCU): minRcu,
	}, r.cluster); err != nil {
		return errors.Annotate(err, "failed to update infra cluster rcu config")
	}

	r.Info("rcu config is updated", "minRcu", minRcu, "maxRcu", maxRcu)

	return nil
}

func (r *reconcileContext) ensureRuVersion() error {
	if r.cluster.Status.Phase != infra.ClusterPhaseRunning &&
		r.cluster.Status.Phase != infra.ClusterPhaseCreating {
		return nil
	}

	ruVersion := strconv.FormatInt(int64(r.cluster.Spec.RuVersion), 10)

	if r.cluster.Annotations[types.AnnoKeyPrevRuVersion] == ruVersion {
		return nil
	}

	if r.cluster.Spec.RuVersion > 0 {
		unsupportedRuVersion := false
		// Falls back gracefully if PD does not support it yet (404).
		if err := r.PdClient.SetKeyspaceRuVersion(r.ctx, r.keyspaceToken, &pd.SetKeyspaceRuVersionRequest{
			RuVersion: r.cluster.Spec.RuVersion,
		}); err != nil {
			if strings.Contains(err.Error(), "404 Not Found") {
				r.Error(err, "PD does not support ru-version API yet, skipping")
				unsupportedRuVersion = true
				if r.cluster.Spec.RuVersion == 2 {
					metrics.RuVersionUnsupportedTotal.WithLabelValues(r.cluster.Name, r.specTidbPool).Inc()
				}
			} else {
				return errors.Annotate(err, "failed to set keyspace ru-version")
			}
		}
		if !unsupportedRuVersion {
			if err := r.updateKeyspaceConfig(); err != nil {
				return errors.Annotate(err, "failed to update keyspace config")
			}
		}
	}

	if err := r.InfraClusterCli.Patch(r.ctx, map[string]any{
		fmt.Sprintf("annotations.'%s'", types.AnnoKeyPrevRuVersion): ruVersion,
	}, r.cluster); err != nil {
		return errors.Annotate(err, "failed to update infra cluster ru version annotation")
	}

	r.Info("ru config is updated", "ruVersion", ruVersion)

	return nil
}

func (r *reconcileContext) ensureResourceLimit() error {
	if r.cluster.Status.Phase != infra.ClusterPhaseRunning &&
		r.cluster.Status.Phase != infra.ClusterPhaseCreating {
		return nil
	}

	mode := GetSpeedLimitFromCluster(r.cluster).Type.String()
	if r.cluster.Status.RuLimitMode == mode {
		return nil
	}

	if mode == SpeedLimitEssentialName {
		// handle in ensureRcuConfig.
		return nil
	}

	r.Info("ru limit mode need update", "old", r.cluster.Status.RuLimitMode, "now", mode)
	if err := r.updateResourceGroup(); err != nil {
		return err
	}
	r.Info("update cluster event", "UpdateRuLimitMode", "ru limit mode is set to "+mode)

	return errors.Annotate(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.RuLimitMode = mode
		return nil
	}), "failed to ensureStatus in ensureResourceLimit")
}

func (r *reconcileContext) reconcileRestoreStatus() error {
	if r.cluster.Status.Phase != infra.ClusterPhaseRunning {
		return errors.New("only allow reconcile restore when cluster is running")
	}

	if r.cluster.Annotations[types.AnnoKeyKeyspaceRestore] == types.AnnoKeyspaceRestoreStart {
		for pool, tManager := range r.TidbManagers {
			if err := tManager.KillAll(r.ctx, r.keyspaceToken, defaultKillAllTimeout); err != nil {
				return errors.Annotatef(err, "failed to kill all tidbs for keyspaceToken in pool %s", pool)
			}
			for rgName := range r.cluster.Status.ResourceGroups {
				if err := tManager.KillAll(r.ctx, rgName, defaultKillAllTimeout); err != nil {
					return errors.Annotatef(err, "failed to kill all tidbs for keyspaceToken in pool %s, rgName %s", pool, rgName)
				}
			}
		}
		if err := r.TidbWorkerApi.Clean(r.ctx, r.keyspaceToken); err != nil {
			return errors.Annotate(err, "failed to clean tidb worker")
		}
		if err := r.InfraClusterCli.Patch(r.ctx, map[string]any{
			fmt.Sprintf("annotations.'%s'", types.AnnoKeyKeyspaceRestore): types.AnnoKeyspaceRestartTiDB,
		}, r.cluster); err != nil {
			return errors.Annotate(err, "failed to update infra cluster")
		}
		r.Info("reconcileRestoreStatus kill all tidb succeed")
	}

	return nil
}

func (r *reconcileContext) reconcileServerless() error {
	var reqErrs error

	if err := utils.IgnoreRequeue(&reqErrs, r.ensureResourceGroup()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.reconcileResumeEtcdToken()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.reconcileTidbPool()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.reconcileTidbReplicas()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.reconcileAuditLog()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.ensureResourceLimit()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.ensureZoneRedundant()); err != nil {
		return err
	}

	if r.O11yAgentEnabled {
		if err := utils.IgnoreRequeue(&reqErrs, r.ensureO11yCluster()); err != nil {
			return err
		}
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.reconcileRestoreStatus()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.ensureRcuConfig()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.ensureRuVersion()); err != nil {
		return err
	}

	if err := utils.IgnoreRequeue(&reqErrs, r.ensureSqlEndpoints()); err != nil {
		return err
	}

	return reqErrs
}

func (r *reconcileContext) observeTidbClusterStatus() error {
	return errors.Annotate(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.Phase = infra.ClusterPhaseRunning
		status.Reason = "cluster is running"
		status.ProxyEndpoint = fmt.Sprintf("%s/proxy/namespace/free-tier-%s/cluster/%s",
			r.BaseProxy, r.cluster.Namespace, "db"+r.cluster.Name)
		return nil
	}), "failed to ensureStatus in observeTidbClusterStatus")
}

func (r *reconcileContext) ensureSqlEndpoints() error {
	publicStatus := &infra.PublicStatus{}
	sqlEndpoints := r.cluster.Spec.SqlEndpoints
	if sqlEndpoints.Public != nil && sqlEndpoints.Public.Enabled {
		var err error
		if publicStatus, err = r.genPublicEndpointStatus(); err != nil {
			return errors.Annotate(err, "failed to genPublicEndpointStatus")
		}
	}
	awsPrivateLinkServiceStatus := &infra.AwsPrivatelinkServiceStatus{}
	if sqlEndpoints.AwsPrivatelinkService != nil && sqlEndpoints.AwsPrivatelinkService.Enabled {
		var err error
		if awsPrivateLinkServiceStatus, err = r.genPrivatelinkStatus(); err != nil {
			return errors.Annotate(err, "failed to genPrivatelinkStatus")
		}
	}

	return errors.Annotate(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.SqlEndpoints.Public = publicStatus
		status.SqlEndpoints.AwsPrivatelinkService = awsPrivateLinkServiceStatus
		return nil
	}), "failed to ensureStatus in ensureSqlEndpoints")
}

type SpeedLimitType int

const (
	SpeedLimitStandard SpeedLimitType = iota
	SpeedLimitThrottled
	SpeedLimitEssential
)

const (
	SpeedLimitStandardName  = "standard"
	SpeedLimitThrottledName = "throttled"
	SpeedLimitEssentialName = "essential"
)

func (t SpeedLimitType) String() string {
	switch t {
	case SpeedLimitStandard:
		return SpeedLimitStandardName
	case SpeedLimitThrottled:
		return SpeedLimitThrottledName
	case SpeedLimitEssential:
		return SpeedLimitEssentialName
	}
	return "unknown"
}

func (t SpeedLimitType) IsThrottled() bool {
	return t == SpeedLimitThrottled
}

type SpeedLimit struct {
	Type   SpeedLimitType
	MinRCU uint32
	MaxRCU uint32
}

var SpeedLimitSettings = []pd.SpeedLimitSetting{
	SpeedLimitStandard:  {FillRate: 2_000_000, BurstLimit: 50_000_000_000}, //nolint:mnd // ignore
	SpeedLimitThrottled: {FillRate: 10, BurstLimit: 50_000},                //nolint:mnd // ignore
}

func (l SpeedLimit) GetPDSetting() pd.SpeedLimitSetting {
	if l.Type == SpeedLimitEssential {
		return pd.SpeedLimitSetting{
			FillRate:   uint64(l.MaxRCU),
			BurstLimit: int64(l.MaxRCU),
		}
	}
	return SpeedLimitSettings[l.Type]
}

func GetSpeedLimitFromCluster(cluster *infra.Cluster) SpeedLimit {
	if infrautil.IsEssentialCluster(cluster) {
		return SpeedLimit{
			Type:   SpeedLimitEssential,
			MinRCU: cluster.Spec.TiDB.MinRcu,
			MaxRCU: cluster.Spec.TiDB.MaxRcu,
		}
	}

	switch cluster.Spec.RuLimitMode {
	case SpeedLimitStandardName:
		return SpeedLimit{Type: SpeedLimitStandard}
	case SpeedLimitThrottledName:
		return SpeedLimit{Type: SpeedLimitThrottled}
	}
	return SpeedLimit{Type: SpeedLimitStandard}
}

func (r *reconcileContext) updateResourceGroup() error {
	limit := GetSpeedLimitFromCluster(r.cluster)
	if err := r.sendResourceLimit(r.PdClient.PutResourceGroup, limit); err != nil {
		// try to create for new created cluster or legacy cluster
		if err1 := r.sendResourceLimit(r.PdClient.PostResourceGroup, limit); err1 != nil {
			r.Error(err1, "try to create resource limit")
			return err
		}
	}
	return nil
}

const ResourceGroupModeRU = 1

func (r *reconcileContext) sendResourceLimit(
	writer func(ctx context.Context, group *pd.ResourceGroup) error,
	limit SpeedLimit,
) error {
	if r.keyspaceID == 0 {
		err := errors.New("keyspace id is 0")
		r.Error(err, "keyspace id is unset")
		return err
	}

	group := &pd.ResourceGroup{
		Name: fmt.Sprintf("%d", r.keyspaceID),
		Mode: ResourceGroupModeRU,
		RUSettings: pd.RUSettings{
			RU: pd.RU{
				Settings: limit.GetPDSetting(),
			},
		},
	}
	return writer(r.ctx, group)
}

func (r *reconcileContext) deleteResourceGroup() error {
	err := r.PdClient.DeleteResourceGroup(r.ctx, fmt.Sprintf("%d", r.keyspaceID))
	return errors.Annotatef(err, "delete resource group failed")
}

func (r *reconcileContext) reconcileAuditLog() error {
	if r.cluster.Status.Phase != infra.ClusterPhaseRunning {
		return nil
	}
	if r.auditLogScheduleRepo == nil {
		r.Info("audit log disabled")
		return nil
	}

	// mock schedule to precheck to avoid unnecessary db operation
	if !r.buildAuditLogSchedule(&entity.AuditLogSchedule{}) {
		r.Info("no need to update audit log schedule, skip")
		return nil
	}

	// create or update audit log schedule
	schedule, err := r.auditLogScheduleRepo.GetByClusterId(r.ctx, r.cluster.Name)
	if err != nil {
		return errors.Annotate(err, "failed to get audit log schedule")
	}
	if schedule == nil {
		err = r.createAuditLog()
	} else {
		err = r.updateAuditLog(schedule)
	}
	if err != nil {
		return err
	}

	return errors.Annotate(r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.AuditLogConfig = infra.AuditLogConfigStatus{
			Enabled:                 r.cluster.Spec.AuditLogConfig.Enabled,
			Unredacted:              r.cluster.Spec.AuditLogConfig.Unredacted,
			CloudStorage:            r.cluster.Spec.AuditLogConfig.CloudStorage,
			RotationIntervalSeconds: r.cluster.Spec.AuditLogConfig.RotationIntervalSeconds,
			RotationSizeBytes:       r.cluster.Spec.AuditLogConfig.RotationSizeBytes,
		}
		return nil
	}), "failed to ensureStatus in reconcileAuditLog")
}

func (r *reconcileContext) createAuditLog() error {
	if !r.cluster.Spec.AuditLogConfig.Enabled {
		r.Info("audit log is disabled, skip creating schedule entity")
		return nil
	}
	schedule := &entity.AuditLogSchedule{
		OrgId:                   r.cluster.TagsTenant,
		ProjectId:               r.cluster.TagsProject,
		ClusterId:               r.cluster.Name,
		State:                   entity.AuditLogScheduleStateEnabled,
		RotationIntervalSeconds: r.cluster.Spec.AuditLogConfig.RotationIntervalSeconds,
		RotationSizeBytes:       r.cluster.Spec.AuditLogConfig.RotationSizeBytes,
		Unredacted:              r.cluster.Spec.AuditLogConfig.Unredacted,
	}
	r.refreshAuditLogScheduleCloudStorage(schedule)
	schedule.ResetRotationTime()
	r.Info("create audit log schedule", "schedule", schedule)
	return r.auditLogScheduleRepo.Create(r.ctx, schedule)
}

func (r *reconcileContext) updateAuditLog(schedule *entity.AuditLogSchedule) error {
	r.buildAuditLogSchedule(schedule)
	r.Info("update audit log schedule", "schedule", schedule)
	if err := r.auditLogScheduleRepo.Save(r.ctx, schedule); err != nil {
		r.Error(err, "failed to save audit log schedule")
		return err
	}
	return nil
}

func (r *reconcileContext) refreshAuditLogScheduleCloudStorage(auditLogSchedule *entity.AuditLogSchedule) {
	auditLogSchedule.CloudStorageType = ""
	auditLogSchedule.CloudStorage = nil
	cloudStorage := r.cluster.Spec.AuditLogConfig.CloudStorage
	if cloudStorage == nil {
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeTiDBCloud
		return
	}
	switch cloudStorage.CloudStorageType {
	case "", infra.CloudStorageTypeTiDBCloud:
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeTiDBCloud
	case infra.CloudStorageTypeS3:
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeS3
		auditLogSchedule.CloudStorage = &entity.AuditLogCloudStorage{
			S3: &entity.AuditLogCloudStorageS3{
				Uri:             cloudStorage.S3.Uri,
				Region:          cloudStorage.S3.Region,
				AccessKeyId:     cloudStorage.S3.AccessKeyId,
				SecretAccessKey: cloudStorage.S3.SecretAccessKey,
				RoleArn:         cloudStorage.S3.RoleArn,
				ExternalId:      cloudStorage.S3.ExternalId,
			},
		}
	case infra.CloudStorageTypeOSS:
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeOSS
		auditLogSchedule.CloudStorage = &entity.AuditLogCloudStorage{
			OSS: &entity.AuditLogCloudStorageOSS{
				Uri:             cloudStorage.OSS.Uri,
				Region:          cloudStorage.OSS.Region,
				AccessKeyId:     cloudStorage.OSS.AccessKeyId,
				AccessKeySecret: cloudStorage.OSS.AccessKeySecret,
				RoleArn:         cloudStorage.OSS.RoleArn,
				ExternalId:      cloudStorage.OSS.ExternalId,
			},
		}
	case infra.CloudStorageTypeGCS:
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeGCS
		auditLogSchedule.CloudStorage = &entity.AuditLogCloudStorage{
			GCS: &entity.AuditLogCloudStorageGCS{
				Uri:               cloudStorage.GCS.Uri,
				Region:            cloudStorage.GCS.Region,
				ServiceAccountKey: cloudStorage.GCS.ServiceAccountKey,
			},
		}
	case infra.CloudStorageTypeAzureBlob:
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeAzureBlob
		auditLogSchedule.CloudStorage = &entity.AuditLogCloudStorage{
			AzureBlob: &entity.AuditLogCloudStorageAzureBlob{
				Uri:      cloudStorage.AzureBlob.Uri,
				Region:   cloudStorage.AzureBlob.Region,
				SASToken: cloudStorage.AzureBlob.SASToken,
			},
		}
	default:
		auditLogSchedule.CloudStorageType = entity.AuditLogCloudStorageTypeTiDBCloud
	}
}

func (r *reconcileContext) deleteAuditLog() error {
	if r.auditLogScheduleRepo == nil {
		return nil
	}
	deleteTime := time.Now()
	if err := r.auditLogScheduleRepo.UpdateDeleteTime(r.ctx, r.cluster.Name, &deleteTime); err != nil {
		if strings.Contains(err.Error(), "'prod_serverless_rds.audit_log_schedules' doesn't exist") {
			r.Info("skip update delete time in audit log schedule since table doesn't exist")
			return nil
		}
		return errors.Annotate(err, "failed to update delete time in audit log schedule")
	}
	return nil
}