enum.go

package main

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

	"gopkg.in/yaml.v3"
)

// Config represents the YAML configuration structure
type Config struct {
	Package string               `yaml:"package"`
	Output  string               `yaml:"output"`
	Enums   map[string]yaml.Node `yaml:"enums"`
}

// Enum defines the structure for enum values
type Enum struct {
	IsMap    bool
	Elements []string
	KeyValue map[string]string
}

func (e *Enum) UnmarshalYAML(value *yaml.Node) error {
	switch value.Kind {
	case yaml.SequenceNode:
		e.IsMap = false
		e.Elements = make([]string, 0)
		for _, n := range value.Content {
			e.Elements = append(e.Elements, n.Value)
		}
	case yaml.MappingNode:
		e.IsMap = true
		e.KeyValue = make(map[string]string)
		for i := 0; i < len(value.Content); i += 2 {
			key := value.Content[i].Value
			val := value.Content[i+1].Value
			e.KeyValue[key] = val
		}
	default:
		return fmt.Errorf("unsupported enum type: %v", value.Kind)
	}
	return nil
}

// ValueInfo contains metadata for generated values
type ValueInfo struct {
	TypeName  string
	ConstName string
	Value     string
	Methods   map[string]bool
}

// InterfaceData holds interface generation information
type InterfaceData struct {
	Name   string
	Method string
}

// TemplateData contains data for code generation template
type TemplateData struct {
	PackageName string
	Interfaces  []InterfaceData
	Types       []string
	Constants   []ConstantData
	Methods     []MethodData
}

// ConstantData defines constant generation parameters
type ConstantData struct {
	Name  string
	Type  string
	Value string
}

// MethodData contains method implementation details
type MethodData struct {
	Type   string
	Method string
}

const templateContent = `package {{.PackageName}}

// Code generated by generator; DO NOT EDIT.

{{range .Types}}
type {{.}} string
{{end}}

{{range .Types}}
func (v {{.}}) String() string {
	return string(v)
}
{{end}}

const (
{{- range $i, $c := .Constants}}
	{{if $i}}{{"\n	"}}{{end}}{{$c.Name}} {{$c.Type}} = {{printf "%q" $c.Value}}
{{- end}}
)

{{range .Interfaces}}
type {{.Name}} interface {
	{{.Method}}()
	String() string
}
{{end}}

{{range .Methods}}
func ({{.Type}}) {{.Method}}() {}
{{end}}
`

// convertToExportedName formats string to PascalCase
func convertToExportedName(s string) string {
	return strings.ReplaceAll(strings.Title(strings.ReplaceAll(s, "-", " ")), " ", "")
}

// generateMethodName creates method names from field names
func generateMethodName(field string) string {
	return strings.Join(strings.Split(strings.ToLower(field), "-"), "")
}

// createInterfaceName generates interface names
func createInterfaceName(field string) string {
	return convertToExportedName(field) + "Enum"
}

// parseYAMLContent handles different YAML node types
func parseYAMLContent(node *yaml.Node) (Enum, error) {
	result := Enum{}

	switch node.Kind {
	case yaml.SequenceNode:
		result.IsMap = false
		for _, n := range node.Content {
			result.Elements = append(result.Elements, n.Value)
		}
	case yaml.MappingNode:
		result.IsMap = true
		result.KeyValue = make(map[string]string)
		for i := 0; i < len(node.Content); i += 2 {
			key := node.Content[i].Value
			value := node.Content[i+1].Value
			result.KeyValue[key] = value
		}
	default:
		return result, fmt.Errorf("unsupported node type: %v", node.Kind)
	}
	return result, nil
}

func main() {
	// Load configuration
	configFile := "enums.yaml"
	if _, err := os.Stat(configFile); os.IsNotExist(err) {
		fmt.Printf("Error: Configuration file %s not found\n", configFile)
		os.Exit(1)
	}

	data, err := os.ReadFile(configFile)
	if err != nil {
		fmt.Printf("Error reading config file: %v\n", err)
		os.Exit(1)
	}

	var config Config
	if err := yaml.Unmarshal(data, &config); err != nil {
		fmt.Printf("Error parsing YAML: %v\n", err)
		os.Exit(1)
	}

	enums := make(map[string]Enum)
	for name, node := range config.Enums {
		var e Enum
		if err := node.Decode(&e); err != nil {
			fmt.Printf("Error decoding enum '%s': %v\n", name, err)
			os.Exit(1)
		}
		enums[name] = e
	}

	if config.Package == "" {
		config.Package = "enums"
	}
	if config.Output == "" {
		config.Output = "generated_enums.go"
	}

	// Create output directory if needed
	outputDir := filepath.Dir(config.Output)
	if outputDir != "." {
		if err := os.MkdirAll(outputDir, 0755); err != nil {
			fmt.Printf("Error creating output directory: %v\n", err)
			os.Exit(1)
		}
	}

	interfaceSet := make(map[string]InterfaceData)
	valueMap := make(map[string]*ValueInfo)

	for field, enum := range enums {
		methodName := generateMethodName(field)
		interfaceName := createInterfaceName(field)

		interfaceSet[interfaceName] = InterfaceData{
			Name:   interfaceName,
			Method: methodName,
		}

		processEnumValues(field, enum, methodName, valueMap)
	}

	// Prepare template data
	templateData := prepareTemplateData(config.Package, interfaceSet, valueMap)

	// Generate code
	if err := generateCode(config.Output, templateData); err != nil {
		fmt.Printf("Code generation failed: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Successfully generated: %s\n", config.Output)
}

func processEnumValues(field string, enum Enum, methodName string, valueMap map[string]*ValueInfo) {
	if enum.IsMap {
		for key, value := range enum.KeyValue {
			baseName := convertToExportedName(key)
			typeName := strings.ToLower(baseName[:1]) + baseName[1:] + "Case"
			valueKey := strings.ToLower(key)

			if _, exists := valueMap[valueKey]; !exists {
				valueMap[valueKey] = &ValueInfo{
					TypeName:  typeName,
					ConstName: baseName,
					Value:     value,
					Methods:   make(map[string]bool),
				}
			}
			valueMap[valueKey].Methods[methodName] = true
		}
	} else {
		for _, v := range enum.Elements {
			baseName := convertToExportedName(v)
			typeName := strings.ToLower(baseName[:1]) + baseName[1:] + "Case"
			valueKey := strings.ToLower(v)

			if _, exists := valueMap[valueKey]; !exists {
				valueMap[valueKey] = &ValueInfo{
					TypeName:  typeName,
					ConstName: baseName,
					Value:     strings.ToLower(v),
					Methods:   make(map[string]bool),
				}
			}
			valueMap[valueKey].Methods[methodName] = true
		}
	}
}

func prepareTemplateData(pkgName string, interfaceSet map[string]InterfaceData, valueMap map[string]*ValueInfo) TemplateData {
	// Sort interfaces
	var interfaces []InterfaceData
	for _, v := range interfaceSet {
		interfaces = append(interfaces, v)
	}
	sort.Slice(interfaces, func(i, j int) bool {
		return interfaces[i].Name < interfaces[j].Name
	})

	// Prepare template data
	data := TemplateData{
		PackageName: pkgName,
		Interfaces:  interfaces,
	}

	// Process values and types
	var valueKeys []string
	for k := range valueMap {
		valueKeys = append(valueKeys, k)
	}
	sort.Strings(valueKeys)

	seenTypes := make(map[string]bool)
	for _, valueKey := range valueKeys {
		info := valueMap[valueKey]

		// Add type (deduplicated)
		if !seenTypes[info.TypeName] {
			data.Types = append(data.Types, info.TypeName)
			seenTypes[info.TypeName] = true
		}

		// Add constant
		data.Constants = append(data.Constants, ConstantData{
			Name:  info.ConstName,
			Type:  info.TypeName,
			Value: info.Value,
		})

		// Add methods
		var methods []string
		for m := range info.Methods {
			methods = append(methods, m)
		}
		sort.Strings(methods)
		for _, m := range methods {
			data.Methods = append(data.Methods, MethodData{
				Type:   info.TypeName,
				Method: m,
			})
		}
	}

	return data
}

func generateCode(outputPath string, data TemplateData) error {
	tmpl, err := template.New("gen").Parse(templateContent)
	if err != nil {
		return err
	}

	f, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer f.Close()

	return tmpl.Execute(f, data)
}