package backup

import (
	"context"
	"fmt"
	"math"
	"os"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	s3api "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	ststypes "github.com/aws/aws-sdk-go-v2/service/sts/types"
	"github.com/go-logr/logr"
	"github.com/pingcap/errors"
	configsdk "github.com/tidbcloud/config-service/sdk"
	"github.com/tidbcloud/config-service/sdk/api"
	"go.opentelemetry.io/otel/attribute"
	corev1 "k8s.io/api/core/v1"
	scheduling_v1 "k8s.io/api/scheduling/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/pointer"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/predicate"

	awsapi "github.com/tidbcloud/infra-api/pkg/providers/aws"

	"github.com/tidbcloud/infra-api/controllers/common"
	"github.com/tidbcloud/infra-api/pkg/provider"
	tracerutil "github.com/tidbcloud/infra-api/pkg/tracer"

	gcpiam "cloud.google.com/go/iam"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/reconciler"
	"github.com/tidbcloud/infra-api/controllers/utils"
	"github.com/tidbcloud/infra-api/pkg/consts"
	"github.com/tidbcloud/infra-api/pkg/gardener"
	"github.com/tidbcloud/infra-api/pkg/k8s"
	metricsctrl "github.com/tidbcloud/infra-api/pkg/metrics/controller"
	provideraws "github.com/tidbcloud/infra-api/pkg/providers/aws"
	providergcp "github.com/tidbcloud/infra-api/pkg/providers/gcp"
)

const (
	errNotRestore = "not restore resource"

	RestoreFinalizer = "restore.cloud.pingcap.com/finalizer"

	storageAdminRoleName gcpiam.RoleName = "roles/storage.admin"

	SystemCriticalPriorityClass = "system-cluster-critical"
)

// RestoreReconciler reconciles a Restore object
type RestoreReconciler struct {
	reconciler *reconciler.Reconciler
	Log        logr.Logger
}

// NewRestoreReconciler creates a RestoreReconciler instance, clusterCli will be nil for non-shared cluster.
func NewRestoreReconciler(mgr ctrl.Manager, log logr.Logger, g *gardener.GardEnvironment, clusterCli client.Client, configSvc configsdk.Interface, awsCli provideraws.Client, gcpAPI providergcp.GCPAPI, p provider.Restore) *RestoreReconciler {
	return &RestoreReconciler{
		Log: log,
		reconciler: &reconciler.Reconciler{
			NewResource: func() client.Object {
				return &cloudv1.Restore{}
			},
			Client: mgr.GetClient(),
			ExternalObj: &restoreConnector{
				g:          g,
				client:     mgr.GetClient(),
				clusterCli: clusterCli,
				configSvc:  configSvc,
				awsCli:     awsCli,
				gcpAPI:     gcpAPI,
				provider:   p,
			},
			ReconcileTimeout: 3 * time.Minute,
		},
	}
}

func NewRestoreReconcilerCompat(infraCli client.Client, log logr.Logger, g *gardener.GardEnvironment, clusterCli client.Client, configSvc configsdk.Interface, awsCli provideraws.Client, gcpAPI providergcp.GCPAPI, p provider.Restore) *RestoreReconciler {
	return &RestoreReconciler{
		Log: log,
		reconciler: &reconciler.Reconciler{
			NewResource: func() client.Object {
				return &cloudv1.Restore{}
			},
			Client: infraCli,
			ExternalObj: &restoreConnector{
				g:          g,
				client:     infraCli,
				clusterCli: clusterCli,
				configSvc:  configSvc,
				awsCli:     awsCli,
				gcpAPI:     gcpAPI,
				provider:   p,
			},
		},
	}
}

//+kubebuilder:rbac:groups=cloud.pingcap.com,resources=restores,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=cloud.pingcap.com,resources=restores/status,verbs=get;update;patch
//+kubebuilder:rbac:groups=cloud.pingcap.com,resources=restores/finalizers,verbs=update

// Reconcile is part of the main k8sStatus reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
func (r *RestoreReconciler) Reconcile(ctx context.Context, req ctrl.Request) (_ ctrl.Result, errResult error) {
	start := time.Now()

	ctx, span := tracer.Start(ctx, "RestoreReconcile")
	defer func() {
		tracerutil.EndSpan(span, errResult)
		metricsctrl.ReconcileDuration.WithLabelValues("restore").Observe(time.Since(start).Seconds())
	}()

	span.SetAttributes(attribute.String("restore_name", req.Name))

	logger := r.Log.WithValues("restore", req.NamespacedName)
	ctx = log.IntoContext(ctx, logger)

	return r.reconciler.Reconcile(ctx, req, RestoreFinalizer)
}

// SetupWithManager sets up the controller with the Manager.
func (r *RestoreReconciler) SetupWithManager(mgr ctrl.Manager, p ...predicate.Predicate) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&cloudv1.Restore{}, builder.WithPredicates(p...)).
		WithEventFilter(utils.GetMigrationAnnotationSkipPredicator()).
		WithOptions(controller.Options{RateLimiter: common.DefaultRateLimiter(), MaxConcurrentReconciles: 10}).
		Complete(r)
}

type restoreConnector struct {
	g      gardener.GardenerAPI
	client client.Client

	// client to the target cluster, only used for the shared K8s cluster now.
	// if this client is set, no need to try to get a client for a shoot (because we are running for a non-gardener based cluster)
	clusterCli client.Client
	configSvc  configsdk.Interface
	awsCli     provideraws.Client
	gcpAPI     providergcp.GCPAPI

	provider provider.Restore
}

func (rc *restoreConnector) Connect(ctx context.Context, obj client.Object) (_ reconciler.ExternalObjClient, errResult error) {
	ctx, span := tracer.Start(ctx, "restoreConnector.Connect")
	defer func() { tracerutil.EndSpan(span, errResult) }()

	r, ok := obj.(*cloudv1.Restore)
	if !ok {
		return nil, errors.New(errNotRestore)
	}

	// when deleting a cluster, it should wait for all restores to be deleted,
	// so we should always connect cluster and network successfully.
	var c cloudv1.Cluster
	clusterNamespacedName := types.NamespacedName{
		Name:      r.Spec.ClusterName,
		Namespace: r.ObjectMeta.Namespace,
	}
	err := rc.client.Get(ctx, clusterNamespacedName, &c)
	if err != nil {
		return nil, err
	}

	span.SetAttributes(
		attribute.String("restore_name", r.Name),
		attribute.String("restore_phase", string(r.Status.Phase)),
	)
	tracerutil.SetCommonAttributes(span, &c)

	// TODO(aylei,fangfang): check cluster phase before connecting
	var k8sStatus *cloudv1.KubernetesStatus
	switch c.Labels[cloudv1.ProviderTypeLabel] {
	case consts.ProviderTypeAWS:
		awsStatus, err := cloudv1.AwsClusterStatusFromCluster(&c)
		if err != nil {
			return nil, err
		}
		k8sStatus = awsStatus.Kubernetes
	case consts.ProviderTypeGCP:
		gcpStatus, err := cloudv1.GcpClusterStatusFromCluster(&c)
		if err != nil {
			return nil, err
		}
		k8sStatus = gcpStatus.Kubernetes
	case consts.ProviderTypeAWSFreeTier:
		awsFreeTierStatus, err := cloudv1.AwsFreeTierClusterStatusFromCluster(&c)
		if err != nil {
			return nil, err
		}
		k8sStatus = awsFreeTierStatus.Kubernetes
	case consts.ProviderTypeAWSSharedTier, consts.ProviderTypeGCPSharedTier:
		sharedTierStatus, err := cloudv1.SharedTierClusterStatusFromCluster(&c)
		if err != nil {
			return nil, err
		}
		k8sStatus = sharedTierStatus.Kubernetes
	case consts.ProviderTypeEKS, consts.ProviderTypeGKE:
		k8sStatus, err = rc.provider.GetKubernetesStatus(&c)
		if err != nil {
			return nil, err
		}
	default:
		return nil, fmt.Errorf("unknown provider %s", c.Labels[cloudv1.ProviderTypeLabel])
	}
	if k8sStatus == nil {
		return nil, fmt.Errorf("k8sStatus status is not yet set by provider %s", c.Labels[cloudv1.ProviderTypeLabel])
	}

	var n cloudv1.Network
	networkNamespaceName := c.GetNetworkRef()
	err = rc.client.Get(ctx, networkNamespaceName, &n)
	if err != nil {
		return nil, err
	}

	provider := new(cloudv1.Provider)
	err = rc.client.Get(ctx, types.NamespacedName{Name: n.Spec.Provider}, provider)
	if err != nil {
		return nil, err
	}

	var clusterCli client.Client
	switch c.Labels[cloudv1.ProviderTypeLabel] {
	case consts.ProviderTypeAWS, consts.ProviderTypeGCP:
		networkStatus, err := cloudv1.GardenerNetworkStatusFromNetwork(&n)
		if err != nil {
			return nil, err
		}
		shoot, err := rc.g.GetShootContext(ctx, networkStatus.ShootName)
		if err != nil {
			return nil, err
		}
		clusterCli = shoot.Manager.GetClient()
	case consts.ProviderTypeAWSFreeTier, consts.ProviderTypeAWSSharedTier, consts.ProviderTypeGCPSharedTier:
		clusterCli = rc.clusterCli // this should already be set when initialized the `restoreConnector`
	case consts.ProviderTypeEKS, consts.ProviderTypeGKE:
		clusterCli, err = rc.provider.GetClient(ctx, &c)
		if err != nil {
			return nil, err
		}
	}

	b := new(cloudv1.Backup)
	backupNamespaceName := types.NamespacedName{
		Name:      r.Spec.From.Backup,
		Namespace: r.ObjectMeta.Namespace,
	}
	err = rc.client.Get(ctx, backupNamespaceName, b)
	if err != nil && reconciler.IgnoreNotFound(err) != nil {
		return nil, err
	}

	var logBackup *cloudv1.Backup
	if r.Spec.From.LogBackup != nil {
		logBackup = new(cloudv1.Backup)
		logBackupNamespacedName := types.NamespacedName{
			Name:      *r.Spec.From.LogBackup,
			Namespace: r.ObjectMeta.Namespace,
		}
		err = rc.client.Get(ctx, logBackupNamespacedName, logBackup)
		if err != nil && reconciler.IgnoreNotFound(err) != nil {
			return nil, err
		}
	}

	return &clusterRestore{
		cli:             rc.client,
		clusterCli:      clusterCli,
		cluster:         &c,
		network:         &n,
		provider:        provider,
		backup:          b, // this maybe nil if the backup already deleted
		logBackup:       logBackup,
		k8sStatus:       k8sStatus,
		configSvc:       rc.configSvc,
		awsCli:          rc.awsCli,
		gcpAPI:          rc.gcpAPI,
		restoreProvider: rc.provider,
	}, nil
}

type clusterRestore struct {
	cli        client.Client
	clusterCli client.Client
	cluster    *cloudv1.Cluster
	network    *cloudv1.Network
	provider   *cloudv1.Provider
	backup     *cloudv1.Backup
	logBackup  *cloudv1.Backup

	k8sStatus *cloudv1.KubernetesStatus
	configSvc configsdk.Interface
	awsCli    provideraws.Client
	gcpAPI    providergcp.GCPAPI

	restoreProvider provider.Restore
}

func (cr *clusterRestore) inNewArch() bool {
	return networkInNewArch(cr.network)
}

func (cr *clusterRestore) Observe(ctx context.Context, obj client.Object) (_ reconciler.ExternalObservation, errResult error) {
	ctx, span := tracer.Start(ctx, "clusterRestore.Observe")
	defer func() { tracerutil.EndSpan(span, errResult) }()

	r, ok := obj.(*cloudv1.Restore)
	if !ok {
		return reconciler.ExternalObservation{}, errors.New(errNotRestore)
	}
	if !cr.restoreProvider.IsDPCompatMode() {
		if r.OwnerReferences == nil {
			// adopt the obj to cluster
			err := controllerutil.SetOwnerReference(cr.cluster, r, cr.cli.Scheme())
			if err != nil {
				return reconciler.ExternalObservation{}, err
			}
			err = cr.cli.Update(ctx, r)
			if err != nil {
				return reconciler.ExternalObservation{}, err
			}
		}
	}

	span.SetAttributes(
		attribute.String("restore_name", r.Name),
		attribute.String("restore_phase", string(r.Status.Phase)),
	)
	tracerutil.SetCommonAttributes(span, cr.cluster)

	inFinalState := r.Status.Phase == cloudv1.RestoreComplete || r.Status.Phase == cloudv1.RestoreFailed

	if inFinalState && cr.needUseNodeGroup(r) {
		ngName := restoreNodeGroupName(r, cr.inNewArch())
		log.FromContext(ctx).Info(fmt.Sprintf("clean node group for restore ng name: %s", ngName))
		if err := cr.restoreProvider.DeleteBRNode(ctx, cr.network.DeepCopy(), ngName); err != nil {
			log.FromContext(ctx).Error(
				err, fmt.Sprintf("clean node group for restore error, ng name: %s", ngName))
			return
		}
	}

	if inFinalState {
		// final state, nothing to do
		return reconciler.ExternalObservation{ResourceUpToDate: true, ResourceExists: true}, nil
	}

	var tidbRestore v1alpha1.Restore
	if err := cr.clusterCli.Get(ctx, types.NamespacedName{
		Name:      r.ObjectMeta.Name,
		Namespace: cr.k8sStatus.Namespace,
	}, &tidbRestore); err != nil {
		if reconciler.IgnoreNotFound(err) != nil {
			return reconciler.ExternalObservation{}, err
		}

		r.Status.Phase = cloudv1.RestorePrepare
		return reconciler.ExternalObservation{ResourceExists: false}, cr.cli.Status().Update(ctx, r)
	}

	return reconciler.ExternalObservation{ResourceExists: true}, nil
}

func convertRestorePhase(phase v1alpha1.RestoreConditionType) cloudv1.RestoreConditionType {
	switch phase {
	case v1alpha1.RestoreScheduled:
		return cloudv1.RestoreScheduled
	case v1alpha1.RestoreRunning:
		return cloudv1.RestoreRunning
	case v1alpha1.RestoreComplete:
		return cloudv1.RestoreComplete
	case v1alpha1.RestoreFailed, v1alpha1.RestoreInvalid:
		return cloudv1.RestoreFailed
	case v1alpha1.RestoreRetryFailed:
		return cloudv1.RestorePrepare // retryable error as prepare
	default:
		return cloudv1.RestorePrepare // empty phase as prepare
	}
}

func syncRestoreProgresses(restore *cloudv1.Restore, tidbRestoreProgresses []v1alpha1.Progress) {
	i := 0
	if len(restore.Status.Progresses) > len(tidbRestoreProgresses) {
		restore.Status.Progresses = restore.Status.Progresses[:len(tidbRestoreProgresses)]
	}
	for ; i < len(restore.Status.Progresses); i++ {
		if restore.Status.Progresses[i].LastTransitionTime.Equal(&tidbRestoreProgresses[i].LastTransitionTime) {
			continue
		}
		restore.Status.Progresses[i].Step = tidbRestoreProgresses[i].Step
		restore.Status.Progresses[i].Progress = convertProgress(tidbRestoreProgresses[i].Progress)
		restore.Status.Progresses[i].LastTransitionTime = tidbRestoreProgresses[i].LastTransitionTime
	}
	for ; i < len(tidbRestoreProgresses); i++ {
		restore.Status.Progresses = append(restore.Status.Progresses, cloudv1.RestoreProgress{
			Step:               tidbRestoreProgresses[i].Step,
			Progress:           convertProgress(tidbRestoreProgresses[i].Progress),
			LastTransitionTime: tidbRestoreProgresses[i].LastTransitionTime,
		})
	}
}

func convertProgress(tidbProgress float64) int32 {
	return int32(math.Floor(tidbProgress))
}

func (cr *clusterRestore) Create(ctx context.Context, obj client.Object) (errResult error) {
	ctx, span := tracer.Start(ctx, "clusterRestore.Create")
	defer func() { tracerutil.EndSpan(span, errResult) }()

	r, ok := obj.(*cloudv1.Restore)
	if !ok {
		return errors.New(errNotRestore)
	}

	span.SetAttributes(
		attribute.String("restore_name", r.Name),
		attribute.String("restore_phase", string(r.Status.Phase)),
	)
	tracerutil.SetCommonAttributes(span, cr.cluster)

	if cr.backup == nil { // backup deleted during restoring creating
		r.Status.Phase = cloudv1.RestoreFailed
		return cr.cli.Status().Update(ctx, r)
	}

	providerType := cr.cluster.ObjectMeta.Labels[cloudv1.ProviderTypeLabel]
	var tidbRestore *v1alpha1.Restore
	var err error
	if providerType == consts.ProviderTypeEKS || providerType == consts.ProviderTypeGKE {
		tidbRestore, err = cr.buildTiDBRestoreForNewArch(r)
		if err != nil {
			return errors.Annotate(err, "build tidb restore for new arch")
		}
	} else {
		tidbRestore, err = cr.buildTiDBRestoreForGardenerArch(ctx, r)
		if err != nil {
			return errors.Annotate(err, "build tidb restore for gardener arch")
		}
	}

	tidbRestore.Annotations = commonAnnotationsForBR()
	// we need to update the status before create the underlying CRD, otherwise if the status update failed,
	// then this `Create` will be skipped (without status update again) in the next round of reconcile after `Observe`.
	r.Status.Phase = cloudv1.RestorePrepare
	if err = cr.cli.Status().Update(ctx, r); err != nil {
		return err
	}

	if len(r.Spec.Env) > 0 {
		tidbRestore.Spec.Env = r.Spec.Env
	}
	if r.Spec.ResourceRequirements != nil {
		tidbRestore.Spec.ResourceRequirements = *r.Spec.ResourceRequirements
	}
	if cr.needUseNodeGroup(r) {
		logger := log.FromContext(ctx).WithName("createNodeForRestore")
		region, err := cr.provider.FindRegion(cr.network.Spec.Region)
		if err != nil {
			logger.Error(err, "find region failed")
			return err
		}
		ngSpec, err := buildRestoreNodeGroupSpecWithTaints(restoreNodeGroupName(r, cr.inNewArch()), r, cr.provider, &region)
		if err != nil {
			logger.Error(err, "build restore node group spec failed")
		}
		logger.Info(fmt.Sprintf("find ngSpec success the machine: %s", ngSpec.MachineSpec.MachineType))
		if err := cr.restoreProvider.CreateBRNode(ctx, cr.network.DeepCopy(), ngSpec); err != nil {
			logger.Error(err, fmt.Sprintf("create restore node failed, node group spec name %s", ngSpec.Name))
			return err
		}
		tidbRestore.Spec.Affinity = utils.NodeAffinityWithUK(restoreNodeAffinityKey(r), restoreNodeAffinityValue(r))
		tidbRestore.Spec.Tolerations = utils.TolerationWithUK(cloudv1.ComponentRestore, r.Name)
		logger.Info(fmt.Sprintf("create node for restore with taints success taints %s", ngSpec.Taints))
	}

	// add priority class to restore pod to avoid node-pressure eviction
	// ref: https://kubernetes.io/docs/concepts/scheduling-eviction/node-pressure-eviction/#pod-selection-for-kubelet-eviction
	if err := cr.ensurePriorityClass(ctx); err != nil {
		return err
	}
	tidbRestore.Spec.PriorityClassName = ClusterCriticalPriorityClass
	_, err = controllerutil.CreateOrUpdate(ctx, cr.clusterCli, tidbRestore, func() error {
		return nil
	})
	return err
}

func (cr *clusterRestore) buildTiDBRestoreForGardenerArch(
	ctx context.Context, r *cloudv1.Restore,
) (*v1alpha1.Restore, error) {
	vars := TidbRestoreVars{
		RestoreName:      r.ObjectMeta.Name,
		RestoreNamespace: cr.k8sStatus.Namespace,
		ClusterName:      utils.ClusterPrefix,
		Region:           cr.network.Spec.Region,
		Bucket:           cr.backup.Status.BucketName,
		BucketPrefix:     cr.backup.ObjectMeta.Name, // same as backup name
	}
	if cr.logBackup != nil {
		vars.BucketLogBackupPrefix = cr.logBackup.ObjectMeta.Name
	}
	if r.CheckBucketInfoIsExisted() {
		logger := log.FromContext(ctx).WithName("queryBucketFromSpec")
		logger.Info("restore from bucket info")
		vars.BucketCloudProvider = r.Spec.From.RestoreBucketInfo.CloudProvider
		vars.Region = r.Spec.From.RestoreBucketInfo.Region
		vars.Bucket = r.Spec.From.RestoreBucketInfo.Name
		vars.BucketPrefix = r.Spec.From.RestoreBucketInfo.BackupPrefix
		vars.BucketLogBackupPrefix = r.Spec.From.RestoreBucketInfo.LogBackupPrefix
	}
	if r.Spec.From.PiTRRestoredTS != nil {
		vars.RestoredTS = *r.Spec.From.PiTRRestoredTS
	}

	if cr.backup != nil && cr.backup.Spec.Mode != nil && *cr.backup.Spec.Mode == cloudv1.BackupModeVolumeSnapshot {
		vars.IsVolumeBackup = true
	}

	tenantId := utils.ParsingLegacyId(cr.cluster, consts.TenantTagKey)
	if tenantId == nil {
		dummyId := uint64(0)
		tenantId = &dummyId
	}
	clusterId := utils.ParsingLegacyId(cr.cluster, consts.ClusterTagKey)
	brConfigOutput, err := cr.configSvc.DescribeRestoreConfiguration(&api.DescribeBRConfigInput{
		TenantID:  tenantId,
		ClusterID: clusterId,
	})
	if err != nil {
		return nil, errors.Annotate(err, "get gitops restore config")
	}
	brConfig, err := BRConfig(brConfigOutput)
	if err != nil {
		return nil, errors.Annotate(err, "convert restore config")
	}

	labels := utils.CopyLabelsForWorkloads(r.Labels)
	annotations := make(map[string]string)
	var tidbRestore *v1alpha1.Restore
	providerType := cr.cluster.ObjectMeta.Labels[cloudv1.ProviderTypeLabel]
	switch providerType {
	case consts.ProviderTypeGCP:
		bucketName := cr.backup.Status.BucketName
		if !r.CheckBucketInfoIsExisted() {
			bucketRegion := cr.backup.Labels["cloud.pingcap.com/region"]
			vars.Region = bucketRegion
		}
		// If bucket region is not same with cluster region. we need attach s3 policy to the cluster.
		isCrossRegion, network, err := cr.isCrossRegionRestore(ctx, vars.Region)
		if err != nil {
			return nil, fmt.Errorf("failed check the restore is cross region, error: %v", err)
		}
		// If cross region is true, we need authorize the restore cluster to be able to access the gcp bucket
		// if bucket info exists, it means we have granted permissions in dataflow
		if isCrossRegion && !r.CheckBucketInfoIsExisted() {
			// Get node role arn from network status
			networkStatus, err := cloudv1.GardenerNetworkStatusFromNetwork(network)
			if err != nil {
				return nil, fmt.Errorf("failed get node role arn from network status, error: %v", err)
			}
			nodeRoleServiceAccount := networkStatus.NodeRole
			// Authorize access to GCP storage
			err = attachGCPStorageAdminPermissionToSA(cr.gcpAPI, nodeRoleServiceAccount, bucketName)
			if err != nil {
				return nil, fmt.Errorf("failed attach storage admin permission to gcp service account, error: %v", err)
			}
		}
		providerSpec, err := cloudv1.GardenerProviderSpecForProvider(cr.provider)
		if err != nil {
			return nil, errors.Annotate(err, "get gardener provider spec")
		}
		tidbRestore = assembleGcpRestoreCRD(vars, providerSpec.GCP.ProjectID, annotations, labels)
		tidbRestore = mergeRestoreConfig(r, tidbRestore, brConfig, brConfigOutput)
	case consts.ProviderTypeAWS:
		if r.CheckBucketInfoIsExisted() {
			logger := log.FromContext(ctx).WithName("queryBucketFromSpec")
			logger.Info("restore from bucket info")
			tidbRestore = assembleAwsRestoreCRD(vars, annotations, labels)
			tidbRestore = mergeRestoreConfig(r, tidbRestore, brConfig, brConfigOutput)
			break
		}
		bucketRegion := cr.backup.Labels["cloud.pingcap.com/region"]
		bucketName := cr.backup.Status.BucketName
		vars.Region = bucketRegion

		// If bucket region is not same with cluster region. we need attch s3 policy to the cluster.
		isCrossRegion, network, err := cr.isCrossRegionRestore(ctx, vars.Region)
		if err != nil {
			return nil, fmt.Errorf("failed check the restore is cross region, error: %v", err)
		}
		// If cross region is true, we need authorize the restore cluster to be able to access the s3 bucket
		if isCrossRegion {
			// Get node role arn from network status
			networkStatus, err := cloudv1.GardenerNetworkStatusFromNetwork(network)
			if err != nil {
				return nil, fmt.Errorf("failed get node role arn from network status, error: %v", err)
			}
			nodeRoleArn := networkStatus.NodeRole

			if nodeRoleArn == "" {
				return nil, fmt.Errorf("node role arn can not be empty")
			}
			// Authorize access to s3 bucket
			if err = attachS3BucketPermissionToNodeRole(ctx, cr.awsCli, bucketName, nodeRoleArn); err != nil {
				return nil, fmt.Errorf("failed attach s3 bucket policy to node role, error: %v", err)
			}
			roleName := roleNameFromARN(nodeRoleArn)
			infraAPIRoleArn := os.Getenv(k8s.AWSRoleARNEnvKey)
			if infraAPIRoleArn != "" && !utils.InIntegrationTest() {
				// AssumeRole require role have trust policy, we will make infra-api can assume node role
				if err = updateNodeRoleAssumePolicy(ctx, cr.awsCli, roleName, []string{infraAPIRoleArn}); err != nil {
					return nil, fmt.Errorf("update node role assume policy, error: %v", err)
				}
				// infra-api will try assume node role and verify permission work fine
				if err := assumeRoleAndVerifyS3Permission(ctx, cr.awsCli, bucketName, bucketRegion, nodeRoleArn); err != nil {
					return nil, fmt.Errorf("failed assume role or verify s3 permission, error: %v", err)
				}
			}
		}
		tidbRestore = assembleAwsRestoreCRD(vars, annotations, labels)
		tidbRestore = mergeRestoreConfig(r, tidbRestore, brConfig, brConfigOutput)
	default:
		return nil, fmt.Errorf("unknown provider %s", providerType)
	}
	return tidbRestore, nil
}

func (cr *clusterRestore) buildTiDBRestoreForNewArch(r *cloudv1.Restore) (*v1alpha1.Restore, error) {
	return cr.restoreProvider.BuildRestore(r, cr.cluster)
}

// create node group for restore see more in https://pingcap.feishu.cn/wiki/wikcnTC7NOM6Gavdhzq7wmg5fMc
func (cr *clusterRestore) needUseNodeGroup(restore *cloudv1.Restore) bool {
	if cr.network == nil {
		return false
	}
	// the default memory limit and the value 100 Mi
	// https://github.com/tidbcloud/dataflow-service/blob/461fa309267d25f2d45fc6987dd7f3409db817d1/internal/core_ctx/infra/facade_imp/infra_imp.go#L44
	hasMemoryRequest := restore.Spec.ResourceRequirements != nil &&
		restore.Spec.ResourceRequirements.Requests.Memory() != nil
	needMoreMemory := (hasMemoryRequest &&
		restore.Spec.ResourceRequirements.Requests.Memory().Cmp(resource.MustParse("100Mi")) > 0)
	return needMoreMemory
}

func (cr *clusterRestore) ensurePriorityClass(ctx context.Context) error {
	priorityClass := &scheduling_v1.PriorityClass{
		ObjectMeta: v1.ObjectMeta{
			Name: ClusterCriticalPriorityClass,
		},
		Value:         ClusterCriticalPriorityValue,
		GlobalDefault: false,
		Description:   "Priority class for br",
	}
	_, err := controllerutil.CreateOrUpdate(ctx, cr.clusterCli, priorityClass, func() error { return nil })
	if err != nil {
		return errors.Annotate(err, "fail to create or update priority class")
	}
	return nil
}

func (cr *clusterRestore) isCrossRegionRestore(ctx context.Context, bucketRegion string) (bool, *cloudv1.Network, error) {
	n := new(cloudv1.Network)

	networkNamespaceName := cr.cluster.GetNetworkRef()
	err := cr.cli.Get(ctx, networkNamespaceName, n)
	if err != nil {
		return false, nil, err
	}
	if n.Spec.Region != bucketRegion {
		return true, n, nil
	}
	return false, nil, nil
}

func attachGCPStorageAdminPermissionToSA(gcpapi providergcp.GCPAPI, nodeRoleServiceAccount string, bucketName string) error {
	member := fmt.Sprintf("serviceAccount:%s", nodeRoleServiceAccount)

	err := gcpapi.AddBucketIAMMember(bucketName, storageAdminRoleName, member)
	if err != nil {
		return err
	}
	return nil
}

func roleNameFromARN(arn string) string {
	idx := strings.LastIndex(arn, "/")
	return arn[idx+1:]
}

func updateNodeRoleAssumePolicy(ctx context.Context, client provideraws.Client, nodeRoleName string, trusts []string) error {
	p := &awsapi.PolicyDocument{
		Version: "2012-10-17",
		Statement: []awsapi.PolicyStatement{
			{
				Effect: awsapi.StatementEffectAllow,
				Principal: &awsapi.PolicyPrincipal{
					Service: "ec2.amazonaws.com",
					AWS:     trusts,
				},
				Action: []string{"sts:AssumeRole"},
			},
		},
	}

	assumePolicy, err := p.ToString()
	if err != nil {
		return err
	}
	_, err = client.IAM().UpdateAssumeRolePolicy(
		ctx, &iam.UpdateAssumeRolePolicyInput{
			PolicyDocument: &assumePolicy,
			RoleName:       &nodeRoleName,
		},
	)
	return err
}

func assumeRole(ctx context.Context, client provideraws.Client, bucketName string, roleToAssumeArn string) (*ststypes.Credentials, error) {
	sessionName := fmt.Sprintf("cross-region-%s", bucketName)
	res, err := client.Sts().AssumeRole(ctx, &sts.AssumeRoleInput{
		RoleArn:         &roleToAssumeArn,
		RoleSessionName: &sessionName,
	})
	if err != nil {
		return nil, err
	}
	return res.Credentials, nil
}

// assumeRoleAndVerifyPermission assume role to get temp static aws credentials and use it list bucket object
func assumeRoleAndVerifyS3Permission(ctx context.Context, client provideraws.Client, bucketName string, bucketRegion, roleToAssumeArn string) error {
	creds, err := assumeRole(ctx, client, bucketName, roleToAssumeArn)
	if err != nil {
		fmt.Println(err)
		return err
	}

	cred := credentials.NewStaticCredentialsProvider(
		*creds.AccessKeyId, *creds.SecretAccessKey, *creds.SessionToken)

	cfg := aws.Config{
		Credentials: cred,
		Region:      bucketRegion,
	}
	c := s3api.NewFromConfig(cfg)
	_, err = c.ListObjectsV2(ctx, &s3api.ListObjectsV2Input{Bucket: &bucketName})
	if err != nil {
		return err
	}
	return nil
}

func attachS3BucketPermissionToNodeRole(
	ctx context.Context,
	client provideraws.Client,
	bucketName string,
	nodeRoleArn string) error {
	pd := provideraws.PolicyDocument{
		Version: provideraws.PolicyVersion,
		Statement: []provideraws.PolicyStatement{
			{
				Effect:   "Allow",
				Action:   []string{"s3:*"},
				Resource: []string{fmt.Sprintf("arn:aws:s3:::%s", bucketName), fmt.Sprintf("arn:aws:s3:::%s/*", bucketName)},
			},
		},
	}
	pdStr, err := pd.ToString()
	if err != nil {
		return err
	}
	role := nodeRoleArn[strings.LastIndex(nodeRoleArn, "/")+1:]
	policyInput := &iam.PutRolePolicyInput{
		PolicyDocument: aws.String(pdStr),
		PolicyName:     aws.String(bucketName),
		RoleName:       aws.String(role),
	}

	_, err = client.IAM().PutRolePolicy(ctx, policyInput)
	return err
}

func mergeRestoreConfig(cr *cloudv1.Restore, restore *v1alpha1.Restore, brConfig *v1alpha1.BRConfig, output *api.DescribeBRConfigOutput) *v1alpha1.Restore {
	tag := output.Version
	if output.Image.Tag != nil {
		tag = *output.Image.Tag
	}
	restore.Spec.ToolImage = fmt.Sprintf("%s:%s", output.Image.BaseImage, tag)

	restore.Spec.BR.LogLevel = brConfig.LogLevel
	restore.Spec.BR.StatusAddr = brConfig.StatusAddr
	restore.Spec.BR.TimeAgo = brConfig.TimeAgo
	restore = AddRestoreOption(cr, restore)
	if brConfig.Concurrency != nil {
		restore.Spec.BR.Concurrency = brConfig.Concurrency
	}
	if brConfig.RateLimit != nil {
		restore.Spec.BR.RateLimit = brConfig.RateLimit
	}
	if brConfig.Checksum != nil {
		restore.Spec.BR.Checksum = brConfig.Checksum
	}
	if brConfig.OnLine != nil {
		restore.Spec.BR.OnLine = brConfig.OnLine
	}
	return restore
}

func (cr *clusterRestore) Update(ctx context.Context, obj client.Object) (errResult error) {
	ctx, span := tracer.Start(ctx, "clusterRestore.Update")
	defer func() { tracerutil.EndSpan(span, errResult) }()

	r, ok := obj.(*cloudv1.Restore)
	if !ok {
		return errors.New(errNotRestore)
	}

	span.SetAttributes(
		attribute.String("restore_name", r.Name),
		attribute.String("restore_phase", string(r.Status.Phase)),
	)
	tracerutil.SetCommonAttributes(span, cr.cluster)

	var tidbRestore v1alpha1.Restore
	if err := cr.clusterCli.Get(ctx, types.NamespacedName{
		Name:      r.ObjectMeta.Name,
		Namespace: cr.k8sStatus.Namespace,
	}, &tidbRestore); err != nil {
		return err
	}

	r.Status.Phase = convertRestorePhase(tidbRestore.Status.Phase)
	syncRestoreProgresses(r, tidbRestore.Status.Progresses)

	if r.Status.Phase == cloudv1.RestoreComplete || r.Status.Phase == cloudv1.RestoreFailed {
		if err := cr.restoreProvider.PostRestoreFinished(ctx, r); err != nil {
			return err
		}
	}

	return cr.cli.Status().Update(ctx, r)
}

func (cr *clusterRestore) Delete(ctx context.Context, obj client.Object) (_ reconciler.ExternalDeletion, errResult error) {
	ctx, span := tracer.Start(ctx, "clusterRestore.Delete")
	defer func() { tracerutil.EndSpan(span, errResult) }()

	r, ok := obj.(*cloudv1.Restore)
	if !ok {
		return reconciler.ExternalDeletion{}, errors.New(errNotRestore)
	}

	span.SetAttributes(
		attribute.String("restore_name", r.Name),
		attribute.String("restore_phase", string(r.Status.Phase)),
	)
	tracerutil.SetCommonAttributes(span, cr.cluster)

	var tidbRestore v1alpha1.Restore
	if err := cr.clusterCli.Get(ctx, types.NamespacedName{
		Name:      r.ObjectMeta.Name,
		Namespace: cr.k8sStatus.Namespace,
	}, &tidbRestore); err != nil {
		if reconciler.IgnoreNotFound(err) != nil {
			return reconciler.ExternalDeletion{}, err
		}

		return reconciler.ExternalDeletion{ResourceDeleted: true}, nil
	}

	if !tidbRestore.GetDeletionTimestamp().IsZero() {
		return reconciler.ExternalDeletion{ResourceDeleted: false}, nil
	}

	deleteErr := cr.clusterCli.Delete(ctx, &tidbRestore)
	if cr.needUseNodeGroup(r) {
		ngName := restoreNodeGroupName(r, cr.inNewArch())
		log.FromContext(ctx).Info(fmt.Sprintf("clean node group for restore ng name: %s", ngName))
		if err := cr.restoreProvider.DeleteBRNode(ctx, cr.network.DeepCopy(), ngName); err != nil {
			log.FromContext(ctx).Error(
				err, fmt.Sprintf("clean node group for restore error, ng name: %s", ngName))
			return
		}
	}
	return reconciler.ExternalDeletion{ResourceDeleted: false}, deleteErr
}

type TidbRestoreVars struct {
	RestoreName           string
	RestoreNamespace      string
	ClusterName           string
	Region                string
	Bucket                string
	BucketPrefix          string
	BucketCloudProvider   string
	BucketLogBackupPrefix string
	RestoredTS            string
	IsVolumeBackup        bool
	PodSecurityContext    *corev1.PodSecurityContext
}

func assembleAwsRestoreCRD(vars TidbRestoreVars, annotations map[string]string, labels map[string]string) *v1alpha1.Restore {
	restore := &v1alpha1.Restore{
		ObjectMeta: v1.ObjectMeta{
			Name:        vars.RestoreName,
			Namespace:   vars.RestoreNamespace,
			Annotations: annotations,
			Labels:      labels,
		},
		Spec: v1alpha1.RestoreSpec{
			Mode:           v1alpha1.RestoreModeSnapshot,
			Type:           v1alpha1.BackupTypeFull,
			PitrRestoredTs: vars.RestoredTS,
			StorageProvider: v1alpha1.StorageProvider{
				S3: &v1alpha1.S3StorageProvider{
					Provider: v1alpha1.S3StorageProviderTypeAWS,
					Region:   vars.Region,
					Bucket:   vars.Bucket,
					Prefix:   vars.BucketPrefix,
				},
			},
			BR: &v1alpha1.BRConfig{
				Cluster:          vars.ClusterName,
				ClusterNamespace: vars.RestoreNamespace,
				SendCredToTikv:   pointer.Bool(false),
			},
			UseKMS:             false,
			PodSecurityContext: vars.PodSecurityContext,
		},
	}

	if len(vars.BucketLogBackupPrefix) > 0 {
		restore.Spec.Mode = v1alpha1.RestoreModePiTR
		// for PiTR, StorageProvider represents log backup, full backup is represented by PitrFullBackupStorageProvider
		restore.Spec.PitrFullBackupStorageProvider = restore.Spec.StorageProvider
		restore.Spec.StorageProvider = v1alpha1.StorageProvider{
			S3: &v1alpha1.S3StorageProvider{
				Provider: v1alpha1.S3StorageProviderTypeAWS,
				Region:   vars.Region,
				Bucket:   vars.Bucket,
				Prefix:   vars.BucketLogBackupPrefix,
			},
		}
	}

	return restore
}

func assembleGcpRestoreCRD(vars TidbRestoreVars, projectName string, annotations map[string]string, labels map[string]string) *v1alpha1.Restore {
	tidbRestore := &v1alpha1.Restore{
		ObjectMeta: v1.ObjectMeta{
			Name:        vars.RestoreName,
			Namespace:   vars.RestoreNamespace,
			Annotations: annotations,
			Labels:      labels,
		},
		Spec: v1alpha1.RestoreSpec{
			Mode:           v1alpha1.RestoreModeSnapshot,
			Type:           v1alpha1.BackupTypeFull,
			PitrRestoredTs: vars.RestoredTS,
			StorageProvider: v1alpha1.StorageProvider{
				Gcs: &v1alpha1.GcsStorageProvider{
					ProjectId:  projectName,
					Location:   vars.Region,
					Bucket:     vars.Bucket,
					Prefix:     vars.BucketPrefix,
					SecretName: consts.BRGCPSecretName,
				},
			},
			BR: &v1alpha1.BRConfig{
				Cluster:          vars.ClusterName,
				ClusterNamespace: vars.RestoreNamespace,
				SendCredToTikv:   pointer.Bool(true),
			},
			UseKMS:             false,
			PodSecurityContext: vars.PodSecurityContext,
		},
	}

	if len(vars.BucketLogBackupPrefix) > 0 {
		tidbRestore.Spec.Mode = v1alpha1.RestoreModePiTR
		// for PiTR, StorageProvider represents log backup, full backup is represented by PitrFullBackupStorageProvider
		tidbRestore.Spec.PitrFullBackupStorageProvider = tidbRestore.Spec.StorageProvider
		tidbRestore.Spec.StorageProvider = v1alpha1.StorageProvider{
			Gcs: &v1alpha1.GcsStorageProvider{
				ProjectId:  projectName,
				Location:   vars.Region,
				Bucket:     vars.Bucket,
				Prefix:     vars.BucketLogBackupPrefix,
				SecretName: consts.BRGCPSecretName,
			},
		}
	}

	return tidbRestore
}