package br

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	dataflowpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protowire"

	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/dataflow"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/workflow/common"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/infra/repoimpl/schema"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
	wf "github.com/tidbcloud/tidb-management-service/utilities/workflow"
)

const backupStorageAWSKMSARNFieldNumber = 5

// backupStorageAWSKMSARN reads the newly added field by protobuf number so TMS
// can consume it before the updated Dataflow Go module is published. Once the
// dependency is upgraded, the descriptor branch reads the generated field.
func backupStorageAWSKMSARN(info *dataflowpb.BackupStorageInfo) string {
	message := info.ProtoReflect()
	if field := message.Descriptor().Fields().ByNumber(backupStorageAWSKMSARNFieldNumber); field != nil {
		return message.Get(field).String()
	}

	unknown := message.GetUnknown()
	for len(unknown) > 0 {
		number, wireType, tagLength := protowire.ConsumeTag(unknown)
		if tagLength < 0 {
			return ""
		}
		unknown = unknown[tagLength:]
		if number == backupStorageAWSKMSARNFieldNumber && wireType == protowire.BytesType {
			value, valueLength := protowire.ConsumeBytes(unknown)
			if valueLength < 0 {
				return ""
			}
			return string(value)
		}
		fieldLength := protowire.ConsumeFieldValue(number, wireType, unknown)
		if fieldLength < 0 {
			return ""
		}
		unknown = unknown[fieldLength:]
	}
	return ""
}

type restoreFromDedicatedAdapter struct {
	restoreAdapterClassicBase

	dataflowCli dataflow.DataFlow
}

func (a *restoreFromDedicatedAdapter) PrepareRestore(
	ctx context.Context,
	req *globalngv1.RestoreClusterRequest,
) (submitRestore, bizerr.BizError) {
	targetClusterId, err := strconv.ParseUint(req.GetTargetClusterId(), 10, 64)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to parse target cluster id")
	}

	targetCluster, berr := a.clusterRepo.Get(ctx, targetClusterId, false)
	if berr != nil {
		return nil, berr.Chain("failed to get target cluster")
	}

	metadataJson, err := metadata.MarshalRequestMetaFromCtx(ctx)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to marshal metadata from context")
	}

	backupId, err := strconv.ParseUint(req.GetBackupId(), 10, 64)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to parse backup id")
	}
	backup, err := a.dataflowCli.GetBackup(ctx, &dataflowpb.GetBackupReq{
		BackupId: backupId,
	})
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to get backup")
	}
	restoreEntity, err := buildEntityForDedicatedRestore(req, targetCluster, backup)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to build classic restore entity")
	}

	return func(ctx context.Context) bizerr.BizError {
		err := a.transaction.Transaction(ctx, func(ctx context.Context) error {
			if berr := lockResourcePoolForRestore(
				ctx,
				a.resourcePoolRepo,
				a.clusterRepo,
				restoreEntity.TargetSharedPoolId,
				restoreEntity.TargetClusterId,
			); berr != nil {
				return berr
			}
			berr = a.restoreRepo.Create(ctx, restoreEntity)
			if berr != nil {
				return berr.Chain("failed to create classic restore record")
			}
			workflowErr := a.workflow.RunDagWithBizId(
				ctx,
				common.RestoreClusterFromDedicated,
				fmt.Sprintf("%d", targetCluster.GetOrganizationId()),
				"0", // nextgen cluster don't have project id
				req.TargetClusterId,
				map[string]string{
					common.KeyTargetClusterId: req.TargetClusterId,
					wf.KeyTraceId:             wf.GetTraceId(ctx),
					wf.KeyRequestMeta:         string(metadataJson),
				},
			)
			if workflowErr != nil {
				return bizerr.WrapErrorWithCode(
					workflowErr,
					bizerr.ErrCodeUnknown,
					"create %s workflow failed",
					common.RestoreClusterFromDedicated,
				)
			}
			a.writeRestoreStartedDisplayEvent(ctx, restoreEntity)
			return nil
		})
		if err != nil {
			return bizerr.WrapError(err, "failed to create classic restore record or run %s dag", common.RestoreClusterFromDedicated)
		}
		return nil
	}, nil
}

func (a *restoreFromDedicatedAdapter) StopRestore(ctx context.Context, req *globalngv1.StopRestoreClusterRequest) bizerr.BizError {
	targetClusterID, err := strconv.ParseUint(req.GetTargetClusterId(), 10, 64)
	if err != nil {
		return bizerr.WrapError(err, "failed to parse target cluster id")
	}
	restoreEntity, berr := a.restoreRepo.Get(ctx, targetClusterID)
	if berr != nil {
		return berr.Chain("failed to get restore before stop")
	}
	if berr := a.restoreAdapterClassicBase.StopRestore(ctx, req, common.RestoreClusterFromDedicated); berr != nil {
		return berr
	}
	if restoreEntity.SourceRoleArn == "" {
		return nil
	}
	err = a.dataflowCli.CleanupClassicRestoreAccess(ctx, &dataflow.CleanupClassicRestoreAccessRequest{
		BucketID: restoreEntity.SourceBucketId, RestoreID: req.GetTargetClusterId(),
	})
	if err != nil {
		tilog.Error(ctx, "failed to clean up dedicated restore access after cancellation", zap.Error(err),
			zap.Uint64("target_cluster_id", targetClusterID), zap.Uint64("bucket_id", restoreEntity.SourceBucketId))
		return bizerr.WrapError(err, "failed to clean up dedicated restore access after cancellation")
	}
	return nil
}

func buildEntityForDedicatedRestore(
	req *globalngv1.RestoreClusterRequest,
	targetCluster *entity.Cluster,
	backup *dataflowpb.GetBackupResp,
) (*entity.Restore, error) {
	stringToUint64 := func(s string) uint64 {
		id, _ := strconv.ParseUint(s, 10, 64)
		return id
	}

	var backupStorageInfo *dataflowpb.BackupStorageInfo
	for _, storageInfo := range backup.StorageInfos {
		if strings.EqualFold(storageInfo.Region, targetCluster.GetRegionId().String()) {
			backupStorageInfo = storageInfo
			break
		}
	}
	if backupStorageInfo == nil {
		return nil, fmt.Errorf("backup storage info not found in target region")
	}

	restoreEntity := &entity.Restore{
		Restore: &schema.Restore{
			State:                 "pending",
			TargetClusterId:       targetCluster.GetClusterId(),
			TenantId:              targetCluster.GetOrganizationId(),
			ExpectRestoreDataSize: backup.DataSizeBytes,
			RestoreMode:           schema.RestoreModeDedicatedSnapshot,
			BackupId:              stringToUint64(req.GetBackupId()),

			TargetSharedPoolId:  stringToUint64(targetCluster.GetSharedPoolId()),
			TargetRegionId:      targetCluster.GetRegionId().String(),
			TargetCloudProvider: targetCluster.GetCloudProviderId().String(),
			TargetServicePlan:   targetCluster.GetServicePlan().String(),
			TargetMinRcu:        targetCluster.GetMinRcu(),
			TargetMaxRcu:        targetCluster.GetMaxRcu(),
			SourceClusterId:     backup.ClusterId,
			SourceRegionId:      targetCluster.GetRegionId().String(),
			SourceCloudProvider: strings.ToLower(backup.CloudProvider),
			SourceServicePlan:   "Dedicated",
			SourceBackupStorage: fmt.Sprintf("s3://%s/%s", backupStorageInfo.BucketName, backupStorageInfo.Prefix),
			SourceBucketId:      backupStorageInfo.BucketId,
			SourceBackupKmsArn:  backupStorageAWSKMSARN(backupStorageInfo),
		},
	}

	return restoreEntity, nil
}