package branch

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/pingcap/errors"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/controllers/cluster"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"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"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	clientv3 "go.etcd.io/etcd/client/v3"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	retryTidbPoolChanging  = 40 * time.Second
	waitNewTidbPoolRunning = 5 * time.Second
)

func (r *reconcileContext) removeOldEtcdToken() error {
	for _, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
		if etcdTokenPrefix != r.etcdTokenPrefix {
			_, 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) removeEtcdToken(tidbPool string) error {
	etcdTokenPrefix := r.EtcdTokenPrefixMap[tidbPool]
	if etcdTokenPrefix == "" {
		return errors.Errorf("etcdTokenPrefix not found in tidbPool %s", tidbPool)
	}

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

	return nil
}

func (r *reconcileContext) setEtcdToken() error {
	tidbPool, replicas := r.specTidbPool, r.branch.Spec.Cluster.TiDB.Replicas
	etcdTokenPrefix := r.EtcdTokenPrefixMap[tidbPool]
	if etcdTokenPrefix == "" {
		return errors.Errorf("etcdTokenPrefix not found in tidbPool %s", tidbPool)
	}

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

	r.Info("put etcd", "etcdTokenPrefix", etcdTokenPrefix, "tidbPool", tidbPool, "replicas", replicas)
	isPut := false
	keyEndpoints := etcdTokenPrefix + r.keyspaceToken + "/" + 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 {
		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 + r.keyspaceToken + "/" + 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) setEtcdReplicas(tidbPool 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, "replicas", replicas)
	keyEndpoints := etcdTokenPrefix + r.keyspaceToken + "/" + 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 + r.keyspaceToken + "/" + 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() error {
	for tidbPool, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
		if etcdTokenPrefix == r.etcdTokenPrefix {
			continue
		}
		keyReplicas := etcdTokenPrefix + r.keyspaceToken + "/" + 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, n-1); err != nil {
			return err
		}
		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")
	}

	keyEndpoints := r.etcdTokenPrefix + r.keyspaceToken + "/" + utils.KeyEndpoints
	r.Info("put etcd", "etcdTokenPrefix", r.etcdTokenPrefix, "keyspaceToken", r.keyspaceToken)
	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 := r.etcdTokenPrefix + r.keyspaceToken + "/" + utils.KeyReplicas
	op := clientv3.WithIgnoreLease()
	if isPut {
		op = clientv3.WithLease(leaseID)
	}
	if _, err := r.EtcdKV.Put(r.ctx, keyReplicas, strconv.Itoa(r.branch.Spec.Cluster.TiDB.Replicas), op); err != nil {
		return errors.Annotatef(err, "failed to put etcd path %s", keyReplicas)
	}

	return r.removeOldEtcdToken()
}

func (r *reconcileContext) updateKeyspaceConfig() error {
	metadata, err := utils.GetBranchMetadata(r.branch)
	if err != nil {
		return errors.Annotate(err, "failed to get branch metadata")
	}

	cfg := &pd.KeyspaceConfig{
		TenantID:             strconv.FormatUint(metadata.TenantID, 10),
		ProjectID:            strconv.FormatUint(metadata.ProjectID, 10),
		ClusterID:            r.branch.Name,
		IsBranch:             "True",
		IsBranchBootstrapped: "False",
		MainClusterId:        r.branch.Annotations[MainBranchIdAnnKey],
	}
	err = r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, cfg)

	return errors.Annotate(err, "failed to update keyspace config")
}

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 {
		err = client.IgnoreNotFound(r.Client.Delete(r.ctx, &keyspaceList.Items[0]))
	}
	return err
}

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) 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.resetEtcdToken(); err != nil {
		return errors.Annotate(err, "failed to create etcd token")
	}

	return errors.Annotate(r.ensureStatus(func(
		status *cloudv1.BranchStatus,
		fromProvider *cloudv1.AwsFreeTierClusterStatus,
	) error {
		if status.DB == nil {
			status.DB = &cloudv1.DBStatus{
				Version:         r.tidbVersion,
				SecondaryStatus: cloudv1.DBSecondaryStatusHealthy,
			}
		}
		status.DB.TiDB = &cloudv1.TiDBStatus{
			Status:                cloudv1.ComponentNormal,
			UpdatedReplicas:       int32(r.branch.Spec.Cluster.TiDB.Replicas),
			TokenBasedAuthEnabled: true, // always return true for serverless clusters
		}

		fromProvider.TidbPool = r.specTidbPool
		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) reconcileResumeEtcdToken() error {
	if r.branch.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.branch.Annotations, types.AnnoKeyResumeEtcdToken)
	return r.InfraClient.Update(r.ctx, r.branch)
}

func (r *reconcileContext) allOldTidbPoolPods() []client.ObjectKey {
	oldPods := make([]client.ObjectKey, 0)
	for pool, tManager := range r.TidbManagers {
		if pool != r.specTidbPool {
			if endpoints := tManager.GetAssignedTidbByToken(r.keyspaceToken); 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() []client.ObjectKey {
	newPods := make([]client.ObjectKey, 0)
	if tManager := r.TidbManagers[r.specTidbPool]; tManager != nil {
		if endpoints := tManager.GetAssignedTidbByToken(r.keyspaceToken); endpoints != nil {
			for _, endpoint := range endpoints.Endpoints {
				podName, _ := ep.ParsePodFromEndpoint(endpoint)
				newPods = append(newPods, client.ObjectKey{Namespace: r.specTidbPool, Name: podName})
			}
		}
	}

	return newPods
}

func (r *reconcileContext) reconcileTidbPool() error {
	branchStatus, err := r.branch.AwsFreeTierClusterStatusFromBranch()
	if err != nil {
		return errors.Annotate(err, "failed to get AwsFreeTierClusterStatusFromBranch")
	}
	if r.specTidbPool != branchStatus.TidbPool {
		r.Info("tidb pool is changing", "old", branchStatus.TidbPool, "now", r.specTidbPool)
		r.branch.Annotations = utils.UpdateMap(r.branch.Annotations, types.AnnoKeyTidbPoolChanging, time.Now().Format(time.RFC3339))
		if err := r.InfraClient.Update(r.ctx, r.branch); 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 *cloudv1.BranchStatus,
			fromProvider *cloudv1.AwsFreeTierClusterStatus,
		) error {
			fromProvider.TidbPool = r.specTidbPool
			return nil
		}); err != nil {
			return errors.Annotate(err, "failed to ensureStatus in reconcileTidbPool")
		}
	}

	if tidbPoolChanging := r.branch.Annotations[types.AnnoKeyTidbPoolChanging]; tidbPoolChanging != "" {
		if allOldTidbPoolPodLen := len(r.allOldTidbPoolPods()); allOldTidbPoolPodLen > 0 {
			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()) > r.branch.Spec.Cluster.TiDB.Replicas {
				if err := r.reduceOneReplicaFromOldPools(); err != nil {
					return errors.Annotatef(err, "failed to reduce one replica from old pods")
				}
				r.Info("reduce one replica in old tidb pool")
				r.branch.Annotations[types.AnnoKeyTidbPoolChanging] = time.Now().Format(time.RFC3339)
				if err := r.InfraClient.Update(r.ctx, r.branch); 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,
			}
		}

		// 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.branch.Annotations, types.AnnoKeyTidbPoolChanging)
		if err := r.InfraClient.Update(r.ctx, r.branch); err != nil {
			return errors.Annotate(err, "failed to remove infra cluster AnnoKeyTidbPoolChanging")
		}
		r.Info("tidb pool changing end")
	}

	return nil
}

func (r *reconcileContext) reconcileTidbReplicas() error {
	if int32(r.branch.Spec.Cluster.TiDB.Replicas) == r.branch.Status.DB.TiDB.UpdatedReplicas {
		return nil
	}

	r.Info("tidb replicas have changed", "old", r.branch.Status.DB.TiDB.UpdatedReplicas, "now", r.branch.Spec.Cluster.TiDB.Replicas)

	if err := r.setTidbPoolReplicas(r.specTidbPool, r.branch.Spec.Cluster.TiDB.Replicas); err != nil {
		return errors.Annotatef(err, "failed to set tidbPool replicas: %s, %d", r.specTidbPool, r.branch.Spec.Cluster.TiDB.Replicas)
	}

	return errors.Annotatef(r.ensureStatus(func(
		status *cloudv1.BranchStatus,
		fromProvider *cloudv1.AwsFreeTierClusterStatus,
	) error {
		status.DB = &cloudv1.DBStatus{
			Version: r.tidbVersion,
			PD:      &cloudv1.PDStatus{Status: cloudv1.ComponentNormal},
			TiKV:    &cloudv1.TiKVStatus{Status: cloudv1.ComponentNormal},
			TiDB: &cloudv1.TiDBStatus{
				Status:                cloudv1.ComponentNormal,
				UpdatedReplicas:       int32(r.branch.Spec.Cluster.TiDB.Replicas),
				TokenBasedAuthEnabled: true, // always return true for serverless clusters
			},
			TiFlash:         &cloudv1.TiFlashStatus{Status: cloudv1.ComponentNormal},
			SecondaryStatus: cloudv1.DBSecondaryStatusHealthy,
		}
		return nil
	}), "failed to ensureStatus in reconcileTidbReplicas")
}

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

	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.ensureResourceGroup()); err != nil {
		return err
	}

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

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

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

	return reqErrs
}

func (r *reconcileContext) observeTidbClusterStatus() (errs error) {
	clusterProxyBaseUrl := fmt.Sprintf(
		"%s/proxy/namespace/free-tier-%s/cluster/%s",
		r.BaseProxy, r.branch.Namespace, "db"+r.branch.Name)

	objectStoreUrl := fmt.Sprintf("%s/system:serviceaccount:free-tier-%s:%s/",
		r.S3BucketName, r.branch.Namespace, types.S3BackupServiceAccount)

	return errors.Annotate(r.ensureStatus(func(
		status *cloudv1.BranchStatus,
		fromProvider *cloudv1.AwsFreeTierClusterStatus,
	) error {
		status.Phase = cloudv1.BranchPhaseRunning
		status.Reason = "branch is running"
		status.Proxy = &cloudv1.ProxyStatus{Endpoint: clusterProxyBaseUrl}
		status.ObjectStore = &cloudv1.ObjectStoreStatus{S3: &cloudv1.Path{Path: objectStoreUrl}}
		return nil
	}), "failed to ensureStatus in observeTidbClusterStatus")
}

func (r *reconcileContext) ensureSqlEndpoints() error {
	specForProvider, err := types.InfraServerlessSpecForProvider(r.branch)
	if err != nil {
		return errors.Annotate(err, "failed to get specForProvider")
	}

	var publicEndpoint *cloudv1.SqlEndpointStatus
	if r.branch.Spec.SqlEndpoints.Public.IsEnabled() {
		var err error
		if publicEndpoint, err = r.genPublicEndpointStatus(); err != nil {
			return errors.Annotate(err, "failed to genPublicEndpointStatus")
		}
	}
	var privateLinkService *cloudv1.PrivateLinkServiceStatus
	if r.branch.Spec.SqlEndpoints.PrivateLinkService != nil {
		var err error
		if privateLinkService, err = r.genPrivatelinkStatus(specForProvider); err != nil {
			return errors.Annotate(err, "failed to genPrivatelinkStatus")
		}
	}

	return errors.Annotate(r.ensureStatus(func(
		status *cloudv1.BranchStatus,
		_ *cloudv1.AwsFreeTierClusterStatus,
	) error {
		status.SqlEndpoints.PublicEndpoint = publicEndpoint
		status.SqlEndpoints.PrivateLinkService = privateLinkService
		return nil
	}), "failed to ensureStatus in ensureSqlEndpoints")
}

func (r *reconcileContext) ensureZoneRedundant() error {
	if r.branch.Status.Phase != cloudv1.BranchPhaseRunning &&
		r.branch.Status.Phase != cloudv1.BranchPhaseCreating &&
		r.branch.Status.Phase != cloudv1.BranchPhaseRestoring {
		return nil
	}
	clusterStatus, err := r.branch.AwsFreeTierClusterStatusFromBranch()
	if err != nil {
		r.Error(err, "call AwsFreeTierClusterStatusFromBranch err, return false")
		return err
	}

	specForProvider, err := types.InfraServerlessSpecForProvider(r.branch)
	if err != nil {
		r.Error(err, "failed to get specForProvider")
		return err
	}

	if clusterStatus.ZoneRedundantEnabled == specForProvider.EnableZoneRedundant {
		return nil
	}

	if err := r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, &pd.KeyspaceConfig{
		ZoneRedundant: fmt.Sprint(specForProvider.EnableZoneRedundant),
	}); err != nil {
		return errors.Annotatef(err, "failed to update enableZoneRedundant to %v", specForProvider.EnableZoneRedundant)
	}
	r.Info("enableZoneRedundant is updated", "enableZoneRedundant", specForProvider.EnableZoneRedundant)

	return errors.Annotate(r.ensureStatus(func(
		_ *cloudv1.BranchStatus,
		fromProvider *cloudv1.AwsFreeTierClusterStatus,
	) error {
		fromProvider.ZoneRedundantEnabled = specForProvider.EnableZoneRedundant
		return nil
	}), "failed to ensureStatus in zoneRedundantEnabled")
}

func (r *reconcileContext) ensureRcuConfig() error {
	if r.branch.Status.Phase != cloudv1.BranchPhaseRunning &&
		r.branch.Status.Phase != cloudv1.BranchPhaseCreating &&
		r.branch.Status.Phase != cloudv1.BranchPhaseRestoring {
		return nil
	}

	if r.branch.Spec.AutoScaling == nil || r.branch.Spec.AutoScaling.MinRcu == 0 {
		return nil
	}

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

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

	if err := r.updateKeyspaceConfigRcu(); err != nil {
		return err
	}

	if err := r.updateResourceGroup(); err != nil {
		return err
	}

	r.branch.Annotations[types.AnnoKeyPrevMinRCU] = minRcu
	r.branch.Annotations[types.AnnoKeyPrevMaxRCU] = maxRcu
	if err := r.InfraClient.Update(r.ctx, r.branch); err != nil {
		return errors.Annotate(err, "failed to update branch rcu annotations")
	}

	r.Info("rcu annotations is updated", "minRcu", minRcu, "maxRcu", maxRcu)
	return nil
}

func GetSpeedLimitFromBranch(branch *cloudv1.Branch) cluster.SpeedLimit {
	if infra.IsEssentialBranch(branch) {
		return cluster.SpeedLimit{
			Type:   cluster.SpeedLimitEssential,
			MinRCU: branch.Spec.AutoScaling.MinRcu,
			MaxRCU: branch.Spec.AutoScaling.MaxRcu,
		}
	}

	switch branch.Annotations[types.AnnoTidbRuLimitMode] {
	case cluster.SpeedLimitStandardName:
		return cluster.SpeedLimit{Type: cluster.SpeedLimitStandard}
	case cluster.SpeedLimitThrottledName:
		return cluster.SpeedLimit{Type: cluster.SpeedLimitThrottled}
	}
	return cluster.SpeedLimit{Type: cluster.SpeedLimitStandard}
}

func (r *reconcileContext) updateResourceGroup() error {
	limit := GetSpeedLimitFromBranch(r.branch)
	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
}

func (r *reconcileContext) sendResourceLimit(
	writer func(ctx context.Context, group *pd.ResourceGroup) error,
	limit cluster.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: cluster.ResourceGroupModeRU,
		RUSettings: pd.RUSettings{
			RU: pd.RU{
				Settings: limit.GetPDSetting(),
			},
		},
	}
	return writer(r.ctx, group)
}

func (r *reconcileContext) ensureResourceGroup() error {
	if r.branch.Status.Phase != cloudv1.BranchPhaseRunning &&
		r.branch.Status.Phase != cloudv1.BranchPhaseCreating &&
		r.branch.Status.Phase != cloudv1.BranchPhaseRestoring {
		return nil
	}

	clusterStatus, err := r.branch.AwsFreeTierClusterStatusFromBranch()
	if err != nil {
		r.Error(err, "call AwsFreeTierClusterStatusFromBranch err, return false")
		return err
	}

	mode := GetSpeedLimitFromBranch(r.branch).Type.String()
	if clusterStatus.RuLimitMode == mode {
		return nil
	}

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

	r.Info("ru limit mode need update", "old", clusterStatus.RuLimitMode, "now", mode)
	if err := r.updateResourceGroup(); err != nil {
		return err
	}
	r.InfraRecorder.Event(r.branch, corev1.EventTypeNormal, "UpdateRuLimitMode", "ru limit mode is set to "+mode)

	return errors.Annotate(r.ensureStatus(func(
		_ *cloudv1.BranchStatus,
		fromProvider *cloudv1.AwsFreeTierClusterStatus,
	) error {
		fromProvider.RuLimitMode = mode
		return nil
	}), "failed to ensureStatus in ensureResourceGroup")
}

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) setKeyspaceConfigIsBranchBootstrapped() error {
	cfg := &pd.KeyspaceConfig{
		IsBranchBootstrapped: "False",
	}
	err := r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, cfg)
	return errors.Annotate(err, "failed to update keyspace config: isBranchBootstrapped")
}

func (r *reconcileContext) updateKeyspaceConfigRcu() error {
	cfg := &pd.KeyspaceConfig{
		MaxRCU: strconv.FormatInt(int64(r.branch.Spec.AutoScaling.MaxRcu), 10),
		MinRCU: strconv.FormatInt(int64(r.branch.Spec.AutoScaling.MinRcu), 10),
	}
	err := r.PdClient.UpdateKeyspaceConfig(r.ctx, r.keyspaceToken, cfg)
	return errors.Annotate(err, "failed to update keyspace config: rcu")
}