package hook

import (
	"context"

	"github.com/shiningrush/fastflow/pkg/entity"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"go.uber.org/zap"
)

// chainBeforeFail installs handle as a fastflow BeforeFail hook, chaining any
// previously-registered hook. These handlers run on the executor's hot path, so
// each link is run under its own recover: a panic in handle, or in any earlier
// hook in the chain, can neither take down the fastflow executor nor prevent the
// other links from running. hookName identifies, in the panic log, which link's
// recover caught the panic.
//
// ownedDagIDs is this hook's routing allowlist: handle is invoked only when the
// failed DAG's ID is in the set. Centralizing the filter here (rather than
// re-implementing it at the top of every Handle) means a new hook physically
// cannot forget it and accidentally act on every failed DAG in the process —
// the allowlist is a required argument.
//
// Both BackupDagFailureHook and BackupExportDagFailureHook register through this
// helper so the (easy-to-get-wrong) chaining + recover + routing boilerplate
// lives in exactly one place.
func chainBeforeFail(
	hookName string,
	ownedDagIDs map[string]struct{},
	handle func(ctx context.Context, dagIns *entity.DagInstance),
) {
	prev := entity.HookDagInstance.BeforeFail
	entity.HookDagInstance.BeforeFail = func(dagIns *entity.DagInstance) {
		// Run the rest of the chain under its own recover so an earlier hook's
		// panic can't prevent this hook's handle from running below.
		if prev != nil {
			runRecovered(hookName, dagIns, func() { prev(dagIns) })
		}
		// Run this hook's handle under its own recover. Routing lives here so a
		// new hook can't forget it: only this hook's own DAGs reach its handle.
		runRecovered(hookName, dagIns, func() {
			if dagIns == nil {
				return
			}
			if _, ok := ownedDagIDs[dagIns.DagID]; !ok {
				return
			}
			handle(context.Background(), dagIns)
		})
	}
}

// chainBeforeCancel is the BeforeCancel counterpart to chainBeforeFail. It
// preserves earlier hooks, isolates panics, and routes only owned DAG IDs to
// the supplied handler.
func chainBeforeCancel(
	hookName string,
	ownedDagIDs map[string]struct{},
	handle func(ctx context.Context, dagIns *entity.DagInstance),
) {
	prev := entity.HookDagInstance.BeforeCancel
	entity.HookDagInstance.BeforeCancel = func(dagIns *entity.DagInstance) {
		if prev != nil {
			runRecovered(hookName, dagIns, func() { prev(dagIns) })
		}
		runRecovered(hookName, dagIns, func() {
			if dagIns == nil {
				return
			}
			if _, ok := ownedDagIDs[dagIns.DagID]; !ok {
				return
			}
			handle(context.Background(), dagIns)
		})
	}
}

// runRecovered runs fn, swallowing and logging any panic so a single hook can
// neither crash the fastflow executor nor abort the rest of the BeforeFail
// chain. It never derefs a nil dagIns from inside the recover handler itself.
func runRecovered(hookName string, dagIns *entity.DagInstance, fn func()) {
	defer func() {
		if r := recover(); r != nil {
			logger := log.GetLogger(context.Background())
			if dagIns != nil {
				logger = logger.With(
					zap.String("dag_id", dagIns.DagID),
					zap.String("dag_ins_id", dagIns.ID),
				)
			}
			logger.Error(
				"fastflow BeforeFail hook panicked, swallowing to keep executor alive and chain running",
				zap.String("recovering_hook", hookName),
				zap.Any("recover", r),
			)
		}
	}()
	fn()
}