ImaGen/internal/backend/workflow_template.go

package backend

import (
	"embed"
	"encoding/json"
	"fmt"
	"io/fs"
	"maps"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

//go:embed workflows/*.json
var bundledWorkflows embed.FS

// placeholderRE matches a single-token placeholder like "${prompt}" — the
// whole string value must be the placeholder, leading/trailing whitespace
// allowed. This lets us preserve types (a numeric substitution becomes a
// JSON number, not a stringified one) instead of round-tripping through
// strings.Replace which would force everything into a string.
var placeholderRE = regexp.MustCompile(`^\s*\$\{([a-zA-Z][a-zA-Z0-9_]*)\}\s*$`)

// LoadWorkflowTemplate returns the parsed JSON for a workflow template.
// `name` is resolved in this order:
//
//  1. exact filesystem path that exists on disk (absolute or relative);
//  2. one of the bundled templates under internal/backend/workflows/
//     (with or without the .json suffix).
//
// The returned map is a fresh deep copy of the template; callers can mutate
// it freely.
func LoadWorkflowTemplate(name string) (map[string]any, error) {
	if name == "" {
		return nil, fmt.Errorf("workflow template name is empty")
	}
	raw, err := readWorkflowBytes(name)
	if err != nil {
		return nil, err
	}
	var wf map[string]any
	if err := json.Unmarshal(raw, &wf); err != nil {
		return nil, fmt.Errorf("workflow %s: parse: %w", name, err)
	}
	return wf, nil
}

// BundledWorkflowNames returns the names of templates compiled into the
// binary, sorted. Each name is the basename without the .json suffix.
func BundledWorkflowNames() []string {
	entries, err := fs.ReadDir(bundledWorkflows, "workflows")
	if err != nil {
		return nil
	}
	out := make([]string, 0, len(entries))
	for _, e := range entries {
		n := e.Name()
		if !strings.HasSuffix(n, ".json") {
			continue
		}
		out = append(out, strings.TrimSuffix(n, ".json"))
	}
	sort.Strings(out)
	return out
}

func readWorkflowBytes(name string) ([]byte, error) {
	// Filesystem path wins if it points at a real file. Lets a user override
	// a bundled template by passing an absolute path in yaml.
	if strings.ContainsRune(name, os.PathSeparator) || strings.HasSuffix(name, ".json") {
		if b, err := os.ReadFile(name); err == nil {
			return b, nil
		} else if !os.IsNotExist(err) {
			return nil, fmt.Errorf("workflow %s: %w", name, err)
		}
	}
	// Bundled lookup. Try the literal name as a file inside workflows/, then
	// with the .json suffix appended.
	candidates := []string{
		filepath.Join("workflows", name),
		filepath.Join("workflows", name+".json"),
	}
	for _, c := range candidates {
		if b, err := bundledWorkflows.ReadFile(c); err == nil {
			return b, nil
		}
	}
	return nil, fmt.Errorf("workflow %q not found (bundled templates: %v)", name, BundledWorkflowNames())
}

// SubstituteWorkflow walks wf and replaces every "${key}" string with the
// matching value from subs, preserving JSON types. Returns the set of
// placeholder keys it actually touched, so the caller can detect missing
// substitutions even when a key is defined in subs but never referenced in
// the workflow (typical when a yaml block sets a knob a different template
// would consume).
//
// Unknown placeholders (referenced in the workflow but absent from subs)
// produce an error so we never submit a workflow with raw "${foo}" tokens.
func SubstituteWorkflow(wf map[string]any, subs map[string]any) (used map[string]struct{}, err error) {
	used = make(map[string]struct{})
	walked, err := substituteValue(wf, subs, used)
	if err != nil {
		return nil, err
	}
	// substituteValue returns the replacement for the top-level value, which
	// should still be the same map (just with mutated children).
	if m, ok := walked.(map[string]any); ok {
		// Copy back into wf so the caller's reference reflects the result.
		for k := range wf {
			delete(wf, k)
		}
		maps.Copy(wf, m)
	}
	return used, nil
}

func substituteValue(v any, subs map[string]any, used map[string]struct{}) (any, error) {
	switch x := v.(type) {
	case map[string]any:
		out := make(map[string]any, len(x))
		for k, child := range x {
			replaced, err := substituteValue(child, subs, used)
			if err != nil {
				return nil, err
			}
			out[k] = replaced
		}
		return out, nil
	case []any:
		out := make([]any, len(x))
		for i, child := range x {
			replaced, err := substituteValue(child, subs, used)
			if err != nil {
				return nil, err
			}
			out[i] = replaced
		}
		return out, nil
	case string:
		if m := placeholderRE.FindStringSubmatch(x); m != nil {
			key := m[1]
			val, ok := subs[key]
			if !ok {
				return nil, fmt.Errorf("workflow placeholder ${%s} has no substitution", key)
			}
			used[key] = struct{}{}
			return val, nil
		}
		return x, nil
	default:
		return v, nil
	}
}