package br

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

	"github.com/aws/smithy-go/ptr"
	clusterngv1beta "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	eventdata "github.com/tidbcloud/eventbus-datastruct/event"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protojson"

	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/eventbus"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/br"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/cluster"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/region"
	resourcepoolrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/resourcepool"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/displayevent"
	"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/nextgen-global-service/pkg/config"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/dynamicconfig"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
	regionutil "github.com/tidbcloud/tidb-management-service/utilities/region"
	wf "github.com/tidbcloud/tidb-management-service/utilities/workflow"
)

type restoreFromPremiumAdapter struct {
	cfg              *config.Config
	workflow         wf.Workflow
	transaction      db.Transactioner
	backupRepo       br.BackupRepo
	restoreRepo      br.RestoreRepo
	clusterRepo      cluster.Repo
	regionRepo       region.Repo
	regionalCli      regionalserver.Client
	eventCommitter   eventbus.WALCommitter
	resourcePoolRepo resourcepoolrepo.Repo
}

func (a *restoreFromPremiumAdapter) PrepareRestore(
	ctx context.Context,
	req *globalngv1.RestoreClusterRequest,
) (submitRestore, bizerr.BizError) {
	sourceCluster, targetCluster, backup, pointInTime, berr := a.getRestoreContextInfo(ctx, req)
	if berr != nil {
		return nil, berr.Chain("failed to get restore context info")
	}

	if berr := validatePremiumSnapshotBackupRestorable(backup); berr != nil {
		return nil, berr
	}

	restoreEntity := buildLightWeightRestoreEntity(sourceCluster, targetCluster, backup, pointInTime)

	createLightweightRestoreRequest := buildLightWeightRestoreRequest(
		ctx,
		sourceCluster,
		targetCluster,
		backup,
		pointInTime,
		a.cfg.DynamicConfig.Loader,
	)
	tilog.Info(ctx, "createLightweightRestoreRequest will be sent to cluster service after workflow run",
		zap.Any("cluster_id", targetCluster.GetClusterId()),
		zap.Any("create_lightweight_restore_request", createLightweightRestoreRequest))

	requestJson, err := protojson.Marshal(createLightweightRestoreRequest)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to marshal create lightweight restore request")
	}

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

	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 restore record")
			}
			workflowErr := a.workflow.RunDagWithBizId(
				ctx,
				common.RestoreCluster,
				fmt.Sprintf("%d", targetCluster.GetOrganizationId()),
				"0", // nextgen cluster don't have project id
				fmt.Sprintf("%d", targetCluster.GetClusterId()),
				map[string]string{
					common.KeyRequest: string(requestJson),
					wf.KeyTraceId:     wf.GetTraceId(ctx),
					wf.KeyRequestMeta: string(metadataJson),
				},
			)
			if workflowErr != nil {
				return bizerr.WrapErrorWithCode(workflowErr, bizerr.ErrCodeUnknown, "create async workflow failed")
			}
			displayEv := newRestoreDisplayEvent(ctx, restoreEntity, eventdata.EventStatusStarted, "Restore started.")
			displayevent.Write(ctx, a.eventCommitter, displayEv)
			return nil
		})
		if err != nil {
			return bizerr.WrapError(err, "failed to create restore record or run restore dag")
		}
		return nil
	}, nil
}

func (a *restoreFromPremiumAdapter) GetRestoreRunningStatus(
	ctx context.Context,
	clusterId uint64,
	restore *entity.Restore,
) (*globalngv1.RestoreStatusResponse, bizerr.BizError) {
	targetCluster, berr := a.clusterRepo.Get(ctx, clusterId, false)
	if berr != nil {
		return nil, berr.Chain("failed to get cluster")
	}

	targetRegion, berr := a.regionRepo.Get(ctx, targetCluster.RegionId)
	if berr != nil {
		return nil, berr.Chain("failed to get region")
	}

	req := &clusterngv1beta.GetLightweightRestoreRequest{
		TargetSharedPoolId: targetCluster.SharedPoolId,
		RestoreId:          fmt.Sprintf("%d", restore.TargetClusterId),
	}

	restoreResp, berr := a.regionalCli.GetLightweightRestore(ctx, targetRegion, req)
	if berr != nil {
		return nil, berr.Chain("failed to get lightweight restore")
	}

	tilog.Info(
		ctx,
		"get lightweight restore in cluster service",
		zap.Uint64("cluster id", clusterId),
		zap.Any("restore", restoreResp),
	)

	switch restoreResp.Status {
	case clusterngv1beta.LightweightRestoreStatus_PROVISIONING:
		return &globalngv1.RestoreStatusResponse{
			State:   globalngv1.RestoreState_RESTORE_STATE_PENDING,
			Message: "Restore is pending for provisioning resources",
		}, nil
	case clusterngv1beta.LightweightRestoreStatus_DATA_RESTORING:
		progress := min(int32(restoreResp.DataRestoreProgressPercentage), 90)
		return &globalngv1.RestoreStatusResponse{
			State:    globalngv1.RestoreState_RESTORE_STATE_RUNNING,
			Message:  "Data restore is running",
			Progress: progress,
		}, nil
	case clusterngv1beta.LightweightRestoreStatus_DATA_RESTORING_ERROR:
		return &globalngv1.RestoreStatusResponse{
			State:   globalngv1.RestoreState_RESTORE_STATE_FAILED,
			Message: "Restore is failed during data restoring: " + restoreResp.Message,
		}, nil
	case clusterngv1beta.LightweightRestoreStatus_POST_RESTORE_ACTION:
		return &globalngv1.RestoreStatusResponse{
			State:    globalngv1.RestoreState_RESTORE_STATE_RUNNING,
			Message:  "Data restore is done, waiting for cluster service to execute post-restore operations",
			Progress: 95,
		}, nil
	case clusterngv1beta.LightweightRestoreStatus_SUCCEED:
		return &globalngv1.RestoreStatusResponse{
			State:    globalngv1.RestoreState_RESTORE_STATE_RUNNING,
			Message:  "Restore is done in regional control plane, waiting for global control plane to execute post-restore operations",
			Progress: 99,
		}, nil
	case clusterngv1beta.LightweightRestoreStatus_FAILED:
		return &globalngv1.RestoreStatusResponse{
			State:   globalngv1.RestoreState_RESTORE_STATE_FAILED,
			Message: "Restore is failed: " + restoreResp.Message,
		}, nil
	}
	return nil, bizerr.New(
		bizerr.GeneralInternalErrorCode,
		"%s",
		"unexpected restore status in cluster service: "+restoreResp.Status.String(),
	)
}

func (a *restoreFromPremiumAdapter) StopRestore(ctx context.Context, req *globalngv1.StopRestoreClusterRequest) bizerr.BizError {
	err := a.workflow.DeleteDagInstancesWithDagId(ctx, req.TargetClusterId, common.RestoreCluster)
	if err != nil {
		return bizerr.WrapError(err, "failed to delete restore dag instances")
	}
	tilog.Info(ctx, "delete restore dag instances success", zap.String("cluster_id", req.TargetClusterId))

	clusterId, err := strconv.ParseUint(req.TargetClusterId, 10, 64)
	if err != nil {
		return bizerr.WrapError(err, "failed to parse target cluster id")
	}
	targetCluster, berr := a.clusterRepo.Get(ctx, clusterId, true)
	if berr != nil {
		return berr.Chain("failed to get cluster")
	}

	targetRegion, berr := a.regionRepo.Get(ctx, targetCluster.RegionId)
	if berr != nil {
		return berr.Chain("failed to get region")
	}

	berr = a.regionalCli.DeleteLightweightRestore(ctx, targetRegion, &clusterngv1beta.DeleteLightweightRestoreRequest{
		TargetSharedPoolId: targetCluster.SharedPoolId,
		RestoreId:          req.TargetClusterId,
	})
	if berr != nil {
		if berr.ErrCode() == bizerr.ErrCodeRecordNotFound {
			tilog.Info(
				ctx,
				"lightweight restore is not found, stop restore done",
				zap.String("cluster_id", req.TargetClusterId),
			)
			return nil
		}
		return berr.Chain("failed to delete lightweight restore")
	}
	tilog.Info(
		ctx,
		"delete lightweight restore in cluster service success",
		zap.String("cluster_id", req.TargetClusterId),
	)
	return nil
}

func validatePremiumSnapshotBackupRestorable(backup *entity.Backup) bizerr.BizError {
	if backup == nil {
		return nil
	}
	if backup.State != schema.BackupStateSuccess {
		return bizerr.New(bizerr.ErrCodeParams, "backup %d is not restorable, state: %s", backup.ID, backup.State)
	}
	if backup.KernelBackupId == 0 || strings.TrimSpace(backup.KernelBackupName) == "" {
		return bizerr.New(bizerr.ErrCodeParams, "backup %d is missing kernel backup metadata", backup.ID)
	}
	return nil
}

func (a *restoreFromPremiumAdapter) getRestoreContextInfo(
	ctx context.Context,
	req *globalngv1.RestoreClusterRequest,
) (
	sourceCluster *entity.Cluster,
	targetCluster *entity.Cluster,
	backup *entity.Backup,
	pointInTime *time.Time,
	berr bizerr.BizError,
) {
	sourceClusterId, err := strconv.ParseUint(req.GetSourceClusterId(), 10, 64)
	if err != nil {
		return nil, nil, nil, nil, bizerr.WrapError(err, "failed to parse source cluster id")
	}
	sourceCluster, berr = a.clusterRepo.Get(ctx, sourceClusterId, true)
	if berr != nil {
		return nil, nil, nil, nil, berr.Chain("failed to get source cluster")
	}

	targetClusterId, err := strconv.ParseUint(req.GetTargetClusterId(), 10, 64)
	if err != nil {
		return nil, nil, nil, nil, bizerr.WrapError(err, "failed to parse target cluster id")
	}

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

	if req.GetRestoreMode() == globalngv1.RestoreMode_RESTORE_MODE_SNAPSHOT {
		backupId, err := strconv.ParseUint(req.GetBackupId(), 10, 64)
		if err != nil {
			return nil, nil, nil, nil, bizerr.WrapError(err, "failed to parse backup id")
		}
		backups, berr := a.backupRepo.UnscopedQuery(ctx, &br.QueryBackupsFilter{
			Ids:     []uint64{backupId},
			Deleted: ptr.Bool(false),
		}, nil)
		if berr != nil {
			return nil, nil, nil, nil, berr.Chain("failed to get backup")
		}
		if len(backups) == 0 {
			return nil, nil, nil, nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "backup not found")
		}
		backup = backups[0]
		if backup.TriggerType == schema.BackupTriggerTypeAuto &&
			!backup.ExpirationTime.IsZero() &&
			!time.Now().UTC().Before(backup.ExpirationTime) {
			return nil, nil, nil, nil, bizerr.New(bizerr.ErrCodeParams, "backup has expired")
		}
	} else {
		if req.GetPointInTime() == nil {
			return nil, nil, nil, nil, bizerr.New(bizerr.ErrCodeParams, "PointInTime is required for PiTR restore")
		}
		pointInTime = ptr.Time(req.GetPointInTime().AsTime())
		// Workround: kernel has a flaw can not handle ms correctly
		pointInTime = ptr.Time(pointInTime.Truncate(time.Second))

		backup = nil
	}

	return sourceCluster, targetCluster, backup, pointInTime, nil
}

func buildLightWeightRestoreEntity(
	sourceCluster *entity.Cluster,
	targetCluster *entity.Cluster,
	backup *entity.Backup,
	pointInTime *time.Time,
) *entity.Restore {
	stringToUint64 := func(s string) uint64 {
		id, _ := strconv.ParseUint(s, 10, 64)
		return id
	}
	restoreEntity := &entity.Restore{
		Restore: &schema.Restore{
			State:           "pending",
			PointInTime:     pointInTime,
			TargetClusterId: targetCluster.GetClusterId(),
			TenantId:        targetCluster.GetOrganizationId(),
			RestoreMode:     schema.RestoreModePitr,

			TargetSharedPoolId:  stringToUint64(targetCluster.GetSharedPoolId()),
			TargetRegionId:      targetCluster.GetRegionId().String(),
			TargetCloudProvider: targetCluster.GetCloudProviderId().String(),
			TargetServicePlan:   targetCluster.GetServicePlan().String(),
			TargetMinRcu:        targetCluster.GetMinRcu(),
			TargetMaxRcu:        targetCluster.GetMaxRcu(),
			SourceClusterId:     sourceCluster.GetClusterId(),
			SourceSharedPoolId:  stringToUint64(sourceCluster.GetSharedPoolId()),
			SourceRegionId:      sourceCluster.GetRegionId().String(),
			SourceCloudProvider: sourceCluster.GetCloudProviderId().String(),
			SourceServicePlan:   sourceCluster.GetServicePlan().String(),
		},
	}

	if backup != nil {
		restoreEntity.BackupId = backup.ID
		restoreEntity.ExpectRestoreDataSize = backup.BackupSizeBytes
		restoreEntity.RestoreMode = schema.RestoreModeSnapshot
	}

	return restoreEntity
}

func buildLightWeightRestoreRequest(
	ctx context.Context,
	sourceCluster *entity.Cluster,
	targetCluster *entity.Cluster,
	backup *entity.Backup, // if pointInTime is not nil, this is the nearest successful backup
	pointInTime *time.Time,
	loader dynamicconfig.Loader,
) *clusterngv1beta.CreateLightweightRestoreRequest {
	res := &clusterngv1beta.CreateLightweightRestoreRequest{
		RestoreId:          fmt.Sprintf("%d", targetCluster.GetClusterId()),
		SourceSharedPoolId: sourceCluster.GetSharedPoolId(),
		SourceClusterId:    fmt.Sprintf("%d", sourceCluster.GetClusterId()),
		TargetSharedPoolId: targetCluster.GetSharedPoolId(),
		TargetClusterId:    fmt.Sprintf("%d", targetCluster.GetClusterId()),
	}

	// PiTR restore
	if pointInTime != nil {
		res.PointInTime = pointInTime.Format("2006-01-02 15:04:05.000")
	} else {
		// Snapshot restore
		res.KernelBackupId = backup.KernelBackupId
		res.KernelBackupName = backup.KernelBackupName
		res.BackupDataSizeBytes = backup.BackupSizeBytes
	}

	// add optional resources config
	// if this is not configured, it will used the same resource as the tikv node in cluster service
	if loader != nil {
		values := loader.GetConfig().GetByDimensions(ctx, config.ConfigNextgenRestorePodResource, map[string]string{
			config.DimensionTenantId: fmt.Sprintf("%d", targetCluster.GetOrganizationId()),
			config.DimensionRegion: regionutil.ToUserFace(
				fmt.Sprintf("%s-%s", targetCluster.GetCloudProviderId().String(), targetCluster.GetRegionId().String()),
			),
			config.DimensionCluster: fmt.Sprintf("%d", targetCluster.GetClusterId()),
		})
		if len(values) != 0 {
			if cpu, ok := values[config.ValuesKeyCpu]; ok {
				res.ResourceCpu = cpu
			}
			if memory, ok := values[config.ValuesKeyMemory]; ok {
				res.ResourceMemory = memory
			}
			if storage, ok := values[config.ValuesKeyStorage]; ok {
				res.ResourceStorage = storage
			}
		}
	}

	return res
}