template_utils.go

package exec

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"sync"
	"text/template"
	"text/template/parse"
	"time"

	"github.com/Masterminds/sprig/v3"
	"github.com/go-viper/mapstructure/v2"
	"github.com/hairyhenderson/gomplate/v3"
	"github.com/hairyhenderson/gomplate/v3/data"
	_ "github.com/hairyhenderson/gomplate/v4"
	"github.com/samber/lo"

	log "github.com/cloudposse/atmos/pkg/logger"
	"github.com/cloudposse/atmos/pkg/merge"
	"github.com/cloudposse/atmos/pkg/perf"
	"github.com/cloudposse/atmos/pkg/schema"
	u "github.com/cloudposse/atmos/pkg/utils"
)

// Cache for sprig function maps to avoid repeated expensive allocations.
// Sprig function maps are immutable once created, so caching is safe.
// Note: Gomplate functions are NOT cached as they may have state and context dependencies.
var (
	sprigFuncMapCache     template.FuncMap
	sprigFuncMapCacheOnce sync.Once
)

const (
	logKeyTemplate = "template"
)

// getSprigFuncMap returns a cached copy of the sprig hermetic function map.
// Sprig function maps are expensive to create (173MB+ allocations) and immutable,
// so we cache and reuse them across template operations.
// This optimization reduces heap allocations by ~3.76% (173MB) per profile run.
//
// We use HermeticTxtFuncMap instead of FuncMap to exclude non-pure side-effectful
// functions like env/expandenv/getHostByName from the Sprig base. The env and
// expandenv functions are instead provided explicitly via getEnvFuncMap, which
// is a deliberate design decision: stack templates legitimately need env var access
// (e.g. {{ env "GITHUB_TOKEN" }} in vendor.yaml, {{ env "USER" }} in stack manifests).
// For untrusted-template contexts (Aqua registry, asset URL templates) use only
// HermeticTxtFuncMap without getEnvFuncMap.
func getSprigFuncMap() template.FuncMap {
	sprigFuncMapCacheOnce.Do(func() {
		sprigFuncMapCache = sprig.HermeticTxtFuncMap()
	})
	return sprigFuncMapCache
}

// getEnvFuncMap returns a FuncMap with env and expandenv functions.
// These are provided separately from getSprigFuncMap so that:
//  1. The Sprig base (HermeticTxtFuncMap) stays free of other OS/network side-effects.
//  2. Stack templates that legitimately need {{ env "KEY" }} and {{ expandenv "$KEY" }}
//     get those functions via an explicit deliberate provision rather than inheriting
//     them from the full Sprig FuncMap.
//  3. Untrusted-template contexts (asset URL templates, Aqua registry) can safely
//     use only HermeticTxtFuncMap without these functions.
//
// The env and expandenv functions shadow Gomplate's "env" namespace object, which
// takes 0 arguments and is not compatible with {{ env "KEY" }} syntax used in Atmos.
func getEnvFuncMap() template.FuncMap {
	return template.FuncMap{
		"env":       os.Getenv,
		"expandenv": os.ExpandEnv,
	}
}

// ProcessTmpl parses and executes Go templates.
func ProcessTmpl(
	atmosConfig *schema.AtmosConfiguration,
	tmplName string,
	tmplValue string,
	tmplData any,
	ignoreMissingTemplateValues bool,
) (string, error) {
	defer perf.Track(atmosConfig, "exec.ProcessTmpl")()

	d := data.Data{}
	ctx := context.TODO()

	// Add Gomplate, Sprig and Atmos template functions.
	cfg := atmosConfig
	if cfg == nil {
		cfg = &schema.AtmosConfiguration{}
	}
	funcs := lo.Assign(gomplate.CreateFuncs(ctx, &d), getSprigFuncMap(), getEnvFuncMap(), FuncMap(cfg, &schema.ConfigAndStacksInfo{}, ctx, &d))

	t, err := template.New(tmplName).Funcs(funcs).Parse(tmplValue)
	if err != nil {
		return "", err
	}

	// Control the behavior during execution if a map is indexed with a key that is not present in the map.
	// If the template context (`tmplData`) does not provide all the required variables, the following errors would be thrown:
	// template: catalog/terraform/eks_cluster_tmpl_hierarchical.yaml:17:12: executing "catalog/terraform/eks_cluster_tmpl_hierarchical.yaml" at <.flavor>: map has no entry for key "flavor"
	// template: catalog/terraform/eks_cluster_tmpl_hierarchical.yaml:12:36: executing "catalog/terraform/eks_cluster_tmpl_hierarchical.yaml" at <.stage>: map has no entry for key "stage"

	option := "missingkey=error"

	if ignoreMissingTemplateValues {
		option = "missingkey=default"
	}

	t.Option(option)

	var res bytes.Buffer
	err = t.Execute(&res, tmplData)
	if err != nil {
		return "", err
	}

	return res.String(), nil
}

// convertRawEnvToStringMap converts a raw env value (any) to map[string]string.
// Handles map[string]any (with non-string values skipped) and map[string]string.
// Returns nil if the input is nil, an unsupported type, or produces an empty map.
func convertRawEnvToStringMap(envRaw any) map[string]string {
	if envRaw == nil {
		return nil
	}

	result := make(map[string]string)

	switch env := envRaw.(type) {
	case map[string]any:
		for k, v := range env {
			if s, ok := v.(string); ok {
				result[k] = s
			}
		}
	case map[string]string:
		for k, v := range env {
			result[k] = v
		}
	}

	if len(result) == 0 {
		return nil
	}

	return result
}

// extractEnvFromRawMap extracts the env map from a raw settings map[string]any.
// This is needed because mapstructure:"-" on TemplatesSettings.Env causes
// mapstructure.Decode to silently drop the env field.
// The path navigated is: templates -> settings -> env.
func extractEnvFromRawMap(settingsMap map[string]any) map[string]string {
	templates, ok := settingsMap["templates"].(map[string]any)
	if !ok {
		return nil
	}

	settings, ok := templates["settings"].(map[string]any)
	if !ok {
		return nil
	}

	envRaw, ok := settings["env"]
	if !ok {
		return nil
	}

	return convertRawEnvToStringMap(envRaw)
}

// savedEnvVar stores the original state of an environment variable for restore.
type savedEnvVar struct {
	value   string
	existed bool
}

// setEnvVarsWithRestore sets environment variables and returns a cleanup function
// that restores the original values. This prevents env pollution across components.
func setEnvVarsWithRestore(envVars map[string]string) (func(), error) {
	saved := make(map[string]savedEnvVar, len(envVars))

	for k := range envVars {
		if original, existed := os.LookupEnv(k); existed {
			saved[k] = savedEnvVar{value: original, existed: true}
		} else {
			saved[k] = savedEnvVar{existed: false}
		}
	}

	cleanup := func() {
		for k, original := range saved {
			if original.existed {
				os.Setenv(k, original.value)
			} else {
				os.Unsetenv(k)
			}
		}
	}

	for k, v := range envVars {
		if err := os.Setenv(k, v); err != nil {
			// Return cleanup for vars already set before the failure.
			return cleanup, err
		}
	}

	return cleanup, nil
}

// ProcessTmplWithDatasources parses and executes Go templates with datasources.
func ProcessTmplWithDatasources(
	atmosConfig *schema.AtmosConfiguration,
	configAndStacksInfo *schema.ConfigAndStacksInfo,
	settingsSection schema.Settings,
	tmplName string,
	tmplValue string,
	tmplData any,
	ignoreMissingTemplateValues bool,
) (string, error) {
	defer perf.Track(atmosConfig, "exec.ProcessTmplWithDatasources")()

	if !atmosConfig.Templates.Settings.Enabled {
		log.Debug("ProcessTmplWithDatasources: not processing templates since templating is disabled in 'atmos.yaml'", logKeyTemplate, tmplName)
		return tmplValue, nil
	}

	log.Trace("ProcessTmplWithDatasources", logKeyTemplate, tmplName)

	// Preserve env vars before mapstructure drops them.
	// mapstructure:"-" on TemplatesSettings.Env causes Decode to silently drop the field.
	// We extract env from both sources directly before the encode/decode/merge pipeline.
	cliEnv := atmosConfig.Templates.Settings.Env
	stackEnv := settingsSection.Templates.Settings.Env

	// Merge the template settings from `atmos.yaml` CLI config and from the stack manifests
	var cliConfigTemplateSettingsMap map[string]any
	var stackManifestTemplateSettingsMap map[string]any
	var templateSettings schema.TemplatesSettings

	err := mapstructure.Decode(atmosConfig.Templates.Settings, &cliConfigTemplateSettingsMap)
	if err != nil {
		return "", err
	}

	err = mapstructure.Decode(settingsSection.Templates.Settings, &stackManifestTemplateSettingsMap)
	if err != nil {
		return "", err
	}

	templateSettingsMerged, err := merge.Merge(atmosConfig, []map[string]any{cliConfigTemplateSettingsMap, stackManifestTemplateSettingsMap})
	if err != nil {
		return "", err
	}

	err = mapstructure.Decode(templateSettingsMerged, &templateSettings)
	if err != nil {
		return "", err
	}

	// Restore env vars that mapstructure dropped.
	// Stack manifest env overrides CLI config env (same precedence as the merge above).
	mergedEnv := make(map[string]string)
	for k, v := range cliEnv {
		mergedEnv[k] = v
	}
	for k, v := range stackEnv {
		mergedEnv[k] = v
	}
	if len(mergedEnv) > 0 {
		templateSettings.Env = mergedEnv
	}

	// Add Atmos, Gomplate and Sprig functions and datasources
	funcs := make(map[string]any)

	// Number of processing evaluations/passes
	evaluations, _ := lo.Coalesce(atmosConfig.Templates.Settings.Evaluations, 1)
	result := tmplValue

	// Set environment variables for template processing before the loop.
	// Restore originals when the function returns.
	if len(templateSettings.Env) > 0 {
		cleanup, envErr := setEnvVarsWithRestore(templateSettings.Env)
		if envErr != nil {
			return "", envErr
		}
		defer cleanup()
	}

	// Track extra env keys introduced by the evaluation loop that weren't in the
	// initial set, so we can clean them up when the function returns.
	extraLoopKeys := make(map[string]struct{})
	defer func() {
		for k := range extraLoopKeys {
			os.Unsetenv(k)
		}
	}()

	for i := 0; i < evaluations; i++ {
		log.Trace("ProcessTmplWithDatasources", logKeyTemplate, tmplName, "evaluation", i+1)

		d := data.Data{}

		// Gomplate functions and datasources
		if atmosConfig.Templates.Settings.Gomplate.Enabled {
			// If timeout is not provided in `atmos.yaml` nor in `settings.templates.settings` stack manifest, use 5 seconds
			timeoutSeconds, _ := lo.Coalesce(templateSettings.Gomplate.Timeout, 5)

			ctx, cancelFunc := context.WithTimeout(context.TODO(), time.Second*time.Duration(timeoutSeconds))
			defer cancelFunc()

			d = data.Data{}
			d.Ctx = ctx

			for k, v := range templateSettings.Gomplate.Datasources {
				_, err := d.DefineDatasource(k, v.Url)
				if err != nil {
					return "", err
				}

				// Add datasource headers
				if len(v.Headers) > 0 {
					d.Sources[k].Header = v.Headers
				}
			}

			funcs = lo.Assign(funcs, gomplate.CreateFuncs(ctx, &d))
		}

		// Sprig functions
		if atmosConfig.Templates.Settings.Sprig.Enabled {
			funcs = lo.Assign(funcs, getSprigFuncMap())
			// Explicitly add env/expandenv since HermeticTxtFuncMap excludes them
			// but they are a documented, intentional feature of Atmos stack templates.
			funcs = lo.Assign(funcs, getEnvFuncMap())
		}

		// Atmos functions
		funcs = lo.Assign(funcs, FuncMap(atmosConfig, configAndStacksInfo, context.TODO(), &d))

		// Process the template
		t := template.New(tmplName).Funcs(funcs)

		// Template delimiters
		leftDelimiter := "{{"
		rightDelimiter := "}}"

		if len(atmosConfig.Templates.Settings.Delimiters) > 0 {
			delimiterError := fmt.Errorf("invalid 'templates.settings.delimiters' config in 'atmos.yaml': %v\n"+
				"'delimiters' must be an array with two string items: left and right delimiter\n"+
				"the left and right delimiters must not be an empty string", atmosConfig.Templates.Settings.Delimiters)

			if len(atmosConfig.Templates.Settings.Delimiters) != 2 {
				return "", delimiterError
			}

			if atmosConfig.Templates.Settings.Delimiters[0] == "" {
				return "", delimiterError
			}

			if atmosConfig.Templates.Settings.Delimiters[1] == "" {
				return "", delimiterError
			}

			leftDelimiter = atmosConfig.Templates.Settings.Delimiters[0]
			rightDelimiter = atmosConfig.Templates.Settings.Delimiters[1]
		}

		t.Delims(leftDelimiter, rightDelimiter)

		// Control the behavior during execution if a map is indexed with a key that is not present in the map
		// If the template context (`tmplData`) does not provide all the required variables, the following errors would be thrown:
		// template: catalog/terraform/eks_cluster_tmpl_hierarchical.yaml:17:12: executing "catalog/terraform/eks_cluster_tmpl_hierarchical.yaml" at <.flavor>: map has no entry for key "flavor"
		// template: catalog/terraform/eks_cluster_tmpl_hierarchical.yaml:12:36: executing "catalog/terraform/eks_cluster_tmpl_hierarchical.yaml" at <.stage>: map has no entry for key "stage"

		option := "missingkey=error"

		if ignoreMissingTemplateValues {
			option = "missingkey=default"
		}

		t.Option(option)

		// Parse the template
		t, err = t.Parse(result)
		if err != nil {
			return "", err
		}

		// Execute the template
		var res bytes.Buffer
		err = t.Execute(&res, tmplData)
		if err != nil {
			return "", err
		}

		result = res.String()
		resultMap, err := u.UnmarshalYAML[schema.AtmosSectionMapType](result)
		if err != nil {
			return "", err
		}

		if resultMapSettings, ok := resultMap["settings"].(map[string]any); ok {
			if resultMapSettingsTemplates, ok := resultMapSettings["templates"].(map[string]any); ok {
				if resultMapSettingsTemplatesSettings, ok := resultMapSettingsTemplates["settings"].(map[string]any); ok {
					// Extract env before mapstructure drops it.
					resultEnv := convertRawEnvToStringMap(resultMapSettingsTemplatesSettings["env"])

					err = mapstructure.Decode(resultMapSettingsTemplatesSettings, &templateSettings)
					if err != nil {
						return "", err
					}

					// Restore env after mapstructure dropped it.
					// Also update OS env vars for the next evaluation pass.
					if len(resultEnv) > 0 {
						templateSettings.Env = resultEnv
						for k, v := range resultEnv {
							os.Setenv(k, v)
							// Track keys not in the initial set for deferred cleanup.
							if _, inInitial := mergedEnv[k]; !inInitial {
								extraLoopKeys[k] = struct{}{}
							}
						}
					}
				}
			}
		}
	}

	log.Trace("ProcessTmplWithDatasources: processed", logKeyTemplate, tmplName)

	return result, nil
}

// IsGolangTemplate checks if the provided string is a Go template.
func IsGolangTemplate(atmosConfig *schema.AtmosConfiguration, str string) (bool, error) {
	defer perf.Track(atmosConfig, "exec.IsGolangTemplate")()
	t, err := template.New(str).Parse(str)
	if err != nil {
		return false, err
	}

	isGoTemplate := false

	// Iterate over all nodes in the template and check if any of them is of type `NodeAction` (field evaluation)
	for _, node := range t.Root.Nodes {
		if node.Type() == parse.NodeAction {
			isGoTemplate = true
			break
		}
	}

	return isGoTemplate, nil
}

// Create temporary directory.
func createTempDirectory() (string, error) {
	// Create a temporary directory for the temporary files.
	tempDir, err := os.MkdirTemp("", "atmos-templates-*")
	if err != nil {
		return "", fmt.Errorf("failed to create temp directory: %w", err)
	}
	// Ensure directory permissions are restricted.
	if err := os.Chmod(tempDir, defaultDirPermissions); err != nil {
		return "", fmt.Errorf("failed to set temp directory permissions: %w", err)
	}
	return tempDir, nil
}

// Write merged JSON data to a temporary file and return its final file URL.
func writeMergedDataToFile(tempDir string, mergedData map[string]interface{}) (*url.URL, error) {
	// Write the merged JSON data to a file.
	rawJSON, err := json.Marshal(mergedData)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal merged data to JSON: %w", err)
	}

	// Create a temporary file inside the temp directory.
	tmpfile, err := os.CreateTemp(tempDir, "gomplate-data-*.json")
	if err != nil {
		return nil, fmt.Errorf("failed to create temp data file for gomplate: %w", err)
	}
	tmpName := tmpfile.Name()
	if _, err := tmpfile.Write(rawJSON); err != nil {
		tmpfile.Close()
		return nil, fmt.Errorf("failed to write JSON to temp file: %w", err)
	}
	if err := tmpfile.Close(); err != nil {
		return nil, fmt.Errorf("failed to close temp data file: %w", err)
	}

	fileURL := toFileScheme(tmpName)

	finalFileUrl, err := fixWindowsFileScheme(fileURL)
	if err != nil {
		return nil, err
	}
	return finalFileUrl, nil
}

// Write the 'outer' top-level file and return its final file URL.
func writeOuterTopLevelFile(tempDir string, fileURL string) (*url.URL, error) {
	// Write the 'outer' top-level file.
	topLevel := map[string]interface{}{
		"Env": map[string]interface{}{
			"README_YAML": fileURL,
		},
	}
	outerJSON, err := json.Marshal(topLevel)
	if err != nil {
		return nil, err
	}

	tmpfile2, err := os.CreateTemp(tempDir, "gomplate-top-level-*.json")
	if err != nil {
		return nil, fmt.Errorf("failed to create temp data file for top-level: %w", err)
	}
	tmpName2 := tmpfile2.Name()
	if _, err = tmpfile2.Write(outerJSON); err != nil {
		tmpfile2.Close()
		return nil, fmt.Errorf("failed to write top-level JSON: %w", err)
	}
	if err = tmpfile2.Close(); err != nil {
		return nil, fmt.Errorf("failed to close top-level JSON: %w", err)
	}

	topLevelFileURL := toFileScheme(tmpName2)

	finalTopLevelFileURL, err := fixWindowsFileScheme(topLevelFileURL)
	if err != nil {
		return nil, err
	}
	return finalTopLevelFileURL, nil
}

// ProcessTmplWithDatasourcesGomplate parses and executes Go templates with datasources using Gomplate.
func ProcessTmplWithDatasourcesGomplate(
	atmosConfig *schema.AtmosConfiguration,
	tmplName string,
	tmplValue string,
	mergedData map[string]interface{},
	ignoreMissingTemplateValues bool,
) (string, error) {
	defer perf.Track(atmosConfig, "exec.ProcessTmplWithDatasourcesGomplate")()

	tempDir, err := createTempDirectory()
	if err != nil {
		return "", err
	}
	defer os.RemoveAll(tempDir)

	finalFileUrl, err := writeMergedDataToFile(tempDir, mergedData)
	if err != nil {
		return "", err
	}

	finalTopLevelFileURL, err := writeOuterTopLevelFile(tempDir, finalFileUrl.String())
	if err != nil {
		return "", err
	}

	// Construct Gomplate Options.
	opts := gomplate.Options{
		Context: map[string]gomplate.Datasource{
			".": {
				URL: finalTopLevelFileURL,
			},
			"config": {
				URL: finalFileUrl,
			},
		},
		Funcs: template.FuncMap{},
	}

	// Render the template.
	renderer := gomplate.NewRenderer(opts)
	var buf bytes.Buffer
	tpl := gomplate.Template{
		Name:   tmplName,
		Text:   tmplValue,
		Writer: &buf,
	}

	if err := renderer.RenderTemplates(context.Background(), []gomplate.Template{tpl}); err != nil {
		return "", fmt.Errorf("failed to render template: %w", err)
	}

	return buf.String(), nil
}