attach.go

package auditlog

import (
	"context"
	"time"

	flow "github.com/Azure/go-workflow"
)

// Attach injects audit BeforeStep/AfterStep callbacks into every step in the
// workflow. Call this BEFORE w.Do(ctx).
//
// The callbacks are merged into each step's existing config via State.MergeConfig,
// so user-defined callbacks (Input, Output, BeforeStep, AfterStep) are preserved.
// Audit callbacks are appended last so they observe the final error.
//
// When the Auditor is disabled, Attach is a no-op.
func (a *Auditor) Attach(w *flow.Workflow) *flow.Workflow {
	if !a.config.Enabled || w == nil {
		return w
	}

	for _, step := range w.Steps() {
		state := w.StateOf(step)
		if state == nil {
			continue
		}

		state.MergeConfig(&flow.StepConfig{
			Before: []flow.BeforeStep{a.beforeFn},
			After:  []flow.AfterStep{a.afterFn},
		})
	}

	return w
}

// CaptureDAG traverses the workflow's step structure and pre-populates step
// records with names, types, dependencies, retry/timeout config, and a
// "pending" status. This makes the DAG available in Report() BEFORE w.Do(ctx)
// is called, enabling real-time dashboards to render the full step graph
// immediately on connect.
//
// Call this AFTER Attach(w) and BEFORE w.Do(ctx). Steps that already have
// records (from a prior CaptureDAG or from event capture) are not overwritten.
//
// When the Auditor is disabled, CaptureDAG is a no-op.
func (a *Auditor) CaptureDAG(w *flow.Workflow) {
	if !a.config.Enabled || w == nil {
		return
	}

	a.recorder.captureDAG(w)
}

// Snapshot reads the workflow's final state after Do() to capture the full DAG
// structure, final statuses, and any steps that were skipped or canceled
// (which bypass Before/After callbacks entirely).
//
// Call this AFTER w.Do(ctx) returns.
//
// When the Auditor is disabled, Snapshot is a no-op.
func (a *Auditor) Snapshot(w *flow.Workflow) {
	if !a.config.Enabled || w == nil {
		return
	}

	a.recorder.snapshotWorkflow(w)
}

// makeCallbacks creates the BeforeStep and AfterStep closures that feed the recorder.
func (r *Recorder) makeCallbacks() (flow.BeforeStep, flow.AfterStep) {
	before := func(ctx context.Context, step flow.Steper) (context.Context, error) {
		r.recordBeforeStep(step)

		return ctx, nil
	}
	after := func(_ context.Context, step flow.Steper, err error) error {
		r.recordAfterStep(step, err)

		return err
	}

	return before, after
}

// snapshotWorkflow reads step statuses, dependencies, and retry/timeout config
// from the workflow's post-execution state. Traverses sub-workflows to capture
// inner step statuses that bypass Before/After callbacks.
func (r *Recorder) snapshotWorkflow(w *flow.Workflow) {
	r.mu.Lock()
	defer r.mu.Unlock()

	for _, root := range w.Steps() {
		flow.Traverse(root, func(step flow.Steper, _ []flow.Steper) flow.TraverseDecision {
			// Skip wrapper steps that don't have their own state (NamedStep, etc).
			// Their underlying step is traversed separately.
			if w.StateOf(step) == nil {
				return flow.TraverseEndBranch
			}

			r.snapshotStepLocked(w, step)

			return flow.TraverseContinue
		})
	}
}

// snapshotStepLocked captures a single step's final state from the workflow.
// Caller must hold r.mu.
func (r *Recorder) snapshotStepLocked(w *flow.Workflow, step flow.Steper) {
	name := flow.String(step)

	state := w.StateOf(step)
	if state == nil {
		return
	}

	status := fromFlowStatus(string(state.GetStatus()))
	err := state.GetError()

	now := time.Now()
	rec := r.getOrCreateStepLocked(step, name, now)
	rec.status = status

	if err != nil {
		errStr := err.Error()
		rec.attemptErr = &errStr
	}

	// Capture retry and timeout configuration.
	if opt := state.Option(); opt != nil {
		if opt.RetryOption != nil {
			rec.hasRetry = true
			//nolint:gosec // Attempts is a small retry count, overflow is not realistic.
			rec.maxAttempts = int(opt.RetryOption.Attempts)
		}

		if opt.Timeout != nil {
			rec.hasTimeout = true
		}
	}

	// Capture dependencies (upstream steps).
	upstreams := w.UpstreamOf(step)
	deps := make([]StepRef, 0, len(upstreams))

	for up := range upstreams {
		deps = append(deps, StepRef{Name: flow.String(up), StepType: stepTypeName(up)})
	}

	sortByName(deps)
	rec.dependencies = deps
}

// captureDAG traverses the workflow to pre-populate step records with names,
// types, dependencies, retry/timeout config, and a "pending" status. This
// enables DAG visualization before execution begins.
func (r *Recorder) captureDAG(w *flow.Workflow) {
	r.mu.Lock()
	defer r.mu.Unlock()

	for _, root := range w.Steps() {
		flow.Traverse(root, func(step flow.Steper, _ []flow.Steper) flow.TraverseDecision {
			// Skip wrapper steps that don't have their own state (NamedStep, etc).
			if w.StateOf(step) == nil {
				return flow.TraverseEndBranch
			}

			r.captureDAGStepLocked(w, step)

			return flow.TraverseContinue
		})
	}
}

// captureDAGStepLocked pre-populates a step record with structural information
// (name, type, dependencies, retry/timeout config) and a "pending" status.
// If a record already exists (from event capture or a prior call), it is not
// overwritten — live execution data always wins.
// Caller must hold r.mu.
func (r *Recorder) captureDAGStepLocked(w *flow.Workflow, step flow.Steper) {
	// Skip if already created — execution data should not be overwritten.
	if _, exists := r.steps[step]; exists {
		return
	}

	name := flow.String(step)
	now := time.Now()

	rec := r.getOrCreateStepLocked(step, name, now)

	// Override the default "running" status to "pending" — execution hasn't started.
	rec.status = StepStatusPending

	// Capture dependencies from the workflow structure.
	upstreams := w.UpstreamOf(step)
	deps := make([]StepRef, 0, len(upstreams))

	for up := range upstreams {
		deps = append(deps, StepRef{Name: flow.String(up), StepType: stepTypeName(up)})
	}

	sortByName(deps)
	rec.dependencies = deps

	// Capture retry and timeout configuration.
	state := w.StateOf(step)
	if state == nil {
		return
	}

	opt := state.Option()
	if opt == nil {
		return
	}

	if opt.RetryOption != nil {
		rec.hasRetry = true

		//nolint:gosec // Attempts is a small retry count, overflow is not realistic.
		rec.maxAttempts = int(opt.RetryOption.Attempts)
	}

	if opt.Timeout != nil {
		rec.hasTimeout = true
	}
}