package workflow

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"strconv"
	"testing"
	"time"

	gomysql "github.com/go-sql-driver/mysql"
	"github.com/stretchr/testify/require"
	"github.com/testcontainers/testcontainers-go"
	mysqlcontainer "github.com/testcontainers/testcontainers-go/modules/mysql"
	"github.com/tidbcloud/fastflow/pkg/entity"
	"github.com/tidbcloud/fastflow/pkg/mod"
	mysqlstore "github.com/tidbcloud/fastflow/store/mysql"
	pkgv2db "github.com/tidbcloud/pkgv2/db"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

const (
	watchDogTestDBName   = "workflow_watchdog_test"
	watchDogTestUser     = "root"
	watchDogTestPassword = "password"
	watchDogTestWorker   = "worker-1"
)

var watchDogTestProvider pkgv2db.TransProvider
var watchDogTestMySQLConfig *gomysql.Config

func TestMain(m *testing.M) {
	ctx := context.Background()
	mysqlContainer, host, port, err := startWatchDogMysqlContainer(
		ctx,
		watchDogTestDBName,
		watchDogTestUser,
		watchDogTestPassword,
	)
	if err != nil {
		panic(err)
	}
	defer func() {
		if err := mysqlContainer.Terminate(ctx); err != nil {
			panic(err)
		}
	}()

	cfg := gomysql.NewConfig()
	cfg.User = watchDogTestUser
	cfg.Passwd = watchDogTestPassword
	cfg.Net = "tcp"
	cfg.Addr = fmt.Sprintf("%s:%d", host, port)
	cfg.DBName = watchDogTestDBName
	cfg.ParseTime = true
	watchDogTestMySQLConfig = cfg

	gormDB, err := gorm.Open(mysql.Open(cfg.FormatDSN()), gormConfig)
	if err != nil {
		panic(err)
	}
	watchDogTestProvider = pkgv2db.NewProvider(gormDB)
	if err := newStore(cfg); err != nil {
		panic(err)
	}
	mod.SetKeeper(watchDogTestKeeper{})

	m.Run()
	mod.GetStore().Close()
}

func TestFastflowHasCmdSchemaMigration(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	require.NoError(t, validateFastflowHasCmdSchema(watchDogTestMySQLConfig))

	prefix := "has-cmd-schema"
	dagIDs := seedFailedDagBacklog(ctx, t, prefix, 1, true)
	instances, err := mod.GetStore().ListDagInstance(&mod.ListDagInstanceInput{
		Worker: watchDogTestWorker,
		HasCmd: true,
	})
	require.NoError(t, err)
	require.Len(t, instances, 1)
	require.Equal(t, dagIDs[0], instances[0].ID)
}

func TestFastflowStoreInitCreatesExpectedHasCmdSchemaOnBrandNewDatabase(t *testing.T) {
	ctx := context.Background()
	cfg := createWatchDogTestDatabase(ctx, t, "workflow_fresh_has_cmd")

	store := mysqlstore.NewStore(&mysqlstore.StoreOption{
		MySQLConfig:     cfg,
		MigrationSwitch: true,
		GormConfig:      gormConfig,
	})
	err := store.Init()
	require.NoError(t, err)
	defer store.Close()

	require.NoError(t, validateFastflowHasCmdSchema(cfg))
}

func TestFastflowStoreInitDoesNotRepairIncidentPhysicalHasCmdColumn(t *testing.T) {
	ctx := context.Background()
	cfg := createWatchDogTestDatabase(ctx, t, "workflow_bad_has_cmd")
	createIncidentPhysicalHasCmdDagInstancesTable(ctx, t, cfg)

	store := mysqlstore.NewStore(&mysqlstore.StoreOption{
		MySQLConfig:     cfg,
		MigrationSwitch: true,
		GormConfig:      gormConfig,
	})
	err := store.Init()
	require.NoError(t, err)
	defer store.Close()

	err = validateFastflowHasCmdSchema(cfg)
	require.ErrorContains(t, err, "dag_instances.has_cmd must be a generated column")
}

func TestNewStoreRejectsIncidentPhysicalHasCmdColumn(t *testing.T) {
	ctx := context.Background()
	cfg := createWatchDogTestDatabase(ctx, t, "workflow_bad_has_cmd")
	createIncidentPhysicalHasCmdDagInstancesTable(ctx, t, cfg)

	err := newStore(cfg)
	require.ErrorContains(t, err, "dag_instances.has_cmd must be a generated column")
}

func createWatchDogTestDatabase(ctx context.Context, t *testing.T, namePrefix string) *gomysql.Config {
	t.Helper()

	adminCfg := *watchDogTestMySQLConfig
	dbName := fmt.Sprintf("%s_%d", namePrefix, time.Now().UnixNano())
	adminDB, err := sql.Open("mysql", adminCfg.FormatDSN())
	require.NoError(t, err)
	t.Cleanup(func() {
		_, dropErr := adminDB.ExecContext(ctx, fmt.Sprintf("DROP DATABASE `%s`", dbName))
		require.NoError(t, dropErr)
		require.NoError(t, adminDB.Close())
	})

	_, err = adminDB.ExecContext(ctx, fmt.Sprintf("CREATE DATABASE `%s`", dbName))
	require.NoError(t, err)

	cfg := *watchDogTestMySQLConfig
	cfg.DBName = dbName
	return &cfg
}

func createIncidentPhysicalHasCmdDagInstancesTable(ctx context.Context, t *testing.T, cfg *gomysql.Config) {
	t.Helper()

	badDB, err := sql.Open("mysql", cfg.FormatDSN())
	require.NoError(t, err)
	defer badDB.Close()

	_, err = badDB.ExecContext(ctx, `
		CREATE TABLE dag_instances (
			id VARCHAR(256) PRIMARY KEY,
			worker VARCHAR(256),
			cmd JSON,
			has_cmd TINYINT DEFAULT NULL,
			INDEX idx_worker_has_cmd(worker, has_cmd)
		)
	`)
	require.NoError(t, err)
}

func TestWatchDogRetriesFailedDagWithFailedTaskOnly(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-failed-with-canceled"
	failedTaskID := "task-failed"
	canceledTaskID := "task-canceled"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, failedTaskID, dagID, entity.TaskInstanceStatusFailed, staleUpdatedAt(), "action failed")
	seedTaskInstance(ctx, t, canceledTaskID, dagID, entity.TaskInstanceStatusCanceled, staleUpdatedAt(), "operator canceled")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, string(entity.CommandNameRetry), string(dag.Cmd.Name))
	require.Equal(t, []string{failedTaskID}, dag.Cmd.TargetTaskInsIDs)
	require.Equal(t, entity.TaskInstanceStatusCanceled, getTaskInstance(ctx, t, canceledTaskID).Status)
}

func TestWatchDogFinalizesStaleSuccessTaskTree(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-success-finalizable"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, staleUpdatedAt(), "old failure")
	seedTaskInstance(ctx, t, "task-success", dagID, entity.TaskInstanceStatusSuccess, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, "task-skipped", dagID, entity.TaskInstanceStatusSkipped, staleUpdatedAt(), "")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.Equal(t, entity.DagInstanceStatusSuccess, dag.Status)
	require.Empty(t, dag.Reason)
	require.Nil(t, dag.Cmd)
}

func TestCompetitiveWatchDogFinalizesSuccessTaskTree(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "competitive-dag-success-finalizable"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, time.Now().Unix(), "old failure")
	seedTaskInstance(ctx, t, "competitive-task-success", dagID, entity.TaskInstanceStatusSuccess, time.Now().Unix(), "")
	seedTaskInstance(ctx, t, "competitive-task-skipped", dagID, entity.TaskInstanceStatusSkipped, time.Now().Unix(), "")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeCompetitive).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.Equal(t, entity.DagInstanceStatusSuccess, dag.Status)
	require.Empty(t, dag.Reason)
	require.Nil(t, dag.Cmd)
}

func TestWatchDogRepairsFullyStuckRunningDagBeforeRetry(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-running-repair"
	taskID := "task-running"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusRunning, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusRunning, staleUpdatedAt(), "")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.Equal(t, entity.DagInstanceStatusFailed, dag.Status)
	require.Equal(t, watchDogRepairReason, dag.Reason)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, string(entity.CommandNameRetry), string(dag.Cmd.Name))
	require.Equal(t, []string{taskID}, dag.Cmd.TargetTaskInsIDs)

	task := getTaskInstance(ctx, t, taskID)
	require.Equal(t, entity.TaskInstanceStatusFailed, task.Status)
	require.Equal(t, watchDogRepairReason, task.Reason)
}

func TestWatchDogDoesNotRetryTaskBeforeItsTimeout(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-long-running-task"
	taskID := "task-long-running"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusRunning, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusRunning, staleUpdatedAt(), "")
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.TaskInstance{}).
		Where("id = ?", taskID).
		Update("timeout_secs", int((15*time.Minute).Seconds())).Error)

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.Equal(t, entity.DagInstanceStatusRunning, dag.Status)
	require.Nil(t, dag.Cmd)
	require.Equal(t, entity.TaskInstanceStatusRunning, getTaskInstance(ctx, t, taskID).Status)
}

func TestWatchDogUsesTimeoutOfNonTerminalTaskOnly(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-mixed-task-timeouts"
	completedTaskID := "task-completed-long"
	stuckTaskID := "task-stuck-short"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusRunning, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, completedTaskID, dagID, entity.TaskInstanceStatusSuccess, time.Now().Unix()-1, "")
	seedTaskInstance(ctx, t, stuckTaskID, dagID, entity.TaskInstanceStatusRunning, staleUpdatedAt(), "")
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.TaskInstance{}).
		Where("id = ?", completedTaskID).
		Updates(map[string]any{"timeout_secs": int((15 * time.Minute).Seconds())}).Error)

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.Equal(t, entity.DagInstanceStatusFailed, dag.Status)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, []string{stuckTaskID}, dag.Cmd.TargetTaskInsIDs)
}

func TestWatchDogSkipsWhenAnyRowIsFresh(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-has-fresh-task"
	taskID := "task-fresh"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusRunning, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusRunning, time.Now().Unix(), "")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.Equal(t, entity.DagInstanceStatusRunning, dag.Status)
	require.Nil(t, dag.Cmd)
	require.Equal(t, entity.TaskInstanceStatusRunning, getTaskInstance(ctx, t, taskID).Status)
}

func TestWatchDogSkipsCanceledOnlyAndFreshCommandBearingDags(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	canceledDagID := "dag-canceled-only"
	seedDagInstance(ctx, t, canceledDagID, entity.DagInstanceStatusFailed, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, "task-canceled-only", canceledDagID, entity.TaskInstanceStatusCanceled, staleUpdatedAt(), "")

	commandDagID := "dag-with-command"
	commandUpdatedAtSec := time.Now().Unix()
	seedDagInstance(ctx, t, commandDagID, entity.DagInstanceStatusFailed, commandUpdatedAtSec, "")
	seedTaskInstance(ctx, t, "task-command-failed", commandDagID, entity.TaskInstanceStatusFailed, staleUpdatedAt(), "")
	commandDag := getDagInstance(ctx, t, commandDagID)
	commandDag.Cmd = &entity.Command{Name: entity.CommandNameRetry, TargetTaskInsIDs: []string{"task-command-failed"}}
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(commandDag).
		Select("Cmd").
		Updates(commandDag).Error)
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.DagInstance{}).
		Where("id = ?", commandDagID).
		UpdateColumn("updated_at", commandUpdatedAtSec).Error)

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	require.Nil(t, getDagInstance(ctx, t, canceledDagID).Cmd)
	commandDag = getDagInstance(ctx, t, commandDagID)
	require.NotNil(t, commandDag.Cmd)
	require.Equal(t, []string{"task-command-failed"}, commandDag.Cmd.TargetTaskInsIDs)
	require.Equal(t, commandUpdatedAtSec, commandDag.UpdatedAt)
}

func TestWatchDogReissuesStaleRetryCommand(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-stale-retry-command"
	taskID := "task-stale-retry-command"
	oldUpdatedAt := staleUpdatedAt()
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, oldUpdatedAt, "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusFailed, oldUpdatedAt, "parser update failed")
	commandDag := getDagInstance(ctx, t, dagID)
	commandDag.Cmd = &entity.Command{Name: entity.CommandNameRetry, TargetTaskInsIDs: []string{taskID}}
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(commandDag).
		Select("Cmd").
		Updates(commandDag).Error)
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.DagInstance{}).
		Where("id = ?", dagID).
		UpdateColumn("updated_at", oldUpdatedAt).Error)

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, string(entity.CommandNameRetry), string(dag.Cmd.Name))
	require.Equal(t, []string{taskID}, dag.Cmd.TargetTaskInsIDs)
	require.Greater(t, dag.UpdatedAt, oldUpdatedAt)
	require.Equal(t, entity.TaskInstanceStatusFailed, getTaskInstance(ctx, t, taskID).Status)
}

func TestWatchDogDoesNotReissueStaleRetryCommandWhenTaskIsFresh(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-stale-command-fresh-task"
	taskID := "task-stale-command-fresh-task"
	oldUpdatedAt := staleUpdatedAt()
	freshUpdatedAtSec := time.Now().Unix()
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, oldUpdatedAt, "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusFailed, freshUpdatedAtSec, "parser update failed")
	commandDag := getDagInstance(ctx, t, dagID)
	commandDag.Cmd = &entity.Command{Name: entity.CommandNameRetry, TargetTaskInsIDs: []string{taskID}}
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(commandDag).
		Select("Cmd").
		Updates(commandDag).Error)
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.DagInstance{}).
		Where("id = ?", dagID).
		UpdateColumn("updated_at", oldUpdatedAt).Error)

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, []string{taskID}, dag.Cmd.TargetTaskInsIDs)
	require.Equal(t, oldUpdatedAt, dag.UpdatedAt)
	require.Equal(t, freshUpdatedAtSec, getTaskInstance(ctx, t, taskID).UpdatedAt)
}

func TestWatchDogTreatsJSONNullCommandAsCommandFree(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "dag-json-null-command"
	taskID := "task-json-null-command"
	oldUpdatedAt := staleUpdatedAt()
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, oldUpdatedAt, "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusFailed, oldUpdatedAt, "old failure")
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.DagInstance{}).
		Where("id = ?", dagID).
		Update("cmd", gorm.Expr("CAST('null' AS JSON)")).Error)
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).
		Model(&entity.DagInstance{}).
		Where("id = ?", dagID).
		UpdateColumn("updated_at", oldUpdatedAt).Error)

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, dagID)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, string(entity.CommandNameRetry), string(dag.Cmd.Name))
	require.Equal(t, []string{taskID}, dag.Cmd.TargetTaskInsIDs)
}

func TestWatchDogLargeBacklogDrainsAcrossLoopsWithoutStarvation(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagIDs := seedFailedDagBacklog(ctx, t, "backlog-drain", watchDogCandidateLimit+1, false)

	wd := newWatchDog(watchDogTestProvider, ModeDispatcher)
	require.NoError(t, wd.retryDagInstancesWithContext(ctx))
	require.Equal(t, int64(watchDogCandidateLimit), countCommandedDagInstances(ctx, t, "backlog-drain"))
	require.Nil(t, getDagInstance(ctx, t, dagIDs[watchDogCandidateLimit]).Cmd)

	require.NoError(t, wd.retryDagInstancesWithContext(ctx))
	require.Equal(t, int64(watchDogCandidateLimit+1), countCommandedDagInstances(ctx, t, "backlog-drain"))
	require.NotNil(t, getDagInstance(ctx, t, dagIDs[watchDogCandidateLimit]).Cmd)
}

func TestWatchDogCommandBearingBacklogDoesNotConsumeBatchCapacity(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	seedFailedDagBacklog(ctx, t, "commanded-backlog", watchDogCandidateLimit, true)
	openDagID := "commanded-backlog-open-dag"
	openTaskID := "commanded-backlog-open-task"
	seedDagInstance(ctx, t, openDagID, entity.DagInstanceStatusFailed, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, openTaskID, openDagID, entity.TaskInstanceStatusFailed, staleUpdatedAt(), "")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	dag := getDagInstance(ctx, t, openDagID)
	require.NotNil(t, dag.Cmd)
	require.Equal(t, []string{openTaskID}, dag.Cmd.TargetTaskInsIDs)
}

func TestWatchDogCandidateScanDoesNotHydrateDagInstanceTags(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	dagID := "tag-hydration-guard-dag"
	taskID := "tag-hydration-guard-task"
	seedDagInstance(ctx, t, dagID, entity.DagInstanceStatusFailed, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, taskID, dagID, entity.TaskInstanceStatusFailed, staleUpdatedAt(), "")
	seedDagInstanceTag(ctx, t, "tag-hydration-guard-tag", dagID, "org_id", "org-1")

	db := watchDogTestProvider.WriteDB(ctx)
	require.NoError(t, db.Exec("ALTER TABLE dag_instance_tags RENAME TO dag_instance_tags_guard").Error)
	defer func() {
		require.NoError(t, db.Exec("ALTER TABLE dag_instance_tags_guard RENAME TO dag_instance_tags").Error)
	}()

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))
	require.NotNil(t, getDagInstance(ctx, t, dagID).Cmd)
}

func TestWatchDogBlockedBacklogDoesNotBlockUnblockedRecovery(t *testing.T) {
	ctx := context.Background()
	resetWatchDogTables(ctx, t)

	blockedDagID := "blocked-backpressure-dag"
	seedDagInstance(ctx, t, blockedDagID, entity.DagInstanceStatusBlocked, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, "blocked-backpressure-task", blockedDagID, entity.TaskInstanceStatusBlocked, staleUpdatedAt(), "")

	failedDagID := "blocked-backpressure-failed-dag"
	failedTaskID := "blocked-backpressure-failed-task"
	seedDagInstance(ctx, t, failedDagID, entity.DagInstanceStatusFailed, staleUpdatedAt(), "")
	seedTaskInstance(ctx, t, failedTaskID, failedDagID, entity.TaskInstanceStatusFailed, staleUpdatedAt(), "")

	require.NoError(t, newWatchDog(watchDogTestProvider, ModeDispatcher).retryDagInstancesWithContext(ctx))

	require.Nil(t, getDagInstance(ctx, t, blockedDagID).Cmd)
	require.Equal(t, entity.DagInstanceStatusBlocked, getDagInstance(ctx, t, blockedDagID).Status)

	failedDag := getDagInstance(ctx, t, failedDagID)
	require.NotNil(t, failedDag.Cmd)
	require.Equal(t, []string{failedTaskID}, failedDag.Cmd.TargetTaskInsIDs)
}

func resetWatchDogTables(ctx context.Context, t *testing.T) {
	t.Helper()
	db := watchDogTestProvider.WriteDB(ctx)
	require.NoError(t, db.Exec("DELETE FROM dag_instance_tags").Error)
	require.NoError(t, db.Exec("DELETE FROM task_instances").Error)
	require.NoError(t, db.Exec("DELETE FROM dag_instances").Error)
}

func seedDagInstance(
	ctx context.Context,
	t *testing.T,
	id string,
	status entity.DagInstanceStatus,
	updatedAt int64,
	reason string,
) {
	t.Helper()
	dag := &entity.DagInstance{
		BaseInfo: entity.BaseInfo{
			ID: id,
		},
		DagID:   "test-dag",
		Trigger: entity.TriggerManually,
		Worker:  watchDogTestWorker,
		Status:  status,
		Reason:  reason,
	}
	db := watchDogTestProvider.WriteDB(ctx)
	require.NoError(t, db.Create(dag).Error)
	require.NoError(t, db.Model(&entity.DagInstance{}).
		Where("id = ?", id).
		UpdateColumn("updated_at", updatedAt).Error)
}

func seedTaskInstance(
	ctx context.Context,
	t *testing.T,
	id string,
	dagInsID string,
	status entity.TaskInstanceStatus,
	updatedAt int64,
	reason string,
) {
	t.Helper()
	task := &entity.TaskInstance{
		BaseInfo: entity.BaseInfo{
			ID: id,
		},
		TaskID:      id,
		DagInsID:    dagInsID,
		ActionName:  "test-action",
		TimeoutSecs: int(executorTimeout.Seconds()),
		Status:      status,
		Reason:      reason,
	}
	db := watchDogTestProvider.WriteDB(ctx)
	require.NoError(t, db.Create(task).Error)
	require.NoError(t, db.Model(&entity.TaskInstance{}).
		Where("id = ?", id).
		UpdateColumn("updated_at", updatedAt).Error)
}

func seedFailedDagBacklog(ctx context.Context, t *testing.T, prefix string, count int, withCmd bool) []string {
	t.Helper()
	updatedAt := staleUpdatedAt()
	dags := make([]entity.DagInstance, 0, count)
	tasks := make([]entity.TaskInstance, 0, count)
	dagIDs := make([]string, 0, count)

	for i := range count {
		dagID := fmt.Sprintf("%s-dag-%04d", prefix, i)
		taskID := fmt.Sprintf("%s-task-%04d", prefix, i)
		dagIDs = append(dagIDs, dagID)

		dag := entity.DagInstance{
			BaseInfo: entity.BaseInfo{
				ID:        dagID,
				CreatedAt: updatedAt,
				UpdatedAt: updatedAt,
			},
			DagID:   "test-dag",
			Trigger: entity.TriggerManually,
			Worker:  watchDogTestWorker,
			Status:  entity.DagInstanceStatusFailed,
		}
		if withCmd {
			dag.Cmd = &entity.Command{Name: entity.CommandNameRetry, TargetTaskInsIDs: []string{taskID}}
		}
		dags = append(dags, dag)
		tasks = append(tasks, entity.TaskInstance{
			BaseInfo: entity.BaseInfo{
				ID:        taskID,
				CreatedAt: updatedAt,
				UpdatedAt: updatedAt,
			},
			TaskID:      taskID,
			DagInsID:    dagID,
			ActionName:  "test-action",
			TimeoutSecs: int(executorTimeout.Seconds()),
			Status:      entity.TaskInstanceStatusFailed,
		})
	}

	db := watchDogTestProvider.WriteDB(ctx)
	require.NoError(t, db.CreateInBatches(dags, 100).Error)
	require.NoError(t, db.CreateInBatches(tasks, 100).Error)
	return dagIDs
}

func seedDagInstanceTag(ctx context.Context, t *testing.T, id, dagInsID, key, value string) {
	t.Helper()
	tag := &entity.DagInstanceTag{
		BaseInfo: entity.BaseInfo{
			ID: id,
		},
		DagInsId: dagInsID,
		Key:      key,
		Value:    value,
	}
	require.NoError(t, watchDogTestProvider.WriteDB(ctx).Create(tag).Error)
}

func getDagInstance(ctx context.Context, t *testing.T, id string) *entity.DagInstance {
	t.Helper()
	var dag entity.DagInstance
	require.NoError(t, watchDogTestProvider.ReadDB(ctx).Where("id = ?", id).First(&dag).Error)
	return &dag
}

func getTaskInstance(ctx context.Context, t *testing.T, id string) *entity.TaskInstance {
	t.Helper()
	var task entity.TaskInstance
	require.NoError(t, watchDogTestProvider.ReadDB(ctx).Where("id = ?", id).First(&task).Error)
	return &task
}

func countCommandedDagInstances(ctx context.Context, t *testing.T, prefix string) int64 {
	t.Helper()
	var count int64
	require.NoError(t, watchDogTestProvider.ReadDB(ctx).
		Model(&entity.DagInstance{}).
		Where("id LIKE ? AND cmd IS NOT NULL", prefix+"-dag-%").
		Count(&count).Error)
	return count
}

func staleUpdatedAt() int64 {
	return time.Now().Add(-3 * executorTimeout).Unix()
}

func startWatchDogMysqlContainer(
	ctx context.Context,
	dbName, user, passwd string,
) (c *mysqlcontainer.MySQLContainer, host string, port uint16, err error) {
	if err := os.Setenv("TESTCONTAINERS_RYUK_DISABLED", "true"); err != nil {
		return nil, "", 0, err
	}
	if err := os.Setenv("DOCKER_API_VERSION", "1.44"); err != nil {
		return nil, "", 0, err
	}
	mysqlContainer, err := mysqlcontainer.RunContainer(
		ctx,
		testcontainers.WithImage("gcr.io/pingcap-public/third-party/mysql:8.0"),
		mysqlcontainer.WithDatabase(dbName),
		mysqlcontainer.WithUsername(user),
		mysqlcontainer.WithPassword(passwd),
		testcontainers.CustomizeRequestOption(func(req *testcontainers.GenericContainerRequest) {
			req.Tmpfs = map[string]string{"/var/lib/mysql": "rw"}
		}),
	)
	if err != nil {
		return nil, "", 0, err
	}

	host, err = mysqlContainer.Host(ctx)
	if err != nil {
		return nil, "", 0, terminateWatchDogMysqlContainer(ctx, mysqlContainer, err)
	}
	portS, err := mysqlContainer.MappedPort(ctx, "3306/tcp")
	if err != nil {
		return nil, "", 0, terminateWatchDogMysqlContainer(ctx, mysqlContainer, err)
	}
	portParsed, err := strconv.ParseUint(portS.Port(), 10, 16)
	if err != nil {
		return nil, "", 0, terminateWatchDogMysqlContainer(ctx, mysqlContainer, err)
	}

	return mysqlContainer, host, uint16(portParsed), nil
}

func terminateWatchDogMysqlContainer(
	ctx context.Context,
	mysqlContainer *mysqlcontainer.MySQLContainer,
	cause error,
) error {
	if terminateErr := mysqlContainer.Terminate(ctx); terminateErr != nil {
		return fmt.Errorf("%w; failed to terminate mysql container: %w", cause, terminateErr)
	}
	return cause
}

type watchDogTestKeeper struct{}

func (watchDogTestKeeper) Close() {}

func (watchDogTestKeeper) IsLeader() bool {
	return true
}

func (watchDogTestKeeper) IsAlive(workerKey string) (bool, error) {
	return workerKey == watchDogTestWorker, nil
}

func (watchDogTestKeeper) AliveNodes() ([]string, error) {
	return []string{watchDogTestWorker}, nil
}

func (watchDogTestKeeper) WorkerKey() string {
	return watchDogTestWorker
}

func (watchDogTestKeeper) WorkerNumber() int {
	return 1
}

func (watchDogTestKeeper) NewMutex(key string) mod.DistributedMutex {
	return nil
}