package impl

import (
	"context"
	"strconv"

	"github.com/AlekSi/pointer"
	workflowEntity "github.com/shiningrush/fastflow/pkg/entity"
	appconfig "github.com/tidbcloud/dataflow-service/app/config"
	"github.com/tidbcloud/dataflow-service/app/internal/feature_gate"
	"github.com/tidbcloud/dataflow-service/app/models/br"
	"github.com/tidbcloud/dataflow-service/app/models/domain"
	restore_model "github.com/tidbcloud/dataflow-service/app/models/restore"
	"github.com/tidbcloud/dataflow-service/app/repo"
	"github.com/tidbcloud/dataflow-service/app/repo/backup"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	restore_repo "github.com/tidbcloud/dataflow-service/app/repo/restore"
	"github.com/tidbcloud/dataflow-service/app/workflow"
	"github.com/tidbcloud/dataflow-service/internal/config"
	bizerr_wrapper "github.com/tidbcloud/dataflow-service/pkg/bizerr-wrapper"
	"github.com/tidbcloud/dataflow-service/pkg/blob"

	"github.com/tidbcloud/dataflow-service/pkg/grpc/grpcli"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/db"
	mgmt "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/zap"
	"google.golang.org/protobuf/encoding/protojson"
)

type PiTRestore struct {
	baseRestoreModel

	backupRepo    backup.BackupRepo
	logBackupRepo backup.LogBackupRepo

	workflowClient workflow.Workflow
	brCommon       *br.BRCommon

	getStorageFunc         blob.GetStorageFunc
	getCheckpointForCRPiTR func(
		ctx context.Context,
		logBackup entity.LogBackup,
		region string,
		getStorageFunc blob.GetStorageFunc,
	) (int64, bool, error)
}

func NewPiTRestore(
	dbProvider db.TransProvider,
	backupRepo backup.BackupRepo,
	restoreRepo restore_repo.RestoreRepo,
	bucketRepo repo.BucketRepo,
	logBackupRepo backup.LogBackupRepo,
	coreCli grpcli.CoreGrpcClient,
	mgmtSvcCli grpcli.MgmtSvcGrpcClient,
	dedicatedGlobalClient grpcli.DedicatedGlobalGrpcClient,
	configCenterCli grpcli.ConfigCenterGrpcClient,
	workflowClient workflow.Workflow,
	featureGate feature_gate.FeatureGate,
) restore_model.RestoreModel {
	brCommon := br.NewBRCommon(bucketRepo)
	s := &PiTRestore{
		baseRestoreModel: baseRestoreModel{
			dbProvider:            dbProvider,
			restoreRepo:           restoreRepo,
			bucketRepo:            bucketRepo,
			coreCli:               coreCli,
			mgmtSvcCli:            mgmtSvcCli,
			configCenterCli:       configCenterCli,
			dedicatedGlobalClient: dedicatedGlobalClient,
			featureGate:           featureGate,
		},
		backupRepo:     backupRepo,
		logBackupRepo:  logBackupRepo,
		workflowClient: workflowClient,
		brCommon:       brCommon,

		getStorageFunc:         blob.NewStorage,
		getCheckpointForCRPiTR: brCommon.GetCheckpointForCRPiTR,
	}
	s.statusMatrix = NewRestoreStatusMatrix()
	s.createRestoreHook = s
	s.deleteRestoreHook = s
	s.baseRestoreModel.createRestoreParamBuilders = []createRestoreParamBuilder{
		s.PiTBackupParamBuilder,
		s.bucketParamBuilder,
		s.restoreClusterVersionBuilder,
		s.restoreClusterParamBuilder,
	}
	s.baseRestoreModel.clusterSpecValidators = []clusterSpecValidator{
		s.ComponentTikvChecker,
		s.ComponentTiFlashChecker,
	}
	return s
}

type PiTRTimeRange struct {
	StartTS int64
	EndTS   int64

	Backup    entity.Backup
	LogBackup entity.LogBackup
}

func (p *PiTRestore) restoreClusterVersionBuilder(
	ctx context.Context, req *createRestoreParam,
) bizerr.BizError {
	logger := log.GetLogger(ctx)
	cluster, err := p.mgmtSvcCli.GetAnyDedicatedClusterBasicByID(ctx, req.createRestoreReq.ClusterID)
	if err != nil {
		logger.Error("get cluster error", zap.Error(err))
		return bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}
	dedicatedCluster := cluster.GetDedicatedCluster()
	if dedicatedCluster == nil {
		logger.Error("dedicated cluster info is nil")
		return bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}
	req.createClusterVersion = dedicatedCluster.Version
	return nil
}

func (p *PiTRestore) PiTBackupParamBuilder(
	ctx context.Context, req *createRestoreParam,
) bizerr.BizError {
	logger := log.GetLogger(ctx)
	timeRanges, err := p.genPiTRTimeRanges(ctx, req)
	if err != nil {
		return err
	}

	if len(timeRanges) == 0 {
		logger.Error("no valid time range found")
		return bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}

	srcClusterId := timeRanges[0].Backup.ClusterID
	srcCluster, err := p.mgmtSvcCli.GetAnyDedicatedClusterBasicByID(ctx, srcClusterId.Uint64())
	if err != nil {
		logger.Error("failed to get cluster info", zap.Error(err))
		return bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}
	for _, timeRange := range timeRanges {
		if p.checkIsMeetTimeRanges(ctx, timeRange, req, srcCluster) {
			req.backup = &timeRange.Backup
			req.logBackup = &timeRange.LogBackup
			return nil
		}
	}
	logger.Warn(appconfig.ErrMsgRestoreTsNotAllow)
	return bizerr_wrapper.NewInvalidInputError(appconfig.ErrMsgRestoreTsNotAllow)
}

func (p *PiTRestore) checkIsMeetTimeRanges(
	ctx context.Context, timeRange *entity.PiTRTimeRange, req *createRestoreParam,
	cluster *mgmt.Cluster,
) bool {
	logger := log.GetLogger(ctx)
	if timeRange.Backup.BucketSettingSnapshot == nil {
		logger.Warn("bucket setting snapshot is nil",
			zap.Uint64("backup_id", timeRange.Backup.ID.Uint64()),
			zap.Uint64("log_backup_id", timeRange.LogBackup.ID.Uint64()),
		)
		return false
	}
	backupRegion := req.createRestoreReq.BackupRegion
	effectiveBackupRegion := backupRegion
	if effectiveBackupRegion == "" {
		effectiveBackupRegion = req.GetClusterRegion()
	}
	endTs := timeRange.EndTS
	if effectiveBackupRegion != "" {
		// Resolve the effective region used for bucket selection:
		// - explicit BackupRegion from request, or
		// - cluster region when BackupRegion is omitted.
		bucket, found := timeRange.Backup.GetBucketByRegion(effectiveBackupRegion)
		if !found {
			if backupRegion != "" {
				logger.Info("backup has no bucket for the requested backup region, skipping time range",
					zap.String("backupRegion", backupRegion),
				)
				return false
			}
			logger.Info("backup has no bucket for cluster region, falling back to primary endTS",
				zap.String("clusterRegion", effectiveBackupRegion),
				zap.Uint64("backup_id", timeRange.Backup.ID.Uint64()),
			)
		} else if !bucket.Master {
			// Secondary bucket: use the per-region CRPiTR checkpoint.
			// On any failure (error or file not found), fall back to the
			// primary end timestamp so the region is still surfaced to callers.
			// refer to pit_backup.go:groupPiTRTimeRangesByRegionV2
			if override, ok := config.Global().BackupConfig.CRRFOverrides.GetMatched(
				cluster.GetOrganizationId(), cluster.GetProjectId(), cluster.GetClusterId(),
				cluster.GetCloudProvider(), cluster.GetRegion(),
			); ok && override.CalculatePiTRTimeRangeV2 {
				checkpoint, found, err := p.getCheckpointForCRPiTR(ctx, timeRange.LogBackup, effectiveBackupRegion, p.getStorageFunc)
				if err != nil {
					logger.Warn("get checkpoint for CR PiTR, falling back to primary endTS",
						zap.String("region", effectiveBackupRegion), zap.Error(err))
					endTs = timeRange.EndTS
				} else if found {
					endTs = checkpoint
				}
			}
		}
	}

	return timeRange.StartTS <= req.createRestoreReq.RestoredTs &&
		endTs >= req.createRestoreReq.RestoredTs
}

func (p *PiTRestore) genPiTRTimeRanges(
	ctx context.Context, req *createRestoreParam,
) ([]*entity.PiTRTimeRange, bizerr.BizError) {
	logger := log.GetLogger(ctx)
	logBackups, err := p.logBackupRepo.List(ctx,
		func(qs entity.LogBackupQuerySet) entity.LogBackupQuerySet {
			return qs.ClusterIDEq(entity.ID(req.createRestoreReq.ClusterID)).
				DeletedAtIsNull().
				OrderDescByCreatedAt()
		})
	if err != nil {
		logger.Error("failed to list log backups",
			zap.Any("clusterID", req.createRestoreReq.ClusterID), zap.Error(err))
		return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}

	backups, err := p.backupRepo.ListExternal(ctx, func(qs entity.BackupQuerySet) entity.BackupQuerySet {
		return qs.ClusterIDEq(entity.ID(req.createRestoreReq.ClusterID)).
			StatusEq(entity.BackupStatusSuccess).
			DeletedAtIsNull().
			OrderDescByCreatedAt()
	})
	if err != nil {
		logger.Error("failed to list backups",
			zap.Any("clusterID", req.createRestoreReq.ClusterID), zap.Error(err))
		return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}

	var timeRanges []*entity.PiTRTimeRange
	for _, logBackup := range logBackups {
		ranges := logBackup.GenPiTRTimeRanges(ctx, backups)
		if len(ranges) == 0 {
			continue
		}
		timeRanges = append(timeRanges, ranges...)
	}
	return timeRanges, nil
}

//nolint:nestif
func (p *PiTRestore) makeCreateRestoreEntity(
	ctx context.Context, createParam *createRestoreParam,
) (*entity.Restore, bizerr.BizError) {
	// Layer 1: best-effort PiTR serialization check.
	// Reject the request if another PiTR restore is already PENDING or RUNNING for the same log backup.
	// This is best-effort: a TOCTOU race between two simultaneous creates is caught by Layer 2 (run_compact_backup).
	logBackup := createParam.logBackup
	override, _ := config.Global().BackupConfig.RestoreOverrides.GetMatched(
		logBackup.TenantID.Uint64(), logBackup.ProjectID.Uint64(), logBackup.ClusterID.Uint64(), "", "",
	)
	// Note: provider/region are omitted here because the target cluster may not exist yet (restore-to-new path).
	// Override entries that gate on provider/region won't match at Layer 1; they will still apply at Layer 2.
	if !override.AllowConcurrentPiTR {
		count, countErr := p.restoreRepo.Count(ctx, func(qs entity.RestoreQuerySet) entity.RestoreQuerySet {
			return qs.LogBackupIDEq(logBackup.ID).
				ModeEq(entity.RestoreModePiTR).
				StatusIn(entity.RestoreStatusPENDING, entity.RestoreStatusRUNNING)
		})
		if countErr != nil {
			log.GetLogger(ctx).Error("failed to count running pitr restores", zap.Error(countErr))
			return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
		}
		if count > 0 {
			log.GetLogger(ctx).Info("PiTR restore blocked by serialization constraint: another restore is already PENDING or RUNNING for the same log backup",
				zap.Uint64("log_backup_id", logBackup.ID.Uint64()))
			return nil, bizerr_wrapper.NewConflictErr(appconfig.ErrMsgPiTRConcurrentRestoreNotAllowed)
		}
	}

	var (
		region                  string
		createClusterReqV2Bytes []byte
		err                     error
	)
	if createParam.createClusterReqV2 != nil {
		region = createParam.createClusterReqV2.Region
		createClusterReqV2Bytes, err = protojson.Marshal(createParam.createClusterReqV2)
		if err != nil {
			log.GetLogger(ctx).Error("failed to marshal create cluster request v2", zap.Error(err))
			return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
		}
	} else if createParam.createRestoreReq.ClusterConfig.ClusterID != 0 {
		region = createParam.createRestoreReq.ClusterConfig.Region
	} else {
		return nil, bizerr_wrapper.NewInternalErr(
			"createClusterReq or createClusterReqV2 is required when cluster id is not specified")
	}

	req := createParam.createRestoreReq
	return &entity.Restore{
		// just for pitr
		LogBackupID: createParam.logBackup.ID,
		RestoredTS:  req.RestoredTs,

		Name:          req.ClusterConfig.Name,
		UserID:        entity.ID(req.UserID),
		Mode:          entity.RestoreModePiTR,
		Region:        &region,
		BackupID:      createParam.backup.ID,
		Status:        entity.RestoreStatusRUNNING,
		RequestArgsV2: string(createClusterReqV2Bytes),
		KernelMetrics: &entity.RestoreKernelMetrics{
			Status: pointer.To(entity.RestoreKernelStatusPending),
		},
		TenantID:  entity.ID(req.TenantID),
		ProjectID: entity.ID(req.ProjectID),
		BucketID:  createParam.bucketID,
		ExtraInfo: &entity.RestoreExtra{
			RestoreCRNameVersion: "v1",
		},
	}, nil
}

func (p *PiTRestore) runCreateRestoreDagIns(
	ctx context.Context, createParam *createRestoreParam, restore *entity.Restore,
) (*workflowEntity.DagInstance, bizerr.BizError) {
	logger := log.GetLogger(ctx)
	// start wfs
	dagVars, err := p.buildRestorePitBackupDagVars(ctx, restore, createParam)
	if err != nil {
		logger.Warn("failed to build dag vars", zap.Error(err))
		return nil, err
	}
	dagIns, dagErr := p.workflowClient.RunDagWithTags(
		workflow.DagPitRestore, dagVars,
		map[string]string{
			workflow.BusinessTypeKey:   string(workflow.BusinessTypeRestore),
			workflow.BusinessActionKey: string(workflow.BusinessActionRestorePit),
			workflow.BusinessID:        strconv.FormatUint(uint64(restore.ID), 10),
			workflow.ClusterID:         strconv.FormatUint(uint64(createParam.backup.ClusterID), 10),
		})
	if dagErr != nil {
		logger.Error("failed to run dag ins", zap.Error(dagErr))
		return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}
	logger.Info("create pit restore dag ins started",
		zap.Any("restore id", restore.ID), zap.Any("dag id", dagIns.ID))
	return dagIns, nil
}

func (p *PiTRestore) buildRestorePitBackupDagVars(
	ctx context.Context, restore *entity.Restore, createRestoreParam *createRestoreParam,
) (map[string]string, bizerr.BizError) {
	runVar := map[string]string{}
	runVar[workflow.KeyRestoreID] = strconv.FormatUint(uint64(restore.ID), 10)

	if createRestoreParam.createRestoreReq.ClusterConfig.ClusterID == 0 {
		runVar[workflow.KeyRestoreToClusterType] = string(domain.RestoreToClusterTypeNew)
	} else {
		if createRestoreParam.createRestoreReq.ClusterConfig.Cloned {
			runVar[workflow.KeyRestoreToClusterType] = string(domain.RestoreToClusterTypeClone)
		} else {
			runVar[workflow.KeyRestoreToClusterType] = string(domain.RestoreToClusterTypeExist)
		}
	}

	if createRestoreParam.createRestoreReq.BRVersion != "" {
		runVar[workflow.BRVersion] = createRestoreParam.createRestoreReq.BRVersion
	}
	if createRestoreParam.createRestoreReq.BRImageTag != "" {
		runVar[workflow.BRImageTag] = createRestoreParam.createRestoreReq.BRImageTag
	}
	if createRestoreParam.createRestoreReq.BRCustomOptions != "" {
		runVar[workflow.BRCustomOptions] = createRestoreParam.createRestoreReq.BRCustomOptions
	}
	return runVar, nil
}

func (p *PiTRestore) buildRecyclePitSnapshotBackupDagVars(
	ctx context.Context, restore *entity.Restore,
) (map[string]string, bizerr.BizError) {
	runVar := map[string]string{}
	runVar[workflow.KeyRestoreID] = strconv.FormatUint(uint64(restore.ID), 10)
	return runVar, nil
}

func (p *PiTRestore) runDeleteRestoreDagIns(
	ctx context.Context, restore *entity.Restore,
) (*workflowEntity.DagInstance, bizerr.BizError) {
	logger := log.GetLogger(ctx)
	// start wfs
	dagVars, err := p.buildRecyclePitSnapshotBackupDagVars(ctx, restore)
	if err != nil {
		logger.Error("failed to build dag vars", zap.Error(err))
		return nil, err
	}
	dagIns, dagErr := p.workflowClient.RunDagWithTags(
		workflow.DagGcPitRestore, dagVars,
		map[string]string{
			workflow.BusinessTypeKey:   string(workflow.BusinessTypeRestore),
			workflow.BusinessActionKey: string(workflow.BusinessActionGcRestorePit),
			workflow.BusinessID:        strconv.FormatUint(uint64(restore.ID), 10),
			workflow.ClusterID:         strconv.FormatUint(uint64(restore.ClusterID), 10),
		})
	if dagErr != nil {
		logger.Error("failed to run dag ins", zap.Error(dagErr))
		return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
	}
	logger.Info("delete pit restore dag ins started",
		zap.Any("restore id", restore.ID), zap.Any("dag id", dagIns.ID))
	return dagIns, nil
}