package hook

import (
	"context"
	"strconv"
	"time"

	"github.com/shiningrush/fastflow/pkg/entity"
	"github.com/tidbcloud/dataflow-service/app/container"
	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"
	export_aws "github.com/tidbcloud/dataflow-service/app/workflow/action/backup_export/aws"
	export_base "github.com/tidbcloud/dataflow-service/app/workflow/action/backup_export/base"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"go.uber.org/zap"
)

// backupExportCreateDagIDs are the DAG IDs whose failure should mark the
// associated backup_export row as Failed. Without this hook a DAG abort (e.g.
// the create_*_batch_ops_job step being watchdog-killed before its own error
// handler runs) leaves backup_export.Status at `running` with an empty job_id
// indefinitely, which keeps the BackupExportDurationTooLong alert firing.
var backupExportCreateDagIDs = map[string]struct{}{
	workflow.DagCreateAWSBackupExport: {},
	workflow.DagCreateGCPBackupExport: {},
}

type BackupExportPermissionRevoker interface {
	RevokePermissionIfUnused(ctx context.Context, backupExportID uint64) error
}

// BackupExportDagFailureHook owns the BeforeFail handler for backup export
// create DAGs. Mirrors BackupDagFailureHook: holding the repo as a field keeps
// the core handler unit-testable.
type BackupExportDagFailureHook struct {
	backupExportRepo  backup_repo.BackupExportRepo
	permissionRevoker BackupExportPermissionRevoker
}

// NewBackupExportDagFailureHook wires a hook against the given BackupExportRepo.
func NewBackupExportDagFailureHook(
	backupExportRepo backup_repo.BackupExportRepo,
	permissionRevoker BackupExportPermissionRevoker,
) *BackupExportDagFailureHook {
	return &BackupExportDagFailureHook{
		backupExportRepo:  backupExportRepo,
		permissionRevoker: permissionRevoker,
	}
}

// Register attaches Handle to the global fastflow DagInstance lifecycle. Must
// be called once at startup before any backup export DAG instance can fail.
// If called more than once (or after another BeforeFail hook), it chains with
// any previously registered BeforeFail hook and invokes the earlier hook first.
func (h *BackupExportDagFailureHook) Register() {
	chainBeforeFail("backup-export-dag-failure-hook", backupExportCreateDagIDs, h.Handle)
	chainBeforeCancel("backup-export-dag-cancel-hook", backupExportCreateDagIDs, h.HandleCancel)
}

// Handle marks the backup_export row associated with a failed create DAG as
// Failed (only if it isn't already final). DAG-ID routing and nil-guarding are
// handled by chainBeforeFail, so Handle is only ever called for an owned DAG.
func (h *BackupExportDagFailureHook) Handle(ctx context.Context, dagIns *entity.DagInstance) {
	ctx = log.WithLogID(ctx)
	logger := log.GetLogger(ctx).With(
		zap.String("worker", "backup-export-dag-failure-hook"),
		zap.String("dag_id", dagIns.DagID),
		zap.String("dag_ins_id", dagIns.ID),
		zap.String("dag_reason", dagIns.Reason),
	)

	exportIDStr := ""
	if dagIns.Vars != nil {
		if v, ok := dagIns.Vars[workflow.KeyBackupExportID]; ok {
			exportIDStr = v.Value
		}
	}
	if exportIDStr == "" {
		logger.Warn("dag instance has no backup_export_id var, skip")
		return
	}
	exportID, err := strconv.ParseUint(exportIDStr, 10, 64)
	if err != nil {
		logger.Warn("parse backup_export_id failed", zap.String("raw", exportIDStr), zap.Error(err))
		return
	}
	logger = logger.With(zap.Uint64("backup_export_id", exportID))

	export, loadErr := h.backupExportRepo.Load(ctx, exportID)
	if loadErr != nil {
		logger.Error("load backup export failed", zap.Error(loadErr))
		return
	}

	h.revokePermissionIfNeeded(ctx, logger, dagIns, exportID)

	if export.InFinalStatus() {
		logger.Info("backup export already in final status, no action",
			zap.String("current_status", string(export.Status)))
		return
	}

	reason := "backup export DAG failed without reason"
	if dagIns.Reason != "" {
		reason = "backup export DAG failed: " + dagIns.Reason
	}
	markBackupExportFailed(export, reason)
	if saveErr := h.backupExportRepo.Save(ctx, export); saveErr != nil {
		logger.Error("save backup export status failed", zap.Error(saveErr))
		return
	}
	logger.Info("marked backup export as failed via DAG failure hook",
		zap.String("status_msg", export.StatusMsg))
}

// HandleCancel releases AWS source-bucket access when a create DAG is canceled.
// It deliberately does not mutate backup_export state: CancelBackupExport owns
// the user-visible canceled status, while this hook owns only DAG cleanup.
func (h *BackupExportDagFailureHook) HandleCancel(ctx context.Context, dagIns *entity.DagInstance) {
	ctx = log.WithLogID(ctx)
	logger := log.GetLogger(ctx).With(
		zap.String("worker", "backup-export-dag-cancel-hook"),
		zap.String("dag_id", dagIns.DagID),
		zap.String("dag_ins_id", dagIns.ID),
		zap.String("dag_reason", dagIns.Reason),
	)

	exportIDStr := ""
	if dagIns.Vars != nil {
		if v, ok := dagIns.Vars[workflow.KeyBackupExportID]; ok {
			exportIDStr = v.Value
		}
	}
	if exportIDStr == "" {
		logger.Warn("dag instance has no backup_export_id var, skip")
		return
	}
	exportID, err := strconv.ParseUint(exportIDStr, 10, 64)
	if err != nil {
		logger.Warn("parse backup_export_id failed", zap.String("raw", exportIDStr), zap.Error(err))
		return
	}

	h.revokePermissionIfNeeded(ctx, logger, dagIns, exportID)
}

func (h *BackupExportDagFailureHook) revokePermissionIfNeeded(ctx context.Context, logger *zap.Logger, dagIns *entity.DagInstance, exportID uint64) {
	if dagIns.DagID != workflow.DagCreateAWSBackupExport {
		return
	}

	if revokeErr := h.permissionRevoker.RevokePermissionIfUnused(ctx, exportID); revokeErr != nil {
		logger.Error("revoke backup export source permission failed", zap.Error(revokeErr))
	}
}

// RegisterBackupExportDagFailureHook is the convenience entrypoint used by the
// fastflow bootstrap: pulls the BackupExportRepo from the global container and
// installs the hook.
func RegisterBackupExportDagFailureHook() {
	baseExportAction := export_base.NewBaseExportAction(
		container.Global().TransProvider,
		container.Global().BackupRepo,
		container.Global().BackupExportRepo,
		container.Global().BucketRepo,
	)
	permissionRevoker := export_aws.NewRevokePermissionAction(
		baseExportAction,
		container.Global().AWSBatchOpsClient,
		container.Global().BucketModelProxy,
	)
	NewBackupExportDagFailureHook(
		container.Global().BackupExportRepo,
		permissionRevoker,
	).Register()
}

// markBackupExportFailed sets the terminal failed fields directly. We do not
// use BackupExport.Fail() here because it rejects the `pending` state, and a
// DAG can abort before InitialStep flips the row from pending to running.
func markBackupExportFailed(export *repo_entity.BackupExport, reason string) {
	now := time.Now()
	export.Status = repo_entity.BackupExportStatusFailed
	// Overwrite the prior "Backup export is running" message with the failure
	// reason so operators see why it stopped.
	export.StatusMsg = reason
	if export.FinishedAt == nil {
		export.FinishedAt = &now
	}
}