atomctl/pkg/ast/provider/renderer.go

package provider

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"text/template"
	"time"
)

// Renderer defines the interface for rendering provider code
type Renderer interface {
	// Render renders providers to Go code
	Render(providers []Provider) ([]byte, error)

	// RenderToFile renders providers to a file
	RenderToFile(providers []Provider, filePath string) error

	// RenderToWriter renders providers to an io.Writer
	RenderToWriter(providers []Provider, writer io.Writer) error

	// AddTemplate adds a custom template
	AddTemplate(name, content string) error

	// RemoveTemplate removes a custom template
	RemoveTemplate(name string)

	// GetTemplate returns a template by name
	GetTemplate(name string) (*template.Template, error)

	// SetTemplateFuncs sets custom template functions
	SetTemplateFuncs(funcs template.FuncMap)
}

// GoRenderer implements the Renderer interface for Go code generation
type GoRenderer struct {
	templates       map[string]*template.Template
	templateFuncs   template.FuncMap
	outputConfig    *OutputConfig
	customTemplates map[string]string
}

// OutputConfig represents configuration for output generation
type OutputConfig struct {
	Header       string            // Header comment for generated files
	PackageName  string            // Package name for generated code
	Imports      map[string]string // Additional imports to include
	GeneratedTag string            // Tag to mark generated code
	DateFormat   string            // Date format for timestamps
	TemplateDir  string            // Directory for custom templates
	IndentString string            // String used for indentation
	LineEnding   string            // Line ending style ("\n" or "\r\n")
}

// RenderContext represents the context for rendering
type RenderContext struct {
	Providers   []Provider
	Config      *OutputConfig
	Timestamp   time.Time
	PackageName string
	Imports     map[string]string
	CustomData  map[string]interface{}
}

// NewGoRenderer creates a new GoRenderer with default configuration
func NewGoRenderer() *GoRenderer {
	return &GoRenderer{
		templates:       make(map[string]*template.Template),
		templateFuncs:   defaultTemplateFuncs(),
		outputConfig:    NewOutputConfig(),
		customTemplates: make(map[string]string),
	}
}

// NewGoRendererWithConfig creates a new GoRenderer with custom configuration
func NewGoRendererWithConfig(config *OutputConfig) *GoRenderer {
	if config == nil {
		config = NewOutputConfig()
	}

	return &GoRenderer{
		templates:       make(map[string]*template.Template),
		templateFuncs:   defaultTemplateFuncs(),
		outputConfig:    config,
		customTemplates: make(map[string]string),
	}
}

// NewOutputConfig creates a new OutputConfig with default values
func NewOutputConfig() *OutputConfig {
	return &OutputConfig{
		Header:       "// Code generated by atomctl provider generator. DO NOT EDIT.",
		PackageName:  "main",
		Imports:      make(map[string]string),
		GeneratedTag: "go:generate",
		DateFormat:   "2006-01-02 15:04:05",
		IndentString: "\t",
		LineEnding:   "\n",
	}
}

// Render implements Renderer.Render
func (r *GoRenderer) Render(providers []Provider) ([]byte, error) {
	var buf bytes.Buffer

	// Create render context
	context := r.createRenderContext(providers)

	// Render the main template
	tmpl, err := r.getOrCreateTemplate("provider", defaultProviderTemplate)
	if err != nil {
		return nil, fmt.Errorf("failed to get provider template: %w", err)
	}

	if err := tmpl.Execute(&buf, context); err != nil {
		return nil, fmt.Errorf("failed to execute template: %w", err)
	}

	return buf.Bytes(), nil
}

// RenderToFile implements Renderer.RenderToFile
func (r *GoRenderer) RenderToFile(providers []Provider, filePath string) error {
	// Create directory if it doesn't exist
	if err := os.MkdirAll(filepath.Dir(filePath), 0o755); err != nil {
		return fmt.Errorf("failed to create directory: %w", err)
	}

	// Create file
	file, err := os.Create(filePath)
	if err != nil {
		return fmt.Errorf("failed to create file: %w", err)
	}
	defer file.Close()

	// Render to file
	return r.RenderToWriter(providers, file)
}

// RenderToWriter implements Renderer.RenderToWriter
func (r *GoRenderer) RenderToWriter(providers []Provider, writer io.Writer) error {
	content, err := r.Render(providers)
	if err != nil {
		return err
	}

	_, err = writer.Write(content)
	return err
}

// AddTemplate implements Renderer.AddTemplate
func (r *GoRenderer) AddTemplate(name, content string) error {
	tmpl, err := template.New(name).Funcs(r.templateFuncs).Parse(content)
	if err != nil {
		return fmt.Errorf("failed to parse template %s: %w", name, err)
	}

	r.templates[name] = tmpl
	r.customTemplates[name] = content
	return nil
}

// RemoveTemplate implements Renderer.RemoveTemplate
func (r *GoRenderer) RemoveTemplate(name string) {
	delete(r.templates, name)
	delete(r.customTemplates, name)
}

// GetTemplate implements Renderer.GetTemplate
func (r *GoRenderer) GetTemplate(name string) (*template.Template, error) {
	return r.getOrCreateTemplate(name, "")
}

// SetTemplateFuncs implements Renderer.SetTemplateFuncs
func (r *GoRenderer) SetTemplateFuncs(funcs template.FuncMap) {
	r.templateFuncs = funcs

	// Re-compile all templates with new functions
	for name, content := range r.customTemplates {
		tmpl, err := template.New(name).Funcs(r.templateFuncs).Parse(content)
		if err != nil {
			continue // Keep the old template if compilation fails
		}
		r.templates[name] = tmpl
	}
}

// Helper methods

func (r *GoRenderer) createRenderContext(providers []Provider) *RenderContext {
	context := &RenderContext{
		Providers:   providers,
		Config:      r.outputConfig,
		Timestamp:   time.Now(),
		PackageName: r.outputConfig.PackageName,
		Imports:     make(map[string]string),
		CustomData:  make(map[string]interface{}),
	}

	// Collect all imports from providers
	for _, provider := range providers {
		for alias, path := range provider.Imports {
			context.Imports[path] = alias
		}
	}

	// Add custom imports
	for alias, path := range r.outputConfig.Imports {
		context.Imports[path] = alias
	}

	return context
}

func (r *GoRenderer) getOrCreateTemplate(name, defaultContent string) (*template.Template, error) {
	if tmpl, exists := r.templates[name]; exists {
		return tmpl, nil
	}

	if defaultContent == "" {
		return nil, fmt.Errorf("template %s not found", name)
	}

	tmpl, err := template.New(name).Funcs(r.templateFuncs).Parse(defaultContent)
	if err != nil {
		return nil, fmt.Errorf("failed to parse default template: %w", err)
	}

	r.templates[name] = tmpl
	return tmpl, nil
}

func defaultTemplateFuncs() template.FuncMap {
	return template.FuncMap{
		"toUpper":    strings.ToUpper,
		"toLower":    strings.ToLower,
		"toTitle":    strings.Title,
		"trimPrefix": strings.TrimPrefix,
		"trimSuffix": strings.TrimSuffix,
		"hasPrefix":  strings.HasPrefix,
		"hasSuffix":  strings.HasSuffix,
		"contains":   strings.Contains,
		"replace":    strings.Replace,
		"join":       strings.Join,
		"split":      strings.Split,
		"formatTime": formatTime,
		"quote":      func(s string) string { return fmt.Sprintf("%q", s) },
		"add":        func(a, b int) int { return a + b },
		"sub":        func(a, b int) int { return a - b },
		"mul":        func(a, b int) int { return a * b },
		"div":        func(a, b int) int { return a / b },
		"dict": func(values ...interface{}) (map[string]interface{}, error) {
			if len(values)%2 != 0 {
				return nil, fmt.Errorf("invalid dict call")
			}
			dict := make(map[string]interface{})
			for i := 0; i < len(values); i += 2 {
				key, ok := values[i].(string)
				if !ok {
					return nil, fmt.Errorf("dict keys must be strings")
				}
				dict[key] = values[i+1]
			}
			return dict, nil
		},
	}
}

func formatTime(t time.Time, format string) string {
	if format == "" {
		format = "2006-01-02 15:04:05"
	}
	return t.Format(format)
}

// Default provider template
const defaultProviderTemplate = `{{.Config.Header}}

// Generated at: {{.Timestamp.Format "2006-01-02 15:04:05"}}
// Package: {{.PackageName}}

package {{.PackageName}}

import (
	{{range $path, $alias := .Imports}}"{{$path}}" {{if $alias}}"{{$alias}}"{{end}}
	{{end}}
)

{{range $provider := .Providers}}
// {{.StructName}} provider implementation
// Mode: {{.Mode}}
// Return Type: {{.ReturnType}}
{{if .NeedPrepareFunc}}func (p *{{.StructName}}) Prepare() error {
	// Prepare logic for {{.StructName}}
	return nil
}{{end}}

func New{{.StructName}}({{range $name, $param := .InjectParams}}{{$name}} {{if $param.Star}}*{{end}}{{$param.Type}}{{if ne $name (last $provider.InjectParams)}}, {{end}}{{end}}) {{.ReturnType}} {
	return &{{.StructName}}{
		{{range $name, $param := .InjectParams}}{{$name}}: {{$name}},
		{{end}}
	}
}

{{end}}
`

// Utility functions for template rendering
func last(m map[string]InjectParam) string {
	if len(m) == 0 {
		return ""
	}
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys[len(keys)-1]
}