tmpl/registry/prompt.go

package registry

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"text/template"

	"github.com/AlecAivazis/survey/v2"
	"github.com/pelletier/go-toml/v2"
)

type templatePrompt struct {
	Name      string   `toml:"-"`
	Value     any      `toml:"-"`
	Message   string   `toml:"prompt"`
	Help      string   `toml:"help"`
	Default   any      `toml:"default"`
	DependsOn []string `toml:"depends_on"`
}

func prompt(dir string, defaults bool) (templatePrompts, error) {
	templatePath := filepath.Join(dir, "template.toml")
	if _, err := os.Lstat(templatePath); err != nil {
		return nil, err
	}

	templateBytes, err := ioutil.ReadFile(templatePath)
	if err != nil {
		return nil, err
	}

	// Expand the template with environment variables
	templateContents := os.Expand(string(templateBytes), func(s string) string {
		if strings.HasPrefix(s, "TMPL_PROMPT") {
			return fmt.Sprintf("${%s}", s)
		}
		return os.Getenv(s)
	})

	var templateMap map[string]templatePrompt
	if err := toml.Unmarshal([]byte(templateContents), &templateMap); err != nil {
		return nil, err
	}

	prompts := sortMap(templateMap)
	for idx, prompt := range prompts {
		if prompt.Message == "" {
			prompt.Message = prompt.Name
		}
		if prompt.Default == nil {
			prompt.Default = ""
		}
		prompts[idx] = prompt
	}

	for idx, prompt := range prompts {
		// Check for env variable
		envKey := strings.ToUpper(prompt.Name)
		if e, ok := os.LookupEnv(fmt.Sprintf("TMPL_VAR_%s", envKey)); ok {
			prompts[idx].Value = e
			os.Setenv(fmt.Sprintf("TMPL_PROMPT_%s", envKey), e)
			continue
		}

		// Check if we are using defaults
		if defaults {
			val := prompt.Default
			switch t := prompt.Default.(type) {
			case []string:
				for idy, s := range t {
					t[idy] = os.ExpandEnv(s)
				}
				val = t
			case string:
				val = os.ExpandEnv(t)
			}
			s := fmt.Sprint(val)
			prompts[idx].Value = s
			os.Setenv(fmt.Sprintf("TMPL_PROMPT_%s", envKey), s)
			continue
		}

		var p survey.Prompt
		switch t := prompt.Default.(type) {
		case []string:
			for idy, s := range t {
				t[idy] = os.ExpandEnv(s)
			}
			p = &survey.Select{
				Message: prompt.Message,
				Options: t,
				Help:    prompt.Help,
			}
		case bool:
			p = &survey.Confirm{
				Message: prompt.Message,
				Default: t,
				Help:    prompt.Help,
			}
		case string:
			p = &survey.Input{
				Message: prompt.Message,
				Default: os.ExpandEnv(t),
				Help:    prompt.Help,
			}
		default:
			p = &survey.Input{
				Message: prompt.Message,
				Default: fmt.Sprint(t),
				Help:    prompt.Help,
			}
		}
		var a string
		if err := survey.AskOne(p, &a); err != nil {
			return nil, err
		}
		prompts[idx].Value = a
		os.Setenv(fmt.Sprintf("TMPL_PROMPT_%s", envKey), a)
	}

	return prompts, nil
}

func sortMap(m map[string]templatePrompt) templatePrompts {
	prompts := make(templatePrompts, 0, len(m))
	for name, prompt := range m {
		prompt.Name = name
		prompts = append(prompts, prompt)
	}
	sort.Slice(prompts, func(i, j int) bool {
		return prompts[i].Name < prompts[j].Name
	})
	sort.Slice(prompts, func(i, j int) bool {
		for _, dep := range prompts[i].DependsOn {
			if strings.EqualFold(dep, prompts[j].Name) {
				return false
			}
		}
		return true
	})
	fmt.Println(prompts)
	return prompts
}

type templatePrompts []templatePrompt

// ToMap converts a slice to templatePrompt into a suitable template context
func (t templatePrompts) ToMap() map[string]any {
	m := make(map[string]any)
	for _, p := range t {
		m[p.Name] = p.Value
	}
	return m
}

// ToFuncMap converts a slice of templatePrompt into a suitable template.FuncMap
func (t templatePrompts) ToFuncMap() template.FuncMap {
	m := make(map[string]any)
	for k, v := range t.ToMap() {
		vv := v // Enclosure
		m[k] = func() any {
			return vv
		}
	}
	return m
}