geo/server/internal/tenant/app/kol_variable_renderer.go

package app

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"unicode/utf8"
)

type KolVariableType string

const (
	KolVariableTypeInput    KolVariableType = "input"
	KolVariableTypeTextarea KolVariableType = "textarea"
	KolVariableTypeSelect   KolVariableType = "select"
	KolVariableTypeNumber   KolVariableType = "number"
	KolVariableTypeCheckbox KolVariableType = "checkbox"

	kolDefaultTextMaxLength = 20
	kolDefaultNumberMin     = 0
	kolDefaultNumberMax     = 10000
)

type KolVariableDefinition struct {
	ID            string          `json:"id"`
	Key           string          `json:"key"`
	Label         string          `json:"label"`
	Type          KolVariableType `json:"type"`
	Required      bool            `json:"required"`
	Placeholder   *string         `json:"placeholder,omitempty"`
	Options       []string        `json:"options,omitempty"`
	MaxLength     *int            `json:"max_length,omitempty"`
	DefaultValue  *string         `json:"default_value,omitempty"`
	MinValue      *float64        `json:"min_value,omitempty"`
	MaxValue      *float64        `json:"max_value,omitempty"`
	DefaultNumber *float64        `json:"default_number,omitempty"`
}

type KolSchemaJSON struct {
	Variables []KolVariableDefinition `json:"variables"`
}

var kolPlaceholderRE = regexp.MustCompile(`\{\{\s*([^{}]+?)\s*\}\}`)

func RenderPrompt(content string, schema KolSchemaJSON, values map[string]any) (string, error) {
	schema = normalizeKolSchema(schema)

	byToken := make(map[string]KolVariableDefinition, len(schema.Variables)*2)
	for _, variable := range schema.Variables {
		id := strings.TrimSpace(variable.ID)
		key := strings.TrimSpace(variable.Key)
		if id != "" {
			byToken[id] = variable
		}
		if key != "" {
			byToken[key] = variable
		}
	}

	missingSet := make(map[string]struct{})
	missing := make([]string, 0)
	addMissing := func(id string) {
		if id == "" {
			return
		}
		if _, exists := missingSet[id]; exists {
			return
		}
		missingSet[id] = struct{}{}
		missing = append(missing, id)
	}

	var resolveErr error

	rendered := kolPlaceholderRE.ReplaceAllStringFunc(content, func(match string) string {
		if resolveErr != nil {
			return match
		}

		submatches := kolPlaceholderRE.FindStringSubmatch(match)
		if len(submatches) < 2 {
			return match
		}

		token := strings.TrimSpace(submatches[1])
		variable, exists := byToken[token]
		if !exists {
			addMissing(token)
			return match
		}

		value, hasValue, err := resolveKolVariableValue(values, variable, token)
		if err != nil {
			resolveErr = err
			return match
		}
		if hasValue {
			return fmt.Sprint(value)
		}
		if variable.Required {
			addMissing(kolVariableDisplayName(variable))
		}
		return ""
	})

	if resolveErr != nil {
		return "", resolveErr
	}

	if len(missing) > 0 {
		return "", fmt.Errorf("missing variables: %s", strings.Join(missing, ", "))
	}

	return rendered, nil
}

func HashPrompt(rendered string) string {
	sum := sha256.Sum256([]byte(rendered))
	return hex.EncodeToString(sum[:])
}

func ValidateSchema(schema KolSchemaJSON) error {
	schema = normalizeKolSchema(schema)

	seenIDs := make(map[string]struct{}, len(schema.Variables))
	seenKeys := make(map[string]struct{}, len(schema.Variables))
	for _, variable := range schema.Variables {
		id := strings.TrimSpace(variable.ID)
		if id == "" || !strings.HasPrefix(id, "var_") {
			return fmt.Errorf("variable id must be non-empty and start with var_")
		}
		if _, exists := seenIDs[id]; exists {
			return fmt.Errorf("duplicate variable id: %s", id)
		}
		seenIDs[id] = struct{}{}

		key := strings.TrimSpace(variable.Key)
		if key == "" {
			return fmt.Errorf("variable key must be non-empty")
		}
		if _, exists := seenKeys[key]; exists {
			return fmt.Errorf("duplicate variable key: %s", key)
		}
		seenKeys[key] = struct{}{}

		switch KolVariableType(strings.TrimSpace(string(variable.Type))) {
		case KolVariableTypeInput, KolVariableTypeTextarea, KolVariableTypeSelect, KolVariableTypeNumber, KolVariableTypeCheckbox:
		default:
			return fmt.Errorf("unsupported variable type: %s", variable.Type)
		}

		switch variable.Type {
		case KolVariableTypeInput, KolVariableTypeTextarea:
			maxLength := kolVariableTextMaxLength(variable)
			if maxLength <= 0 {
				return fmt.Errorf("%s max length must be greater than 0", kolVariableDisplayName(variable))
			}
			if variable.DefaultValue != nil && utf8.RuneCountInString(*variable.DefaultValue) > maxLength {
				return fmt.Errorf("%s default value exceeds max length %d", kolVariableDisplayName(variable), maxLength)
			}
		case KolVariableTypeNumber:
			minValue := kolVariableNumberMin(variable)
			maxValue := kolVariableNumberMax(variable)
			if maxValue < minValue {
				return fmt.Errorf("%s max value must be greater than or equal to min value", kolVariableDisplayName(variable))
			}
			if variable.DefaultNumber != nil && (*variable.DefaultNumber < minValue || *variable.DefaultNumber > maxValue) {
				return fmt.Errorf("%s default number must be between %s and %s", kolVariableDisplayName(variable), formatKolNumber(minValue), formatKolNumber(maxValue))
			}
		}
	}

	return nil
}

func JSONEncodeSchema(schema KolSchemaJSON) ([]byte, error) {
	return json.Marshal(normalizeKolSchema(schema))
}

func normalizeKolSchema(schema KolSchemaJSON) KolSchemaJSON {
	if schema.Variables == nil {
		schema.Variables = []KolVariableDefinition{}
		return schema
	}

	normalizedVariables := make([]KolVariableDefinition, 0, len(schema.Variables))
	for _, variable := range schema.Variables {
		normalizedVariables = append(normalizedVariables, normalizeKolVariableDefinition(variable))
	}
	schema.Variables = normalizedVariables
	return schema
}

func normalizeKolVariableDefinition(variable KolVariableDefinition) KolVariableDefinition {
	variable.ID = strings.TrimSpace(variable.ID)
	variable.Key = strings.TrimSpace(variable.Key)
	variable.Label = strings.TrimSpace(variable.Label)
	variable.Type = normalizeKolVariableType(variable.Type)

	switch variable.Type {
	case KolVariableTypeInput, KolVariableTypeTextarea:
		maxLength := kolVariableTextMaxLength(variable)
		variable.MaxLength = intPtr(maxLength)
		variable.Options = nil
		variable.MinValue = nil
		variable.MaxValue = nil
		variable.DefaultNumber = nil
	case KolVariableTypeNumber:
		minValue := kolVariableNumberMin(variable)
		maxValue := kolVariableNumberMax(variable)
		variable.MinValue = float64Ptr(minValue)
		variable.MaxValue = float64Ptr(maxValue)
		variable.Options = nil
		variable.MaxLength = nil
		variable.DefaultValue = nil
	case KolVariableTypeSelect, KolVariableTypeCheckbox:
		variable.MaxLength = nil
		variable.DefaultValue = nil
		variable.MinValue = nil
		variable.MaxValue = nil
		variable.DefaultNumber = nil
	default:
		variable.Type = KolVariableTypeInput
		maxLength := kolVariableTextMaxLength(variable)
		variable.MaxLength = intPtr(maxLength)
		variable.Options = nil
		variable.MinValue = nil
		variable.MaxValue = nil
		variable.DefaultNumber = nil
	}

	return variable
}

func normalizeKolVariableType(value KolVariableType) KolVariableType {
	switch KolVariableType(strings.TrimSpace(string(value))) {
	case KolVariableTypeTextarea:
		return KolVariableTypeTextarea
	case KolVariableTypeSelect:
		return KolVariableTypeSelect
	case KolVariableTypeNumber:
		return KolVariableTypeNumber
	case KolVariableTypeCheckbox:
		return KolVariableTypeCheckbox
	case KolVariableTypeInput:
		fallthrough
	default:
		return KolVariableTypeInput
	}
}

func resolveKolVariableValue(values map[string]any, variable KolVariableDefinition, token string) (any, bool, error) {
	candidates := []string{
		strings.TrimSpace(token),
		strings.TrimSpace(variable.Key),
		strings.TrimSpace(variable.ID),
	}

	seen := make(map[string]struct{}, len(candidates))
	for _, candidate := range candidates {
		if candidate == "" {
			continue
		}
		if _, exists := seen[candidate]; exists {
			continue
		}
		seen[candidate] = struct{}{}

		if value, ok := values[candidate]; ok {
			return normalizeKolRuntimeValue(variable, value)
		}
	}

	return kolVariableDefaultValue(variable)
}

func normalizeKolRuntimeValue(variable KolVariableDefinition, value any) (any, bool, error) {
	switch variable.Type {
	case KolVariableTypeInput, KolVariableTypeTextarea:
		text := fmt.Sprint(value)
		if strings.TrimSpace(text) == "" {
			return "", false, nil
		}
		maxLength := kolVariableTextMaxLength(variable)
		if utf8.RuneCountInString(text) > maxLength {
			return nil, false, fmt.Errorf("%s exceeds max length %d", kolVariableDisplayName(variable), maxLength)
		}
		return text, true, nil
	case KolVariableTypeNumber:
		minValue := kolVariableNumberMin(variable)
		maxValue := kolVariableNumberMax(variable)
		if maxValue < minValue {
			return nil, false, fmt.Errorf("invalid number range for %s", kolVariableDisplayName(variable))
		}
		number, ok := parseKolNumber(value)
		if !ok {
			return nil, false, fmt.Errorf("%s must be a valid number", kolVariableDisplayName(variable))
		}
		if number < minValue || number > maxValue {
			return nil, false, fmt.Errorf("%s must be between %s and %s", kolVariableDisplayName(variable), formatKolNumber(minValue), formatKolNumber(maxValue))
		}
		return number, true, nil
	case KolVariableTypeCheckbox:
		switch typed := value.(type) {
		case nil:
			return nil, false, nil
		case []string:
			if len(typed) == 0 {
				return typed, false, nil
			}
			return typed, true, nil
		case []any:
			if len(typed) == 0 {
				return typed, false, nil
			}
			return typed, true, nil
		default:
			text := strings.TrimSpace(fmt.Sprint(value))
			if text == "" {
				return "", false, nil
			}
			return value, true, nil
		}
	case KolVariableTypeSelect:
		fallthrough
	default:
		text := fmt.Sprint(value)
		if strings.TrimSpace(text) == "" {
			return "", false, nil
		}
		return text, true, nil
	}
}

func kolVariableDefaultValue(variable KolVariableDefinition) (any, bool, error) {
	switch variable.Type {
	case KolVariableTypeInput, KolVariableTypeTextarea:
		if variable.DefaultValue == nil {
			return nil, false, nil
		}
		return normalizeKolRuntimeValue(variable, *variable.DefaultValue)
	case KolVariableTypeNumber:
		if variable.DefaultNumber == nil {
			return nil, false, nil
		}
		return normalizeKolRuntimeValue(variable, *variable.DefaultNumber)
	default:
		return nil, false, nil
	}
}

func kolVariableTextMaxLength(variable KolVariableDefinition) int {
	if variable.MaxLength == nil || *variable.MaxLength <= 0 {
		return kolDefaultTextMaxLength
	}
	return *variable.MaxLength
}

func kolVariableNumberMin(variable KolVariableDefinition) float64 {
	if variable.MinValue == nil {
		return kolDefaultNumberMin
	}
	return *variable.MinValue
}

func kolVariableNumberMax(variable KolVariableDefinition) float64 {
	if variable.MaxValue == nil {
		return kolDefaultNumberMax
	}
	return *variable.MaxValue
}

func parseKolNumber(value any) (float64, bool) {
	switch typed := value.(type) {
	case float64:
		return typed, true
	case float32:
		return float64(typed), true
	case int:
		return float64(typed), true
	case int8:
		return float64(typed), true
	case int16:
		return float64(typed), true
	case int32:
		return float64(typed), true
	case int64:
		return float64(typed), true
	case uint:
		return float64(typed), true
	case uint8:
		return float64(typed), true
	case uint16:
		return float64(typed), true
	case uint32:
		return float64(typed), true
	case uint64:
		return float64(typed), true
	case json.Number:
		parsed, err := typed.Float64()
		return parsed, err == nil
	case string:
		trimmed := strings.TrimSpace(typed)
		if trimmed == "" {
			return 0, false
		}
		parsed, err := strconv.ParseFloat(trimmed, 64)
		return parsed, err == nil
	default:
		return 0, false
	}
}

func formatKolNumber(value float64) string {
	return strconv.FormatFloat(value, 'f', -1, 64)
}

func intPtr(value int) *int {
	return &value
}

func float64Ptr(value float64) *float64 {
	return &value
}

func kolVariableDisplayName(variable KolVariableDefinition) string {
	if label := strings.TrimSpace(variable.Label); label != "" {
		return label
	}
	if key := strings.TrimSpace(variable.Key); key != "" {
		return key
	}
	return strings.TrimSpace(variable.ID)
}