package premium_dm

import (
	"context"
	"fmt"
	"strconv"

	dmconfig "github.com/pingcap/tiflow/dm/config"
	"github.com/pingcap/tiflow/dm/pkg/checker"
	"github.com/tidbcloud/dataflow-service-ng/app/models/domain"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	apputils "github.com/tidbcloud/dataflow-service-ng/app/utils"
	"github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	"github.com/tidbcloud/dataflow-service-ng/internal/config"
	bizerr_wrapper "github.com/tidbcloud/dataflow-service-ng/pkg/bizerr-wrapper"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"github.com/tidbcloud/pkgv2/bizerr"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/zap"
)

func hasAnySourceBinlogCheckpoint(sources []entity.PremiumDMSource) bool {
	for i := range sources {
		if sources[i].BinlogName != "" || sources[i].BinlogGTID != "" {
			return true
		}
	}
	return false
}

func hasAliCloudRDSSource(sources []entity.PremiumDMSource) bool {
	for i := range sources {
		if sources[i].SourceType == entity.PremiumDBTypeAliCloudRDSMySQL {
			return true
		}
	}
	return false
}

func isFullInstanceMigrateForAllSources(sources []entity.PremiumDMSource) bool {
	for i := range sources {
		if len(sources[i].MigrationRules) > 0 {
			return false
		}
	}
	return true
}

// sourcesForPremiumPrecheckItem returns the sources and passwords relevant to
// one precheck item. The primary-key checker is an AliCloud RDS requirement,
// so it must not run against generic MySQL or Azure sources in a mixed task.
//
// Source passwords are positional in PrecheckDMJobReq. Keep filtering them in
// the same loop as sources so every generated subtask config uses its matching
// password.
func sourcesForPremiumPrecheckItem(
	sources []entity.PremiumDMSource,
	sourcePasswords []string,
	checkItem string,
) ([]entity.PremiumDMSource, []string, error) {
	if len(sources) != len(sourcePasswords) {
		return nil, nil, fmt.Errorf("source and password count mismatch: %d sources, %d passwords", len(sources), len(sourcePasswords))
	}
	if checkItem != PrimaryKeyChecking {
		return sources, sourcePasswords, nil
	}

	aliCloudRDSSources := make([]entity.PremiumDMSource, 0, len(sources))
	aliCloudRDSPasswords := make([]string, 0, len(sources))
	for i := range sources {
		if sources[i].SourceType != entity.PremiumDBTypeAliCloudRDSMySQL {
			continue
		}
		aliCloudRDSSources = append(aliCloudRDSSources, sources[i])
		aliCloudRDSPasswords = append(aliCloudRDSPasswords, sourcePasswords[i])
	}
	if len(aliCloudRDSSources) == 0 {
		return nil, nil, fmt.Errorf("no AliCloud RDS source for primary-key precheck")
	}
	return aliCloudRDSSources, aliCloudRDSPasswords, nil
}

func (w *PremiumDMPrecheckWorker) resolveSourcePasswords(sources []entity.PremiumDMSource) ([]string, bizerr.BizError) {
	if len(sources) == 0 {
		return nil, bizerr_wrapper.NewInternalErr("no source found for precheck")
	}

	passwords := make([]string, 0, len(sources))
	for i := range sources {
		current, err := decryptPersistedPassword(w.taskWorker.Encryptor, sources[i].ConnProfile.Password)
		if err != nil {
			return nil, err
		}
		passwords = append(passwords, current)
	}
	return passwords, nil
}

// PrecheckItemRunner executes a single precheck item and persists its result.
// The returned message is set only for ng-agent checker initialization errors.
// This is a local wire contract between the item runner and RunPrecheck.
type PrecheckItemRunner func(
	ctx context.Context,
	precheckID string,
	item *entity.PremiumDMPrecheckItem,
) (*premiumDMPrecheckUserMessage, bizerr.BizError)

type premiumPrecheckRunState struct {
	failedCount                  int
	firstErrMsg                  string
	userFacingInitFailureMessage *premiumDMPrecheckUserMessage
}

func premiumModeToEntity(mode v1beta2.PremiumDMTaskMode) entity.PremiumDMTaskMode {
	switch mode {
	case v1beta2.PremiumDMTaskMode_ALL:
		return entity.PremiumDMTaskModeAll
	case v1beta2.PremiumDMTaskMode_INCREMENTAL:
		return entity.PremiumDMTaskModeIncremental
	default:
		return entity.PremiumDMTaskModeFull
	}
}

func premiumImportModeToEntity(mode v1beta2.PremiumDMTaskImportMode_ENUM) entity.PremiumDMTaskImportMode {
	switch mode {
	case v1beta2.PremiumDMTaskImportMode_IMPORTMODE_PHYSICAL:
		return entity.PremiumDMTaskImportModePhysical
	default:
		return entity.PremiumDMTaskImportModeLogical
	}
}

func (w *PremiumDMPrecheckWorker) failPrecheck(ctx context.Context, precheckID string, err error) {
	if err == nil {
		return
	}
	_ = w.dao.UpdatePrecheck(ctx, precheckID, map[string]any{
		"stage":   entity.PremiumDMPrecheckStageFailed,
		"err_msg": err.Error(),
	})
}

//nolint:funlen // precheck orchestration has many sequential steps
func (w *PremiumDMPrecheckWorker) runSinglePrecheck(ctx context.Context, precheck *entity.PremiumDMPrecheck) error {
	if precheck == nil {
		return bizerr_wrapper.NewInternalErr("precheck is nil")
	}
	clusterID := precheck.TiDBID
	precheckID := precheck.ID
	logger := GetWorkerLoggerFromContext(ctx).With(
		zap.String("cluster_id", clusterID),
		zap.String("precheck_id", precheckID),
	)

	precheck, bizErr := w.dao.GetPrecheck(ctx, clusterID, precheckID)
	if bizErr != nil {
		logger.Error("failed to load precheck", zap.Error(bizErr))
		return bizErr
	}

	createReq, err := UnmarshalPrecheckRequest(precheck.Request)
	if err != nil {
		logger.Error("failed to unmarshal precheck request", zap.Error(err))
		w.failPrecheck(ctx, precheckID, err)
		return err
	}

	clusterIDUint, err := strconv.ParseUint(clusterID, 10, 64)
	if err != nil {
		logger.Error("invalid cluster id", zap.Error(err))
		w.failPrecheck(ctx, precheckID, err)
		return err
	}

	listResp, listErr := w.taskWorker.MgmtClient.ListClusters(
		ctx,
		&tidbmgmtservicev1.ListClustersRequest{
			Filter: &tidbmgmtservicev1.ClusterFilter{
				ClusterIds: []uint64{clusterIDUint},
			},
		},
	)
	if listErr != nil {
		logger.Error("failed to get cluster", zap.Error(listErr))
		w.failPrecheck(ctx, precheckID, listErr)
		return listErr
	}
	if len(listResp.Clusters) == 0 {
		err = bizerr_wrapper.NewInternalErr("cluster not found")
		logger.Error("failed to get cluster", zap.Error(err))
		w.failPrecheck(ctx, precheckID, err)
		return err
	}
	cluster := listResp.Clusters[0]
	agentClient, err := w.taskWorker.newNGRegionalClientByClusterID(ctx, clusterIDUint)
	if err != nil {
		logger.Error("failed to get ng regional agent client", zap.Error(err))
		w.failPrecheck(ctx, precheckID, err)
		return err
	}

	// Batch-load all PRIVATELINK endpoints for this cluster's shared pool.
	clusterWrapper := apputils.NewCluster(cluster)
	id2endpoint, err := BuildEndpointMapForCluster(ctx, w.taskWorker.BizSinkEndpointRepo, clusterWrapper)
	if err != nil {
		logger.Error("failed to build endpoint map for cluster", zap.Error(err))
		w.failPrecheck(ctx, precheckID, err)
		return err
	}

	sources := ConvertRequestSourcesToEntities(
		clusterID,
		precheck.TenantID,
		createReq.GetSources(),
		id2endpoint,
	)
	if len(sources) == 0 {
		err = bizerr_wrapper.NewInternalErr("no source found for precheck")
		logger.Error("failed to initialize precheck source", zap.Error(err))
		w.failPrecheck(ctx, precheckID, err)
		return err
	}

	taskMode := premiumModeToEntity(createReq.GetMode())
	importMode := premiumImportModeToEntity(createReq.GetImportMode())

	items := GetPremiumPrecheckExecutionItems(
		taskMode,
		hasAnySourceBinlogCheckpoint(sources),
		importMode,
		hasAliCloudRDSSource(sources),
		precheckID,
		config.Global().PremiumDMConfig.PrecheckSkipConfigs,
		parsePrecheckUintID(precheck.TenantID),
		parsePrecheckUintID(precheck.ProjectID),
		clusterIDUint,
	)

	sourcePasswords, bizErr := w.resolveSourcePasswords(sources)
	if bizErr != nil {
		logger.Error("failed to resolve source passwords", zap.Error(bizErr))
		w.failPrecheck(ctx, precheckID, bizErr)
		return bizErr
	}
	targetPassword, bizErr := decryptPersistedPassword(w.taskWorker.Encryptor, createReq.GetTarget().GetPassword())
	if bizErr != nil {
		logger.Error("failed to decrypt target password", zap.Error(bizErr))
		w.failPrecheck(ctx, precheckID, bizErr)
		return bizErr
	}

	w.RunPrecheck(ctx, clusterID, precheckID, items,
		func(
			itemCtx context.Context,
			runPrecheckID string,
			item *entity.PremiumDMPrecheckItem,
		) (*premiumDMPrecheckUserMessage, bizerr.BizError) {
			if item == nil {
				return nil, bizerr_wrapper.NewInternalErr("precheck item is nil")
			}

			itemSources, itemSourcePasswords, filterErr := sourcesForPremiumPrecheckItem(
				sources,
				sourcePasswords,
				item.Description,
			)
			if filterErr != nil {
				logger.Error("failed to select precheck sources", zap.Error(filterErr))
				return nil, bizerr_wrapper.NewInternalErr("%s", filterErr.Error())
			}

			subTaskCfgs := make([]*dmconfig.SubTaskConfig, 0, len(itemSources))
			for i := range itemSources {
				premiumSource := &itemSources[i]
				baseCfg, cfgErr := BuildPremiumPrecheckBaseConfig(cluster, createReq, premiumSource)
				if cfgErr != nil {
					logger.Error("failed to build base precheck config", zap.Error(cfgErr), zap.String("source_id", premiumSource.ID))
					return nil, cfgErr
				}
				subTaskCfg := ApplyPremiumSourcePrecheckConfig(baseCfg, premiumSource, taskMode, importMode)
				subTaskCfgs = append(subTaskCfgs, subTaskCfg)
			}

			agentReq := &agentpb.PrecheckDMJobReq{
				ReqId:               runPrecheckID,
				FullInstanceMigrate: isFullInstanceMigrateForAllSources(itemSources),
				ClusterPassword:     targetPassword,
				SourcePasswords:     itemSourcePasswords,
				ClusterId:           clusterID,
			}

			itemLogger := logger.With(
				zap.String("item_id", item.ID),
				zap.String("check_item", item.Description),
			)
			return w.doPremiumPrecheckDMJob(
				itemCtx,
				agentClient,
				runPrecheckID,
				itemLogger,
				subTaskCfgs,
				item.Description,
				item,
				agentReq,
			)
		})

	// Clean up sink endpoint usages for PRIVATELINK sources after precheck completes
	w.cleanupPrecheckEndpointUsages(ctx, logger, precheckID, sources)

	return nil
}

// cleanupPrecheckEndpointUsages deletes BizSinkEndpointUsage records for
// PRIVATELINK sources after a precheck completes.
func (w *PremiumDMPrecheckWorker) cleanupPrecheckEndpointUsages(
	ctx context.Context,
	logger *zap.Logger,
	precheckID string,
	sources []entity.PremiumDMSource,
) {
	for i := range sources {
		s := &sources[i]
		if s.ConnProfile.Type == entity.PremiumConnTypePrivateLink && s.ConnProfile.SinkEndpointID != 0 {
			delErr := w.taskWorker.SinkEndpointModel.DeleteSinkEndpointUsage(ctx, &domain.DeleteSinkEndpointUsageReq{
				SinkEndpointID: s.ConnProfile.SinkEndpointID,
				OwnerID:        precheckID,
				UsageType:      entity.BizSinkEndpointUsageTypePremiumDMPrecheck,
			})
			if delErr != nil {
				logger.Warn("failed to delete sink endpoint usage for precheck source",
					zap.Error(delErr), zap.String("source_id", s.ID))
			}
		}
	}
}

func (w *PremiumDMPrecheckWorker) RunPrecheck(
	ctx context.Context,
	clusterID string,
	precheckID string,
	items []*entity.PremiumDMPrecheckItem,
	runner PrecheckItemRunner,
) {
	logger := GetWorkerLoggerFromContext(ctx).With(zap.String("precheck_id", precheckID))

	if !w.beginPrecheckRun(ctx, clusterID, precheckID, logger, len(items)) {
		return
	}
	state := &premiumPrecheckRunState{}

	defer func() {
		w.finalizePrecheckRun(ctx, clusterID, precheckID, logger, state, recover())
	}()

	if runner == nil {
		logger.Error("precheck runner is nil")
		state.failedCount++
		state.firstErrMsg = "runner callback is nil"
		return
	}

	for _, item := range items {
		stop, shouldReturn := w.processPrecheckItem(ctx, clusterID, precheckID, item, runner, logger, state)
		if shouldReturn {
			return
		}
		if stop {
			break
		}
	}
}

func (w *PremiumDMPrecheckWorker) beginPrecheckRun(
	ctx context.Context,
	clusterID string,
	precheckID string,
	logger *zap.Logger,
	totalItems int,
) bool {
	precheck, err := w.dao.GetPrecheck(ctx, clusterID, precheckID)
	if err != nil {
		logger.Error("failed to get precheck", zap.Error(err))
		return false
	}
	if precheck.Stage == string(entity.PremiumDMPrecheckStageCanceled) {
		logger.Info("precheck is canceled, skipping")
		return false
	}
	if bizErr := w.dao.UpdatePrecheck(ctx, precheckID, map[string]any{
		"stage": entity.PremiumDMPrecheckStageRunning,
	}); bizErr != nil {
		logger.Error("failed to transition precheck to running", zap.Error(bizErr))
		return false
	}
	logger.Info("precheck started", zap.Int("total_items", totalItems))
	return true
}

func (w *PremiumDMPrecheckWorker) finalizePrecheckRun(
	ctx context.Context,
	clusterID string,
	precheckID string,
	logger *zap.Logger,
	state *premiumPrecheckRunState,
	panicValue any,
) {
	if panicValue != nil {
		logger.Error("precheck orchestration panicked", zap.Any("panic", panicValue))
		state.failedCount++
		if state.firstErrMsg == "" {
			state.firstErrMsg = "orchestration panic"
		}
	}

	finalPrecheck, reloadErr := w.dao.GetPrecheck(ctx, clusterID, precheckID)
	if reloadErr != nil {
		logger.Error("failed to reload precheck for final update", zap.Error(reloadErr))
		return
	}
	if finalPrecheck.Stage == string(entity.PremiumDMPrecheckStageCanceled) {
		logger.Info("precheck was canceled, skipping final status update")
		return
	}

	finalStage, finalErrMsg := resolveFinalPrecheckStatus(finalPrecheck, state.failedCount, state.firstErrMsg)
	updates := map[string]any{"stage": finalStage}
	if finalErrMsg != "" {
		updates["err_msg"] = finalErrMsg
	}
	if bizErr := w.dao.UpdatePrecheck(ctx, precheckID, updates); bizErr != nil {
		logger.Error("failed to update precheck final status", zap.Error(bizErr))
		return
	}

	logger.Info("precheck completed",
		zap.String("final_stage", string(finalStage)),
		zap.Int("failed_items", state.failedCount),
		zap.String("error_message", finalErrMsg),
	)
}

func (w *PremiumDMPrecheckWorker) processPrecheckItem(
	ctx context.Context,
	clusterID string,
	precheckID string,
	item *entity.PremiumDMPrecheckItem,
	runner PrecheckItemRunner,
	logger *zap.Logger,
	state *premiumPrecheckRunState,
) (stop bool, shouldReturn bool) {
	currentPrecheck, reloadErr := w.dao.GetPrecheck(ctx, clusterID, precheckID)
	if reloadErr != nil {
		logger.Error("failed to reload precheck before item processing", zap.Error(reloadErr))
		state.failedCount++
		if state.firstErrMsg == "" {
			state.firstErrMsg = reloadErr.Error()
		}
		return false, false
	}
	if currentPrecheck.Stage == string(entity.PremiumDMPrecheckStageCanceled) {
		logger.Info("precheck was canceled during item processing, stopping")
		return true, false
	}

	itemLogger := logger.With(zap.String("item_id", item.ID), zap.String("item_type", item.Description))
	if state.userFacingInitFailureMessage != nil {
		return false, w.persistSkippedPrecheckItem(
			ctx,
			precheckID,
			item,
			itemLogger,
			state,
		)
	}

	initFailureMessage, itemErr := runner(ctx, precheckID, item)
	if initFailureMessage != nil {
		state.userFacingInitFailureMessage = initFailureMessage
	}
	if itemErr == nil {
		return false, false
	}

	itemLogger.Warn("precheck item failed", zap.Error(itemErr))
	state.failedCount++
	if state.firstErrMsg == "" {
		state.firstErrMsg = itemErr.Error()
	}
	return false, true
}

func (w *PremiumDMPrecheckWorker) persistSkippedPrecheckItem(
	ctx context.Context,
	precheckID string,
	item *entity.PremiumDMPrecheckItem,
	itemLogger *zap.Logger,
	state *premiumPrecheckRunState,
) bool {
	itemErr := w.persistPremiumPrecheckUserFailure(
		ctx,
		precheckID,
		item,
		state.userFacingInitFailureMessage,
		itemLogger,
	)
	if itemErr == nil {
		return false
	}
	itemLogger.Warn("failed to persist skipped precheck item", zap.Error(itemErr))
	state.failedCount++
	if state.firstErrMsg == "" {
		state.firstErrMsg = itemErr.Error()
	}
	return true
}

func resolveFinalPrecheckStatus(
	precheck *entity.PremiumDMPrecheck,
	failedCount int,
	firstErrMsg string,
) (entity.PremiumDMPrecheckStage, string) {
	if failedCount > 0 {
		return entity.PremiumDMPrecheckStageFailed, firstErrMsg
	}
	if precheck == nil {
		return entity.PremiumDMPrecheckStageFinished, ""
	}
	for i := range precheck.DMPrecheckItems {
		item := &precheck.DMPrecheckItems[i]
		if item.Stage == string(entity.PremiumDMPrecheckStageFailed) {
			return entity.PremiumDMPrecheckStageFailed, item.ErrMsg
		}
		if summary, ok := item.SummarizeResults(); ok && summary.State == checker.StateFailure {
			return entity.PremiumDMPrecheckStageFailed, ""
		}
	}
	return entity.PremiumDMPrecheckStageFinished, ""
}

func parsePrecheckUintID(id string) uint64 {
	parsed, err := strconv.ParseUint(id, 10, 64)
	if err != nil {
		return 0
	}
	return parsed
}