package flow

import (
	"context"
	"sync"
	"time"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/google/uuid"
)

// 运行中执行跟踪：优雅关停时取消全部执行（含脱离连接的恢复执行）并等待落完终态再退出，
// 避免"程序停止 → exec 仍卡在 status=1"。
//
// 背景：WS 执行（handleExecute）的 execCtx 派生自连接 closeCtx，Close() 取消连接即随之中止落终态；
// 但恢复执行（recoverExecution）的 ctx 经 context.WithoutCancel 脱离连接，程序关停时若不显式取消，
// 恢复中的 exec 不会落终态（仍 status=1），重启要等心跳陈旧 60s 才能再捞。这里登记所有运行中执行，
// SetShuttingDown 统一取消、main 等待全部落库后再退出。
var (
	execRunMu sync.Mutex
	execRuns  = make(map[string]context.CancelFunc)
)

// trackExecRun 登记一次运行中执行，返回 finish 在落完终态（recordWorkflow 后）调用。
// 每次运行唯一 key（重试/断点续跑复用同一 execId 也不冲突）。
func trackExecRun(cancel context.CancelFunc) (finish func()) {
	key := uuid.NewString()
	execRunMu.Lock()
	execRuns[key] = cancel
	execRunMu.Unlock()
	return func() {
		execRunMu.Lock()
		delete(execRuns, key)
		execRunMu.Unlock()
	}
}

// cancelAllExecRuns 优雅关停：取消所有运行中执行。恢复执行 ctx 经 WithoutCancel 脱离连接，
// 必须显式取消才能随 WS 执行一起落终态（status=3/retryable=1）。
func cancelAllExecRuns() {
	execRunMu.Lock()
	cancels := make([]context.CancelFunc, 0, len(execRuns))
	for _, c := range execRuns {
		cancels = append(cancels, c)
	}
	execRunMu.Unlock()
	for _, c := range cancels {
		c()
	}
}

// WaitExecRunsDrain 等待所有运行中执行落完终态（限时），供 main 优雅关停收尾后退出进程。
// 关停后不再启动新执行（execute/reExecute/recoverExecution 顶部有 IsShuttingDown 守卫），
// 因此运行中集合只减不增，轮询安全。
func WaitExecRunsDrain(timeout time.Duration) {
	deadline := time.Now().Add(timeout)
	for {
		execRunMu.Lock()
		n := len(execRuns)
		execRunMu.Unlock()
		if n == 0 {
			return
		}
		if time.Now().After(deadline) {
			g.Log().Warningf(context.Background(), "优雅关停等待执行落库超时，剩余 %d 个执行", n)
			return
		}
		time.Sleep(100 * time.Millisecond)
	}
}