package hook

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

	"github.com/shiningrush/fastflow/pkg/entity"
	"github.com/shiningrush/fastflow/pkg/mod"
	"github.com/tidbcloud/dataflow-service/app/container"
	"github.com/tidbcloud/dataflow-service/app/models/domain"
	"github.com/tidbcloud/dataflow-service/app/models/tidb"
	backup_repo "github.com/tidbcloud/dataflow-service/app/repo/backup"
	repo_entity "github.com/tidbcloud/dataflow-service/app/repo/entity"
	"github.com/tidbcloud/dataflow-service/app/workflow"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"github.com/tidbcloud/pkgv2/db"
	"go.uber.org/zap"
)

const pitBackupFailureHookName = "pit-backup-dag-failure-hook"

type pitBackupFailureCompensation struct {
	expectedSwitchStatus    repo_entity.PitBackupSwitchStatus
	compensatedSwitchStatus repo_entity.PitBackupSwitchStatus
	logBackupErrMsg         string
}

var pitBackupFailureCompensations = map[string]pitBackupFailureCompensation{
	workflow.DagEnablePitBackup: {
		expectedSwitchStatus:    repo_entity.PitBackupStatusEnabling,
		compensatedSwitchStatus: repo_entity.PitBackupStatusWaitingForEnabling,
		logBackupErrMsg:         "PiTR enable aborted before log backup started",
	},
	workflow.DagRestartPitBackup: {
		expectedSwitchStatus:    repo_entity.PitBackupStatusRestarting,
		compensatedSwitchStatus: repo_entity.PitBackupStatusWaitingForRestarting,
		logBackupErrMsg:         "PiTR restart aborted before log backup started",
	},
}

type taskInstanceLister interface {
	ListTaskInstance(input *mod.ListTaskInstanceInput) ([]*entity.TaskInstance, error)
}

// PitBackupDagFailureHook owns only parent PiTR resources and state.
// The child backup-row lifecycle is separate.
type PitBackupDagFailureHook struct {
	transProvider        db.TransProvider
	logBackupRepo        backup_repo.LogBackupRepo
	pitBackupSwitchRepo  backup_repo.PitBackupSwitchRepo
	clusterGCSwitchModel tidb.ClusterGCSwitchModel
	taskInstanceLister   taskInstanceLister
}

func NewPitBackupDagFailureHook(
	transProvider db.TransProvider,
	logBackupRepo backup_repo.LogBackupRepo,
	pitBackupSwitchRepo backup_repo.PitBackupSwitchRepo,
	clusterGCSwitchModel tidb.ClusterGCSwitchModel,
	taskInstanceLister taskInstanceLister,
) *PitBackupDagFailureHook {
	return &PitBackupDagFailureHook{
		transProvider:        transProvider,
		logBackupRepo:        logBackupRepo,
		pitBackupSwitchRepo:  pitBackupSwitchRepo,
		clusterGCSwitchModel: clusterGCSwitchModel,
		taskInstanceLister:   taskInstanceLister,
	}
}

func (h *PitBackupDagFailureHook) Register() {
	ownedDagIDs := make(map[string]struct{}, len(pitBackupFailureCompensations))
	for dagID := range pitBackupFailureCompensations {
		ownedDagIDs[dagID] = struct{}{}
	}
	chainBeforeFail(pitBackupFailureHookName, ownedDagIDs, h.Handle)
}

func (h *PitBackupDagFailureHook) Handle(ctx context.Context, dagIns *entity.DagInstance) {
	ctx = log.WithLogID(ctx)
	compensation, ok := pitBackupFailureCompensations[dagIns.DagID]
	if !ok {
		return
	}
	logger := log.GetLogger(ctx).With(
		zap.String("worker", pitBackupFailureHookName),
		zap.String("dag_id", dagIns.DagID),
		zap.String("dag_ins_id", dagIns.ID),
	)
	failedTasks, err := h.taskInstanceLister.ListTaskInstance(&mod.ListTaskInstanceInput{
		DagInsID: dagIns.ID,
		Status:   []entity.TaskInstanceStatus{entity.TaskInstanceStatusFailed},
	})
	if err != nil {
		logger.Error("list failed PiTR task instances failed", zap.Error(err))
		return
	}
	if !shouldCompensatePitBackupFailure(failedTasks) {
		logger.Info("skip PiTR compensation for unrelated failed task")
		return
	}

	ids, err := parsePitBackupFailureIDs(dagIns)
	if err != nil {
		logger.Warn("skip PiTR compensation because DAG vars are invalid", zap.Error(err))
		return
	}
	logger = logger.With(
		zap.Uint64("log_backup_id", ids.logBackupID),
		zap.Uint64("pit_backup_switch_id", ids.pitBackupSwitchID),
	)
	err = h.compensate(ctx, ids, compensation)
	if err != nil {
		logger.Error("PiTR compensation transaction failed", zap.Error(err))
		return
	}
	logger.Info("PiTR compensation completed")
}

func shouldCompensatePitBackupFailure(
	failedTasks []*entity.TaskInstance,
) bool {
	for _, task := range failedTasks {
		if task.TaskID == workflow.ActionStopGCForPitBackup || task.TaskID == workflow.ActionDoSnapshotBackup {
			return true
		}
	}
	return false
}

type pitBackupFailureIDs struct {
	logBackupID       uint64
	pitBackupSwitchID uint64
}

func parsePitBackupFailureIDs(dagIns *entity.DagInstance) (pitBackupFailureIDs, error) {
	if dagIns.Vars == nil {
		return pitBackupFailureIDs{}, fmt.Errorf("DAG vars are missing")
	}
	parse := func(key string) (uint64, error) {
		value, ok := dagIns.Vars[key]
		if !ok || value.Value == "" {
			return 0, fmt.Errorf("DAG var %q is missing", key)
		}
		id, err := strconv.ParseUint(value.Value, 10, 64)
		if err != nil || id == 0 {
			return 0, fmt.Errorf("DAG var %q must be a positive integer", key)
		}
		return id, nil
	}
	switchID, err := parse(workflow.KeyPitBackupSwitchID)
	if err != nil {
		return pitBackupFailureIDs{}, err
	}
	logBackupID, err := parse(workflow.KeyEnableLogBackupID)
	if err != nil {
		return pitBackupFailureIDs{}, err
	}
	return pitBackupFailureIDs{
		logBackupID:       logBackupID,
		pitBackupSwitchID: switchID,
	}, nil
}

func (h *PitBackupDagFailureHook) compensate(
	ctx context.Context, ids pitBackupFailureIDs, compensation pitBackupFailureCompensation,
) error {
	err := h.transProvider.Transaction(ctx, func(txCtx context.Context) error {
		pitBackupSwitch, err := h.pitBackupSwitchRepo.Load(txCtx, ids.pitBackupSwitchID)
		if err != nil {
			return err
		}
		logBackup, err := h.logBackupRepo.Load(txCtx, ids.logBackupID)
		if err != nil {
			return err
		}
		if !canAbortPitBackupFailure(pitBackupSwitch, logBackup, compensation) {
			return nil
		}
		now, cleanedUp := time.Now(), true
		logBackup.Status = repo_entity.LogBackupStatusStopped
		logBackup.FinishedAt = &now
		logBackup.CleanedUp = &cleanedUp
		logBackup.ErrMsg = compensation.logBackupErrMsg
		if err := h.logBackupRepo.Save(txCtx, logBackup); err != nil {
			return err
		}
		pitBackupSwitch.Status = compensation.compensatedSwitchStatus
		pitBackupSwitch.Action = repo_entity.PitBackupSwitchAction{}
		if err := h.pitBackupSwitchRepo.Save(txCtx, pitBackupSwitch); err != nil {
			return err
		}

		// Enable must run with the ambient transaction context. Its MetaDB
		// transaction is intentionally nested into this one by the provider, so
		// a failed cluster-side GC update rolls back the compensation state too.
		// This closes the post-commit/pre-Enable crash window while keeping the
		// blocker release scoped to this exact log-backup attempt.
		gcReq := &domain.EnableTiDBClusterGCSwitchReq{
			TenantID:     logBackup.TenantID.Uint64(),
			ProjectID:    logBackup.ProjectID.Uint64(),
			ClusterID:    logBackup.ClusterID.Uint64(),
			BusinessID:   logBackup.ID.Uint64(),
			BusinessType: repo_entity.TiDBClusterGCSwitchBusinessTypePiTBackup,
		}
		if _, err := h.clusterGCSwitchModel.Enable(txCtx, gcReq); err != nil {
			return err
		}
		return nil
	})
	return err
}

func canAbortPitBackupFailure(
	pitBackupSwitch *repo_entity.PitBackupSwitch, logBackup *repo_entity.LogBackup,
	compensation pitBackupFailureCompensation,
) bool {
	return pitBackupSwitch.Status == compensation.expectedSwitchStatus &&
		logBackup.Status == repo_entity.LogBackupStatusScheduled
}

func RegisterPitBackupDagFailureHook() {
	NewPitBackupDagFailureHook(
		container.Global().TransProvider, container.Global().LogBackupRepo,
		container.Global().PitBackupSwitchRepo, container.Global().TiDBClusterGCSwitchModel,
		mod.GetStore(),
	).Register()
}