package workflow

import (
	"context"
	"errors"
	"slices"
	"strings"
	"time"

	"github.com/samber/lo"
	"github.com/tidbcloud/fastflow/pkg/entity"
	"github.com/tidbcloud/fastflow/pkg/mod"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

const watchDogLockingStrength = "UPDATE"

type dispatcherWatchDog struct {
	provider db.TransProvider
	cmd      *mod.DefCommander
}

type watchDogRecoveryAction struct {
	dagID        string
	dagInsID     string
	retryTaskIDs []string
}

func newDispatcherWatchDog(provider db.TransProvider) *dispatcherWatchDog {
	return &dispatcherWatchDog{provider: provider, cmd: &mod.DefCommander{}}
}

func (wd *dispatcherWatchDog) retry(ctx context.Context) error {
	nowSec := time.Now().Unix()
	dagInsIDs, err := wd.listCandidateDagInstanceIDs(ctx)
	if err != nil {
		return err
	}
	staleCommandDagInsIDs, err := wd.listStaleCommandDagInstanceIDs(ctx, nowSec)
	if err != nil {
		return err
	}
	dagInsIDs = append(dagInsIDs, staleCommandDagInsIDs...)

	for _, dagInsID := range dagInsIDs {
		action, err := wd.prepareDagRecovery(ctx, dagInsID, nowSec)
		if err != nil {
			return err
		}
		if len(action.retryTaskIDs) == 0 {
			continue
		}

		if err := wd.cmd.RetryTask(action.retryTaskIDs); err != nil {
			MetricsWatchDogRetryDagInsTotal.WithLabelValues(action.dagID, action.dagInsID, "true").Inc()
			tilog.PlainError(
				"create retry task command failed",
				zap.String("dag_id", action.dagID),
				zap.String("instance_id", action.dagInsID),
				zap.Strings("task_ids", action.retryTaskIDs),
				zap.Error(err),
			)
			continue
		}

		MetricsWatchDogRetryDagInsTotal.WithLabelValues(action.dagID, action.dagInsID, "false").Inc()
	}

	return nil
}

func (wd *dispatcherWatchDog) listCandidateDagInstanceIDs(ctx context.Context) ([]string, error) {
	var dagInsIDs []string
	err := wd.provider.ReadDB(ctx).
		Model(&entity.DagInstance{}).
		Where("status in ?", []entity.DagInstanceStatus{
			entity.DagInstanceStatusFailed,
			entity.DagInstanceStatusRunning,
		}).
		Where("(cmd IS NULL OR has_cmd = 0 OR JSON_TYPE(cmd) = 'NULL')").
		Order("updated_at ASC, id ASC").
		Limit(watchDogCandidateLimit).
		Pluck("id", &dagInsIDs).Error
	return dagInsIDs, err
}

func (wd *dispatcherWatchDog) listStaleCommandDagInstanceIDs(ctx context.Context, now int64) ([]string, error) {
	var dagInsIDs []string
	staleBefore := now - int64((2 * executorTimeout).Seconds())
	err := wd.provider.ReadDB(ctx).
		Model(&entity.DagInstance{}).
		Where("status in ?", []entity.DagInstanceStatus{
			entity.DagInstanceStatusFailed,
			entity.DagInstanceStatusRunning,
		}).
		Where("cmd IS NOT NULL AND JSON_TYPE(cmd) <> 'NULL'").
		Where("updated_at <= ?", staleBefore).
		Order("updated_at ASC, id ASC").
		Limit(watchDogCandidateLimit).
		Pluck("id", &dagInsIDs).Error
	return dagInsIDs, err
}

func (wd *dispatcherWatchDog) prepareDagRecovery(
	ctx context.Context,
	dagInsID string,
	now int64,
) (watchDogRecoveryAction, error) {
	var action watchDogRecoveryAction
	err := wd.provider.Transaction(ctx, func(ctx context.Context) error {
		tx := wd.provider.WriteDB(ctx)

		dagIns, tasks, err := loadWatchDogLockedSnapshot(tx, dagInsID)
		if err != nil {
			if errors.Is(err, gorm.ErrRecordNotFound) {
				tilog.PlainWarn("skip retry dag instance because it no longer exists", zap.String("instance_id", dagInsID))
				return nil
			}
			return err
		}

		action.dagID = dagIns.DagID
		action.dagInsID = dagIns.ID

		if !isWatchDogMutationCandidate(dagIns.Status) {
			return nil
		}
		if len(tasks) == 0 {
			tilog.PlainWarn(
				"skip retry dag instance because it has no task instances",
				zap.String("dag_id", dagIns.DagID),
				zap.String("instance_id", dagIns.ID),
			)
			return nil
		}
		if dagIns.Cmd != nil {
			return prepareStaleRetryCommandRecovery(tx, dagIns, tasks, now, &action)
		}

		failedTaskIDs := collectWatchDogRetryTargetIDs(tasks)
		if dagIns.Status == entity.DagInstanceStatusFailed && len(failedTaskIDs) > 0 {
			action.retryTaskIDs = failedTaskIDs
			return nil
		}

		if !isFullStuck(dagIns, tasks, now) {
			return nil
		}

		if allTasksSuccessOrSkipped(tasks) {
			return tx.Model(&entity.DagInstance{}).
				Where("id = ?", dagIns.ID).
				Updates(map[string]any{
					"status": entity.DagInstanceStatusSuccess,
					"reason": "",
				}).Error
		}

		repairableTaskIDs := collectRepairableTaskIDs(tasks)
		if len(repairableTaskIDs) > 0 {
			if err := tx.Model(&entity.TaskInstance{}).
				Where("id in ?", repairableTaskIDs).
				Updates(map[string]any{
					"status": entity.TaskInstanceStatusFailed,
					"reason": watchDogRepairReason,
				}).Error; err != nil {
				return err
			}
		}

		if len(repairableTaskIDs) > 0 || (dagIns.Status == entity.DagInstanceStatusRunning && len(failedTaskIDs) > 0) {
			if err := tx.Model(&entity.DagInstance{}).
				Where("id = ?", dagIns.ID).
				Updates(map[string]any{
					"status": entity.DagInstanceStatusFailed,
					"reason": watchDogRepairReason,
				}).Error; err != nil {
				return err
			}
		}

		if len(repairableTaskIDs) == 0 && len(failedTaskIDs) == 0 {
			tilog.PlainDebug(
				"skip retry dag instance because it has no watchdog retry targets",
				zap.String("dag_id", dagIns.DagID),
				zap.String("instance_id", dagIns.ID),
			)
			return nil
		}

		action.retryTaskIDs = append(slices.Clone(failedTaskIDs), repairableTaskIDs...)
		return nil
	})
	return action, err
}

func prepareStaleRetryCommandRecovery(
	tx *gorm.DB,
	dagIns *entity.DagInstance,
	tasks []*entity.TaskInstance,
	now int64,
	action *watchDogRecoveryAction,
) error {
	if dagIns.Cmd.Name != entity.CommandNameRetry {
		tilog.PlainWarn(
			"skip stale command dag instance because command is not retry",
			zap.String("dag_id", dagIns.DagID),
			zap.String("instance_id", dagIns.ID),
			zap.String("command", string(dagIns.Cmd.Name)),
		)
		return nil
	}
	if !isFullStuck(dagIns, tasks, now) {
		return nil
	}

	retryTaskIDs := collectRetryCommandTargetIDs(dagIns.Cmd.TargetTaskInsIDs, tasks)
	if len(retryTaskIDs) == 0 {
		tilog.PlainWarn(
			"skip stale retry command because it has no retryable targets",
			zap.String("dag_id", dagIns.DagID),
			zap.String("instance_id", dagIns.ID),
			zap.Strings("task_ids", dagIns.Cmd.TargetTaskInsIDs),
		)
		return nil
	}

	if err := tx.Model(&entity.DagInstance{}).
		Where("id = ?", dagIns.ID).
		Updates(map[string]any{"cmd": nil}).Error; err != nil {
		return err
	}
	action.retryTaskIDs = retryTaskIDs
	return nil
}

func loadWatchDogLockedSnapshot(
	tx *gorm.DB,
	dagInsID string,
) (*entity.DagInstance, []*entity.TaskInstance, error) {
	var dagIns entity.DagInstance
	if err := tx.Clauses(clause.Locking{Strength: watchDogLockingStrength}).
		Where("id = ?", dagInsID).
		First(&dagIns).Error; err != nil {
		return nil, nil, err
	}

	var tasks []*entity.TaskInstance
	if err := tx.Clauses(clause.Locking{Strength: watchDogLockingStrength}).
		Where("dag_ins_id = ?", dagInsID).
		Order("id ASC").
		Find(&tasks).Error; err != nil {
		return nil, nil, err
	}

	return &dagIns, tasks, nil
}

func isWatchDogMutationCandidate(status entity.DagInstanceStatus) bool {
	return status == entity.DagInstanceStatusFailed || status == entity.DagInstanceStatusRunning
}

func isFullStuck(dagIns *entity.DagInstance, tasks []*entity.TaskInstance, now int64) bool {
	latestDeadline := int64(0)
	for _, task := range tasks {
		if task == nil {
			continue
		}
		if task.Status == entity.TaskInstanceStatusSuccess || task.Status == entity.TaskInstanceStatusSkipped {
			// Successfully completed or skipped tasks do not represent work that the watchdog needs to recover.
			continue
		}
		staleAfter := 2 * executorTimeout
		if task.TimeoutSecs > 0 {
			switch task.Status {
			case entity.TaskInstanceStatusRunning,
				entity.TaskInstanceStatusEnding,
				entity.TaskInstanceStatusRetrying,
				entity.TaskInstanceStatusContinue:
				// Do not retry a task while its action can still be running within its own timeout.
				taskTimeout := time.Duration(task.TimeoutSecs) * time.Second
				if taskTimeout > staleAfter {
					staleAfter = taskTimeout
				}
			}
		}
		deadline := task.UpdatedAt + int64(staleAfter.Seconds())
		if deadline > latestDeadline {
			latestDeadline = deadline
		}
	}
	if latestDeadline == 0 {
		latestDeadline = dagIns.UpdatedAt + int64((2 * executorTimeout).Seconds())
	}
	return now > latestDeadline
}

func allTasksSuccessOrSkipped(tasks []*entity.TaskInstance) bool {
	return len(tasks) > 0 && lo.EveryBy(tasks, func(task *entity.TaskInstance) bool {
		return task.Status == entity.TaskInstanceStatusSuccess ||
			task.Status == entity.TaskInstanceStatusSkipped
	})
}

func collectWatchDogRetryTargetIDs(tasks []*entity.TaskInstance) []string {
	return lo.FilterMap(tasks, func(task *entity.TaskInstance, _ int) (string, bool) {
		return task.ID, task.Status == entity.TaskInstanceStatusFailed && !IsTerminalReason(task.Reason)
	})
}

func collectRetryCommandTargetIDs(targetIDs []string, tasks []*entity.TaskInstance) []string {
	taskByID := make(map[string]*entity.TaskInstance, len(tasks))
	for _, task := range tasks {
		taskByID[task.ID] = task
	}

	retryTaskIDs := make([]string, 0, len(targetIDs))
	for _, targetID := range targetIDs {
		task, ok := taskByID[targetID]
		if !ok {
			continue
		}
		if (task.Status == entity.TaskInstanceStatusFailed || task.Status == entity.TaskInstanceStatusCanceled) &&
			!IsTerminalReason(task.Reason) {
			retryTaskIDs = append(retryTaskIDs, task.ID)
		}
	}
	return retryTaskIDs
}

func collectTaskIDsByStatus(tasks []*entity.TaskInstance, statuses ...entity.TaskInstanceStatus) []string {
	if len(tasks) == 0 || len(statuses) == 0 {
		return nil
	}
	statusSet := make(map[entity.TaskInstanceStatus]struct{}, len(statuses))
	for _, status := range statuses {
		statusSet[status] = struct{}{}
	}

	ret := make([]string, 0, len(tasks))
	for _, task := range tasks {
		if task == nil || task.ID == "" {
			continue
		}
		if _, ok := statusSet[task.Status]; ok {
			ret = append(ret, task.ID)
		}
	}
	return ret
}

func uniqueTaskIDs(ids []string) []string {
	if len(ids) == 0 {
		return nil
	}
	seen := map[string]struct{}{}
	ret := make([]string, 0, len(ids))
	for _, id := range ids {
		if id == "" {
			continue
		}
		if _, ok := seen[id]; ok {
			continue
		}
		seen[id] = struct{}{}
		ret = append(ret, id)
	}
	return ret
}

func collectRepairableTaskIDs(tasks []*entity.TaskInstance) []string {
	return lo.FilterMap(tasks, func(task *entity.TaskInstance, _ int) (string, bool) {
		return task.ID, slices.Contains([]entity.TaskInstanceStatus{
			entity.TaskInstanceStatusInit,
			entity.TaskInstanceStatusRunning,
			entity.TaskInstanceStatusEnding,
			entity.TaskInstanceStatusRetrying,
		}, task.Status) && !IsTerminalReason(task.Reason)
	})
}

func (wd *dispatcherWatchDog) cleanCompleteDagInstances(ctx context.Context) error {
	return wd.provider.Transaction(ctx, func(ctx context.Context) error {
		tx := wd.provider.WriteDB(ctx)

		var dagInstances []entity.DagInstance
		if err := tx.Model(&entity.DagInstance{}).
			Where("status = ? and created_at < ?", entity.DagInstanceStatusSuccess, time.Now().Second()-7*24*3600).
			Find(&dagInstances).Error; err != nil {
			return err
		}

		dagIdMap := make(map[string]any)
		lo.ForEach(dagInstances, func(item entity.DagInstance, index int) {
			if strings.HasPrefix(item.DagID, DynamicDagPrefix) {
				tilog.Info(ctx, "delete dynamic dag", zap.String("dag", item.DagID))
				dagIdMap[item.DagID] = struct{}{}
			}
		})

		res := tx.Model(&entity.DagInstance{}).
			Where("status = ? and created_at < ?", entity.DagInstanceStatusSuccess, time.Now().Second()-7*24*3600).
			Delete(&entity.DagInstance{})
		if res.Error != nil {
			return res.Error
		}
		MetricsWatchDogCleanDagInsTotal.Add(float64(res.RowsAffected))

		return tx.Model(&entity.Dag{}).Where("id in ?", lo.Keys(dagIdMap)).Delete(&entity.Dag{}).Error
	})
}

func (wd *watchDog) cleanCompleteDagInstances(ctx context.Context) error {
	if wd.mode == ModeCompetitive {
		return wd.competitive.cleanCompleteDagInstances(ctx)
	}
	return wd.dispatcher.cleanCompleteDagInstances(ctx)
}