membank/auto-check/pkg/workflow/workflow.go

package workflow

import (
	"context"
	"fmt"
	"time"

	"auto-check/pkg/config"
	"auto-check/pkg/stress"
)

// ============================
// Context — 阶段间共享数据
// ============================

// Context 工作流上下文，保存阶段间传递的数据
type Context struct {
	base context.Context // 可取消上下文

	// 配置（由入口设置）
	Config config.Config

	// 循环模式下的状态数据（key 均为设备 MAC）
	Devices map[string]*Device          // 设备列表（MAC → 设备）
	Results map[string]*stress.Report  // 测试结果（nil/失败=需测试）
	Printed map[string]string          // 上次打印的测试状态
}

// NewContext 创建上下文
func NewContext() *Context {
	return &Context{
		base:    context.Background(),
		Devices: make(map[string]*Device),
		Results: make(map[string]*stress.Report),
		Printed: make(map[string]string),
	}
}

// WithCancel 支持取消
func (c *Context) WithCancel() (context.CancelFunc, error) {
	ctx, cancel := context.WithCancel(c.base)
	c.base = ctx
	return cancel, nil
}

// Done 返回取消信号
func (c *Context) Done() <-chan struct{} {
	return c.base.Done()
}

// Err 返回取消原因
func (c *Context) Err() error {
	return c.base.Err()
}

// ============================
// Stage — 阶段接口
// ============================

// Stage 工作流阶段
type Stage interface {
	Name() string
	Execute(ctx *Context) error
}

// ============================
// Workflow — 流水线调度器
// ============================

// OnStageError 阶段失败处理策略
type OnStageError int

const (
	// StopOnError 失败即中断整个工作流
	StopOnError OnStageError = iota
	// ContinueOnError 失败继续执行后续阶段
	ContinueOnError
)

// StageResult 阶段执行结果
type StageResult struct {
	Stage    string
	Start    time.Time
	End      time.Time
	Duration time.Duration
	Err      error
}

// Workflow Pipeline 工作流
type Workflow struct {
	stages    []Stage
	onError   OnStageError
	results   []StageResult
	startTime time.Time
}

// New 创建空工作流
func New() *Workflow {
	return &Workflow{
		onError: StopOnError,
	}
}

// SetOnError 设置失败策略
func (w *Workflow) SetOnError(policy OnStageError) *Workflow {
	w.onError = policy
	return w
}

// AddStage 追加阶段（串行）
func (w *Workflow) AddStage(stage Stage) *Workflow {
	w.stages = append(w.stages, stage)
	return w
}

// Run 依次执行所有阶段
func (w *Workflow) Run(ctx *Context) error {
	w.startTime = time.Now()
	fmt.Printf("══ 工作流启动 (%d 个阶段) ══\n\n", len(w.stages))

	for i, stage := range w.stages {
		// 检查是否被取消
		select {
		case <-ctx.Done():
			fmt.Printf("  [!] 工作流已取消，停止于阶段 %d/%d\n", i+1, len(w.stages))
			return ctx.Err()
		default:
		}

		fmt.Printf("════ 阶段 %d/%d: %s ════\n", i+1, len(w.stages), stage.Name())
		start := time.Now()
		err := stage.Execute(ctx)
		duration := time.Since(start)

		result := StageResult{
			Stage:    stage.Name(),
			Start:    start,
			End:      time.Now(),
			Duration: duration,
			Err:      err,
		}
		w.results = append(w.results, result)

		if err != nil {
			fmt.Printf("  [✗] 阶段失败: %v\n", err)
			switch w.onError {
			case StopOnError:
				fmt.Printf("══ 工作流中止（策略: 失败即停止）══\n")
				return err
			case ContinueOnError:
				fmt.Printf("  [!] 继续执行后续阶段（策略: 失败继续）\n\n")
				continue
			}
		} else {
			fmt.Printf("  [✓] 阶段完成 (%v)\n\n", duration.Round(time.Millisecond))
		}
	}

	fmt.Printf("══ 工作流完成，总耗时 %v ══\n", time.Since(w.startTime).Round(time.Millisecond))
	return nil
}

// Results 阶段执行结果列表
func (w *Workflow) Results() []StageResult {
	return w.results
}