interpres/encode.go

package interpres

import (
	"bytes"
	"fmt"
	"math"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

var (
	localDateTimeType = reflect.TypeOf(LocalDateTime{})
	localDateType     = reflect.TypeOf(LocalDate{})
	localTimeType     = reflect.TypeOf(LocalTime{})
	timeGoType        = reflect.TypeOf(time.Time{})
)

// encoder produces a TOML document from a Go value via a small intermediate
// representation that preserves the order in which fields were declared.
type encoder struct {
	buf bytes.Buffer
}

func newEncoder() *encoder { return &encoder{} }

func (e *encoder) bytes() []byte { return e.buf.Bytes() }

// encode converts v into a TOML document. v must be a struct or a
// map[string]V (or a non-nil pointer to one).
func (e *encoder) encode(v any) error {
	rv := reflect.ValueOf(v)
	if !rv.IsValid() {
		return fmt.Errorf("interpres: cannot marshal nil value")
	}
	if rv.Kind() == reflect.Pointer {
		if rv.IsNil() {
			return fmt.Errorf("interpres: cannot marshal nil pointer")
		}
		rv = rv.Elem()
	}
	doc := &tomlDoc{}
	switch rv.Kind() {
	case reflect.Struct:
		if err := buildStructDoc(rv, doc, ""); err != nil {
			return err
		}
	case reflect.Map:
		if err := buildMapDoc(rv, doc, ""); err != nil {
			return err
		}
	default:
		return fmt.Errorf("interpres: top-level value must be a struct or map[string]V, got %s", rv.Type())
	}
	return e.emitDoc(doc, nil)
}

// --- intermediate representation -----------------------------------------

// tomlDoc holds the entries of one TOML table, partitioned by kind. Fields
// declared at the same TOML level are grouped: scalars come first, then
// sub-tables, then arrays of tables. The order WITHIN each group matches
// either struct field order (for structs) or sorted key order (for maps).
type tomlDoc struct {
	scalars []tomlKV
	tables  []tomlTable
	arrays  []tomlArray
}

type tomlKV struct {
	key string
	val any
}

type tomlTable struct {
	key string
	doc *tomlDoc
}

type tomlArray struct {
	key  string
	docs []*tomlDoc
}

func (d *tomlDoc) addScalar(key string, val any) {
	d.scalars = append(d.scalars, tomlKV{key: key, val: val})
}

func (d *tomlDoc) addTable(key string, sub *tomlDoc) {
	d.tables = append(d.tables, tomlTable{key: key, doc: sub})
}

func (d *tomlDoc) addArray(key string, subs []*tomlDoc) {
	d.arrays = append(d.arrays, tomlArray{key: key, docs: subs})
}

// --- reflection walk: struct ---------------------------------------------

func buildStructDoc(v reflect.Value, doc *tomlDoc, ctx string) error {
	t := v.Type()
	for i := 0; i < t.NumField(); i++ {
		f := t.Field(i)
		if f.PkgPath != "" {
			continue
		}
		if f.Anonymous {
			tag, _ := f.Tag.Lookup("toml")
			if tag == "-" {
				continue
			}
			if tag == "" {
				fv := followPtr(v.Field(i))
				if !fv.IsValid() {
					continue
				}
				switch fv.Kind() {
				case reflect.Struct:
					if isScalarStruct(fv.Type()) {
						name := strings.ToLower(f.Name)
						if err := doc.appendScalar(name, fv.Interface(), ctx); err != nil {
							return err
						}
						continue
					}
					if err := buildStructDoc(fv, doc, ctx); err != nil {
						return err
					}
					continue
				case reflect.Map:
					if err := buildMapDoc(fv, doc, ctx); err != nil {
						return err
					}
					continue
				}
			}
		}
		name := fieldName(f)
		if name == "-" {
			continue
		}
		if err := addField(doc, name, v.Field(i), ctx); err != nil {
			return err
		}
	}
	return nil
}

// fieldName returns the TOML key for a struct field, honouring the `toml`
// tag (name or `-`) and falling back to a lower-cased field name.
func fieldName(f reflect.StructField) string {
	if tag, ok := f.Tag.Lookup("toml"); ok {
		name := strings.Split(tag, ",")[0]
		if name == "-" {
			return "-"
		}
		if name != "" {
			return name
		}
	}
	return strings.ToLower(f.Name)
}

// --- reflection walk: map ------------------------------------------------

func buildMapDoc(v reflect.Value, doc *tomlDoc, ctx string) error {
	if v.Type().Key().Kind() != reflect.String {
		return fmt.Errorf("interpres: map key must be string, got %s", v.Type().Key())
	}
	keys := v.MapKeys()
	sort.Slice(keys, func(i, j int) bool { return keys[i].String() < keys[j].String() })
	for _, k := range keys {
		if err := addField(doc, k.String(), v.MapIndex(k), ctx); err != nil {
			return err
		}
	}
	return nil
}

// --- reflection walk: field dispatch -------------------------------------

func addField(doc *tomlDoc, name string, v reflect.Value, ctx string) error {
	if v.CanInterface() {
		if m, ok := v.Interface().(Marshaler); ok {
			mv, err := m.MarshalTOML()
			if err != nil {
				return fmt.Errorf("interpres: %s.%s: %w", ctx, name, err)
			}
			v = reflect.ValueOf(mv)
		}
	}
	v = followPtr(v)
	if !v.IsValid() {
		return nil
	}
	if v.Kind() == reflect.Interface {
		if v.IsNil() {
			return nil
		}
		v = v.Elem()
	}
	switch v.Kind() {
	case reflect.Struct:
		if isScalarStruct(v.Type()) {
			return doc.appendScalar(name, v.Interface(), ctx)
		}
		return addSubTable(doc, name, v, ctx)
	case reflect.Map:
		return addSubTable(doc, name, v, ctx)
	case reflect.Slice, reflect.Array:
		return addArrayValue(doc, name, v, ctx)
	default:
		val, err := normalizeValue(v)
		if err != nil {
			return fmt.Errorf("interpres: %s.%s: %w", ctx, name, err)
		}
		return doc.appendScalar(name, val, ctx)
	}
}

// appendScalar wraps addScalar with a uniform error path.
func (d *tomlDoc) appendScalar(name string, val any, ctx string) error {
	d.addScalar(name, val)
	return nil
}

func addSubTable(doc *tomlDoc, name string, v reflect.Value, ctx string) error {
	sub := &tomlDoc{}
	switch v.Kind() {
	case reflect.Struct:
		if err := buildStructDoc(v, sub, joinKey(ctx, name)); err != nil {
			return err
		}
	case reflect.Map:
		if err := buildMapDoc(v, sub, joinKey(ctx, name)); err != nil {
			return err
		}
	}
	doc.addTable(name, sub)
	return nil
}

func addArrayValue(doc *tomlDoc, name string, v reflect.Value, ctx string) error {
	if v.Kind() == reflect.Slice && v.IsNil() {
		// A nil slice has no explicit representation in TOML — skip.
		return nil
	}
	n := v.Len()
	if n == 0 {
		if isTableElementType(v.Type().Elem()) {
			// Empty array of tables has no valid TOML form — skip.
			return nil
		}
		return doc.appendScalar(name, []any{}, ctx)
	}

	if isTableElementValue(v.Index(0)) {
		subs := make([]*tomlDoc, n)
		for i := 0; i < n; i++ {
			ev := followPtr(v.Index(i))
			if !ev.IsValid() {
				return fmt.Errorf("interpres: %s.%s[%d]: nil element", ctx, name, i)
			}
			sub := &tomlDoc{}
			switch ev.Kind() {
			case reflect.Struct:
				if isScalarStruct(ev.Type()) {
					return fmt.Errorf("interpres: %s.%s[%d]: heterogeneous array contains scalar", ctx, name, i)
				}
				if err := buildStructDoc(ev, sub, joinKey(ctx, fmt.Sprintf("%s[%d]", name, i))); err != nil {
					return err
				}
			case reflect.Map:
				if err := buildMapDoc(ev, sub, joinKey(ctx, fmt.Sprintf("%s[%d]", name, i))); err != nil {
					return err
				}
			default:
				return fmt.Errorf("interpres: %s.%s: heterogeneous array, expected table", ctx, name)
			}
			subs[i] = sub
		}
		doc.addArray(name, subs)
		return nil
	}

	// Regular array of scalars.
	items := make([]any, n)
	for i := 0; i < n; i++ {
		ev := followPtr(v.Index(i))
		if !ev.IsValid() {
			return fmt.Errorf("interpres: %s.%s[%d]: nil element", ctx, name, i)
		}
		if ev.CanInterface() {
			if m, ok := ev.Interface().(Marshaler); ok {
				mv, err := m.MarshalTOML()
				if err != nil {
					return fmt.Errorf("interpres: %s.%s[%d]: %w", ctx, name, i, err)
				}
				ev = reflect.ValueOf(mv)
				ev = followPtr(ev)
			}
		}
		val, err := normalizeValue(ev)
		if err != nil {
			return fmt.Errorf("interpres: %s.%s[%d]: %w", ctx, name, i, err)
		}
		items[i] = val
	}
	return doc.appendScalar(name, items, ctx)
}

// normalizeValue converts a reflect.Value into one of the canonical scalar or
// nested-array representations the emitter understands. Slices and arrays are
// recursively normalised so that nested arrays (e.g. [][]int) work.
func normalizeValue(v reflect.Value) (any, error) {
	if v.CanInterface() {
		if m, ok := v.Interface().(Marshaler); ok {
			return m.MarshalTOML()
		}
	}
	switch v.Kind() {
	case reflect.String:
		return v.String(), nil
	case reflect.Bool:
		return v.Bool(), nil
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v.Int(), nil
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		u := v.Uint()
		if u > math.MaxInt64 {
			return nil, fmt.Errorf("unsigned value %d overflows int64", u)
		}
		return int64(u), nil
	case reflect.Float32, reflect.Float64:
		return v.Float(), nil
	case reflect.Slice, reflect.Array:
		items := make([]any, v.Len())
		for i := 0; i < v.Len(); i++ {
			val, err := normalizeValue(v.Index(i))
			if err != nil {
				return nil, fmt.Errorf("[%d]: %w", i, err)
			}
			items[i] = val
		}
		return items, nil
	}
	if !v.IsValid() {
		return nil, fmt.Errorf("invalid value")
	}
	return nil, fmt.Errorf("cannot encode %s", v.Type())
}

// followPtr unwraps pointer and interface layers. Returns a zero Value if a
// nil pointer or nil interface is encountered.
func followPtr(v reflect.Value) reflect.Value {
	for {
		switch v.Kind() {
		case reflect.Pointer, reflect.Interface:
			if v.IsNil() {
				return reflect.Value{}
			}
			v = v.Elem()
			continue
		}
		return v
	}
}

// isScalarStruct reports whether t is a struct type that the encoder treats
// as a TOML scalar (time.Time, LocalDateTime, LocalDate, LocalTime).
func isScalarStruct(t reflect.Type) bool {
	return t == timeGoType || isLocalDateType(t)
}

func isLocalDateType(t reflect.Type) bool {
	return t == localDateTimeType || t == localDateType || t == localTimeType
}

func isTableElementType(t reflect.Type) bool {
	switch t.Kind() {
	case reflect.Struct:
		return !isScalarStruct(t)
	case reflect.Map:
		return t.Key().Kind() == reflect.String
	}
	return false
}

func isTableElementValue(v reflect.Value) bool {
	v = followPtr(v)
	if !v.IsValid() {
		return false
	}
	return isTableElementType(v.Type())
}

func joinKey(ctx, name string) string {
	if ctx == "" {
		return name
	}
	return ctx + "." + name
}

// --- emission ------------------------------------------------------------

// writeHeaderSep writes a single newline before a table or array-of-tables
// header so the output has a blank line between sections, unless the buffer
// is empty (i.e. this is the very first header).
func (e *encoder) writeHeaderSep() {
	if e.buf.Len() == 0 {
		return
	}
	e.buf.WriteByte('\n')
}

func (e *encoder) emitDoc(doc *tomlDoc, prefix []string) error {
	for _, kv := range doc.scalars {
		if err := e.writeKV(kv.key, kv.val); err != nil {
			return err
		}
	}
	for _, t := range doc.tables {
		path := append(append([]string{}, prefix...), t.key)
		e.writeHeaderSep()
		e.buf.WriteByte('[')
		writeKeyPath(&e.buf, path)
		e.buf.WriteString("]\n")
		if err := e.emitDoc(t.doc, path); err != nil {
			return err
		}
	}
	for _, a := range doc.arrays {
		path := append(append([]string{}, prefix...), a.key)
		for _, sub := range a.docs {
			e.writeHeaderSep()
			e.buf.WriteString("[[")
			writeKeyPath(&e.buf, path)
			e.buf.WriteString("]]\n")
			if err := e.emitDoc(sub, path); err != nil {
				return err
			}
		}
	}
	return nil
}

func (e *encoder) writeKV(key string, val any) error {
	if !utf8.ValidString(key) {
		return fmt.Errorf("interpres: key %q is not valid UTF-8", key)
	}
	e.writeKey(key)
	e.buf.WriteString(" = ")
	if err := e.writeValue(val); err != nil {
		return err
	}
	e.buf.WriteByte('\n')
	return nil
}

func writeKeyPath(buf *bytes.Buffer, path []string) {
	for i, p := range path {
		if i > 0 {
			buf.WriteByte('.')
		}
		if isBareKey(p) {
			buf.WriteString(p)
			continue
		}
		writeQuotedString(buf, p)
	}
}

func (e *encoder) writeKey(key string) {
	if isBareKey(key) {
		e.buf.WriteString(key)
		return
	}
	writeQuotedString(&e.buf, key)
}

// writeQuotedString writes s as a TOML basic string (double-quoted) to buf.
// Returns an error only if s is not valid UTF-8; invalid byte sequences
// within a valid UTF-8 string are encoded as \ufffd replacement characters.
func writeQuotedString(buf *bytes.Buffer, s string) error {
	if !utf8.ValidString(s) {
		return fmt.Errorf("interpres: string is not valid UTF-8")
	}
	buf.WriteByte('"')
	for i := 0; i < len(s); {
		r, size := utf8.DecodeRuneInString(s[i:])
		if r == utf8.RuneError && size == 1 {
			buf.WriteString(`\ufffd`)
			i++
			continue
		}
		i += size
		writeEscapedRune(buf, r)
	}
	buf.WriteByte('"')
	return nil
}

// writeEscapedRune writes a single rune to buf, escaping it as required by
// TOML basic-string rules.
func writeEscapedRune(buf *bytes.Buffer, r rune) {
	switch r {
	case '\\':
		buf.WriteString(`\\`)
	case '"':
		buf.WriteString(`\"`)
	case '\b':
		buf.WriteString(`\b`)
	case '\t':
		buf.WriteString(`\t`)
	case '\n':
		buf.WriteString(`\n`)
	case '\f':
		buf.WriteString(`\f`)
	case '\r':
		buf.WriteString(`\r`)
	default:
		if r < 0x20 || r == 0x7f {
			fmt.Fprintf(buf, `\u%04X`, r)
		} else {
			buf.WriteRune(r)
		}
	}
}

func isBareKey(s string) bool {
	if s == "" {
		return false
	}
	for i := 0; i < len(s); i++ {
		c := s[i]
		if !((c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') || c == '_' || c == '-') {
			return false
		}
	}
	return true
}

func (e *encoder) writeValue(val any) error {
	switch v := val.(type) {
	case string:
		return e.writeStringVal(v)
	case bool:
		e.buf.WriteString(strconv.FormatBool(v))
		return nil
	case int64:
		e.buf.WriteString(strconv.FormatInt(v, 10))
		return nil
	case float64:
		return e.writeFloat(v)
	case time.Time:
		e.buf.WriteString(v.Format(time.RFC3339Nano))
		return nil
	case LocalDateTime:
		e.buf.WriteString(v.String())
		return nil
	case LocalDate:
		e.buf.WriteString(v.String())
		return nil
	case LocalTime:
		e.buf.WriteString(v.String())
		return nil
	case []any:
		e.buf.WriteByte('[')
		for i, item := range v {
			if i > 0 {
				e.buf.WriteString(", ")
			}
			if err := e.writeValue(item); err != nil {
				return err
			}
		}
		e.buf.WriteByte(']')
		return nil
	case nil:
		return fmt.Errorf("interpres: cannot encode nil value")
	default:
		return fmt.Errorf("interpres: cannot encode %T", val)
	}
}

func (e *encoder) writeStringVal(s string) error {
	return writeQuotedString(&e.buf, s)
}

func (e *encoder) writeFloat(v float64) error {
	switch {
	case math.IsNaN(v):
		e.buf.WriteString("nan")
	case math.IsInf(v, 1):
		e.buf.WriteString("inf")
	case math.IsInf(v, -1):
		e.buf.WriteString("-inf")
	case v == 0:
		// Normalise negative zero to positive zero (TOML has no -0).
		e.buf.WriteString("0.0")
	default:
		s := strconv.FormatFloat(v, 'g', -1, 64)
		// TOML forbids leading zeros in the exponent digits.
		if idx := strings.LastIndexAny(s, "eE"); idx >= 0 {
			mant := s[:idx]
			exp := s[idx+1:] // e.g. "+06", "-05"
			sign := ""
			if len(exp) > 0 && (exp[0] == '+' || exp[0] == '-') {
				sign = string(exp[0])
				exp = exp[1:]
			}
			exp = strings.TrimLeft(exp, "0")
			if exp == "" {
				exp = "0"
			}
			s = mant + "e" + sign + exp
		}
		if !strings.ContainsAny(s, ".eE") {
			s += ".0"
		}
		e.buf.WriteString(s)
	}
	return nil
}