package workflow

import (
	"context"
	"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"
)

type competitiveWatchDog struct {
	provider db.TransProvider
}

func newCompetitiveWatchDog(provider db.TransProvider) *competitiveWatchDog {
	return &competitiveWatchDog{provider: provider}
}

func (wd *competitiveWatchDog) retry(_ context.Context) error {
	dagInsList, err := wd.listRetryCandidates()
	if err != nil {
		return err
	}

	for _, dagIns := range dagInsList {
		if dagIns == nil || dagIns.ID == "" {
			continue
		}
		if dagIns.Worker == "" {
			latest, latestErr := mod.GetStore().GetDagInstance(dagIns.ID)
			if latestErr != nil {
				continue
			}
			dagIns = latest
			if dagIns.Worker == "" {
				newVersion, claimed, claimErr := mod.GetStore().TryClaimDagIns(dagIns.ID, mod.GetKeeper().WorkerKey(), dagIns.Version)
				if claimErr != nil {
					continue
				}
				if claimed {
					dagIns.Worker = mod.GetKeeper().WorkerKey()
					dagIns.Version = newVersion
				}
			}
			if dagIns.Worker == "" {
				continue
			}
		}
		if dagIns.Worker != mod.GetKeeper().WorkerKey() {
			continue
		}

		tasks, listErr := mod.GetStore().ListTaskInstance(&mod.ListTaskInstanceInput{DagInsID: dagIns.ID})
		if listErr != nil || len(tasks) == 0 {
			continue
		}

		nowSec := time.Now().Unix()
		isStale := isFullStuck(dagIns, tasks, nowSec)

		if dagIns.Cmd != nil {
			if !isStale || dagIns.Cmd.Name != entity.CommandNameRetry {
				continue
			}
			retryTargets := collectRetryCommandTargetIDs(dagIns.Cmd.TargetTaskInsIDs, tasks)
			if len(retryTargets) == 0 {
				continue
			}
			if err := mod.GetStore().PatchDagIns(&entity.DagInstance{BaseInfo: entity.BaseInfo{ID: dagIns.ID}, Cmd: nil}, "Cmd"); err != nil {
				MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "true").Inc()
				continue
			}
			if err := wd.retryTargetsForDag(dagIns, retryTargets); err != nil {
				MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "true").Inc()
				continue
			}
			MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "false").Inc()
			continue
		}

		failedOrCanceled := collectTaskIDsByStatus(tasks, entity.TaskInstanceStatusFailed, entity.TaskInstanceStatusCanceled)
		taskByID := lo.KeyBy(tasks, func(task *entity.TaskInstance) string {
			if task == nil {
				return ""
			}
			return task.ID
		})
		failedOrCanceled = lo.Filter(failedOrCanceled, func(taskID string, _ int) bool {
			task, ok := taskByID[taskID]
			return !ok || !IsTerminalReason(task.Reason)
		})
		if len(failedOrCanceled) == 0 {
			failedOrCanceled = collectTaskIDsByStatus(tasks, entity.TaskInstanceStatusInit)
		}

		if dagIns.Status == entity.DagInstanceStatusFailed && len(failedOrCanceled) > 0 {
			if err := wd.retryTargetsForDag(dagIns, failedOrCanceled); err != nil {
				MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "true").Inc()
				continue
			}
			MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "false").Inc()
			continue
		}

		if dagIns.Status == entity.DagInstanceStatusFailed && allTasksSuccessOrSkipped(tasks) {
			if err := mod.GetStore().PatchDagIns(&entity.DagInstance{
				BaseInfo: entity.BaseInfo{ID: dagIns.ID},
				Status:   entity.DagInstanceStatusSuccess,
				Reason:   "",
			}, "Reason"); err != nil {
				MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "true").Inc()
				continue
			}
			MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "false").Inc()
			continue
		}

		if !isStale {
			continue
		}

		repairable := collectRepairableTaskIDs(tasks)
		for _, taskID := range repairable {
			_ = mod.GetStore().PatchTaskIns(&entity.TaskInstance{
				BaseInfo: entity.BaseInfo{ID: taskID},
				Status:   entity.TaskInstanceStatusFailed,
				Reason:   watchDogRepairReason,
			})
		}

		retryTargets := uniqueTaskIDs(append(failedOrCanceled, repairable...))
		if len(retryTargets) == 0 {
			continue
		}

		if dagIns.Status == entity.DagInstanceStatusRunning {
			if err := mod.GetStore().PatchDagIns(&entity.DagInstance{
				BaseInfo: entity.BaseInfo{ID: dagIns.ID},
				Status:   entity.DagInstanceStatusFailed,
				Reason:   watchDogRepairReason,
			}, "Reason"); err != nil {
				MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "true").Inc()
				continue
			}
		}

		if err := wd.retryTargetsForDag(dagIns, retryTargets); err != nil {
			MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "true").Inc()
			continue
		}
		MetricsWatchDogRetryDagInsTotal.WithLabelValues(dagIns.DagID, dagIns.ID, "false").Inc()
	}

	return nil
}

func (wd *competitiveWatchDog) listRetryCandidates() ([]*entity.DagInstance, error) {
	statusFilter := []entity.DagInstanceStatus{entity.DagInstanceStatusFailed, entity.DagInstanceStatusRunning}
	owned, err := mod.GetStore().ListDagInstance(&mod.ListDagInstanceInput{
		Worker: mod.GetKeeper().WorkerKey(),
		Status: statusFilter,
		Limit:  watchDogCandidateLimit,
	})
	if err != nil {
		return nil, err
	}

	remaining := watchDogCandidateLimit - len(owned)
	if remaining <= 0 {
		return owned, nil
	}

	emptyOwned, err := mod.GetStore().ListDagInstance(&mod.ListDagInstanceInput{
		WorkerEmpty: true,
		Status:      statusFilter,
		Limit:       int64(remaining),
	})
	if err != nil {
		return nil, err
	}

	ret := make([]*entity.DagInstance, 0, len(owned)+len(emptyOwned))
	ret = append(ret, owned...)
	ret = append(ret, emptyOwned...)
	return ret, nil
}

func (wd *competitiveWatchDog) retryTargetsForDag(dagIns *entity.DagInstance, retryTargets []string) error {
	if len(retryTargets) == 0 {
		return nil
	}

	for _, taskID := range retryTargets {
		if err := mod.GetStore().PatchTaskIns(&entity.TaskInstance{
			BaseInfo: entity.BaseInfo{ID: taskID},
			Status:   entity.TaskInstanceStatusRetrying,
			Reason:   "",
		}); err != nil {
			return err
		}
	}

	latest, err := mod.GetStore().GetDagInstance(dagIns.ID)
	if err != nil {
		return err
	}
	if latest.Worker == "" {
		newVersion, claimed, claimErr := mod.GetStore().TryClaimDagIns(latest.ID, mod.GetKeeper().WorkerKey(), latest.Version)
		if claimErr != nil {
			return claimErr
		}
		if claimed {
			latest.Worker = mod.GetKeeper().WorkerKey()
			latest.Version = newVersion
		}
	}
	if latest.Worker != mod.GetKeeper().WorkerKey() {
		return nil
	}

	latest.Run()
	if err := mod.GetStore().PatchDagIns(&entity.DagInstance{
		BaseInfo: entity.BaseInfo{ID: latest.ID},
		Status:   latest.Status,
		Reason:   "",
		Cmd:      nil,
	}, "Cmd", "Reason"); err != nil {
		return err
	}

	mod.GetParser().InitialDagIns(latest)
	return nil
}

func (wd *competitiveWatchDog) 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)
		for i := range dagInstances {
			if strings.HasPrefix(dagInstances[i].DagID, DynamicDagPrefix) {
				tilog.Info(ctx, "delete dynamic dag", zap.String("dag", dagInstances[i].DagID))
				dagIdMap[dagInstances[i].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
	})
}