rill/internal/config/upgrade.go

package config

import (
	"fmt"
	"io/fs"
	"log/slog"
	"os"
	"path/filepath"
	"slices"
	"strings"

	"github.com/goccy/go-yaml"
	"github.com/goccy/go-yaml/ast"
	"github.com/goccy/go-yaml/parser"
)

// ConfigVersion 当前配置结构版本，新增配置项时递增。
const ConfigVersion = 3

// upgradeResult 描述一次配置自动补全的结果。
type upgradeResult struct {
	Data    []byte
	Version int
	Added   []string
}

// upgradeConfig 比较配置版本，版本落后时按默认模板递归补全缺失项。
// 返回 nil 表示无需变更。
func upgradeConfig(data []byte) (*upgradeResult, error) {
	file, err := parser.ParseBytes(data, parser.ParseComments)
	if err != nil {
		return nil, fmt.Errorf("解析配置失败: %w", err)
	}
	body, err := rootMapping(file)
	if err != nil {
		return nil, err
	}

	version := versionOf(body)
	if version > ConfigVersion {
		slog.Warn("配置文件版本高于当前程序，可能存在未知配置项，请升级程序",
			"version", version, "supported", ConfigVersion)
		return nil, nil
	}
	if version == ConfigVersion {
		return nil, nil
	}

	defaults, err := defaultMapping()
	if err != nil {
		return nil, err
	}

	added := mergeMissing(body, defaults, "")
	setVersion(body, ConfigVersion)
	return &upgradeResult{Data: []byte(file.String()), Version: version, Added: added}, nil
}

// mergeMissing 以用户配置为基底递归补入默认模板中缺失的键，已有值保持不变。
// 返回新增配置项的路径列表。
func mergeMissing(user, defaults *ast.MappingNode, prefix string) []string {
	var added []string

	index := make(map[string]int, len(user.Values))
	for i, v := range user.Values {
		index[v.Key.String()] = i
	}
	delta := mappingKeyColumn(user) - mappingKeyColumn(defaults)
	insertAt := 0

	for _, dv := range defaults.Values {
		key := dv.Key.String()
		full := joinPath(prefix, key)

		i, exists := index[key]
		if !exists {
			node := dv
			node.AddColumn(delta)
			user.Values = slices.Insert(user.Values, insertAt, node)
			shiftIndex(index, insertAt)
			index[key] = insertAt
			insertAt++
			added = append(added, full)
			continue
		}
		if i+1 > insertAt {
			insertAt = i + 1
		}

		uv := user.Values[i]
		switch {
		case isMapping(uv.Value) && isMapping(dv.Value):
			added = append(added, mergeMissing(
				uv.Value.(*ast.MappingNode), dv.Value.(*ast.MappingNode), full)...)
		case isNull(uv.Value) && isMapping(dv.Value):
			dm := dv.Value.(*ast.MappingNode)
			dm.AddColumn(uv.Key.GetToken().Position.Column - dv.Key.GetToken().Position.Column)
			uv.Value = dm
			added = append(added, leafPaths(dm, full)...)
		}
	}
	return added
}

// setVersion 将配置中的 version 更新为当前版本（保留原行内注释）。
func setVersion(body *ast.MappingNode, version int) {
	for _, v := range body.Values {
		if v.Key.String() != "version" {
			continue
		}
		node, err := yaml.ValueToNode(version)
		if err != nil {
			return
		}
		if comment := v.Value.GetComment(); comment != nil {
			_ = node.SetComment(comment)
		}
		v.Value = node
		return
	}
}

// versionOf 读取配置中的 version，缺失或无效时视为 0。
func versionOf(body *ast.MappingNode) int {
	for _, v := range body.Values {
		if v.Key.String() != "version" {
			continue
		}
		var version int
		if err := yaml.NodeToValue(v.Value, &version); err != nil {
			return 0
		}
		return version
	}
	return 0
}

func defaultMapping() (*ast.MappingNode, error) {
	file, err := parser.ParseBytes(defaultConfigYAML, parser.ParseComments)
	if err != nil {
		return nil, fmt.Errorf("解析默认配置模板失败: %w", err)
	}
	return rootMapping(file)
}

func rootMapping(file *ast.File) (*ast.MappingNode, error) {
	if len(file.Docs) == 0 || file.Docs[0].Body == nil {
		return nil, fmt.Errorf("配置内容为空")
	}
	body, ok := file.Docs[0].Body.(*ast.MappingNode)
	if !ok {
		return nil, fmt.Errorf("配置根节点必须是映射")
	}
	return body, nil
}

func mappingKeyColumn(m *ast.MappingNode) int {
	if len(m.Values) > 0 && m.Values[0].Key != nil {
		return m.Values[0].Key.GetToken().Position.Column
	}
	if m.Start != nil {
		return m.Start.Position.Column
	}
	return 1
}

func joinPath(prefix, key string) string {
	if prefix == "" {
		return key
	}
	return prefix + "." + key
}

func leafPaths(m *ast.MappingNode, prefix string) []string {
	var paths []string
	for _, v := range m.Values {
		key := joinPath(prefix, v.Key.String())
		if child, ok := v.Value.(*ast.MappingNode); ok && len(child.Values) > 0 {
			paths = append(paths, leafPaths(child, key)...)
			continue
		}
		paths = append(paths, key)
	}
	return paths
}

func isMapping(node ast.Node) bool {
	_, ok := node.(*ast.MappingNode)
	return ok
}

func isNull(node ast.Node) bool {
	if node == nil {
		return true
	}
	_, ok := node.(*ast.NullNode)
	return ok
}

func shiftIndex(index map[string]int, from int) {
	for key, i := range index {
		if i >= from {
			index[key] = i + 1
		}
	}
}

// applyConfigUpgrade 执行配置检查与补全；失败时仅告警并使用原内容，不阻塞启动。
func applyConfigUpgrade(path string, data []byte) []byte {
	result, err := upgradeConfig(data)
	if err != nil {
		slog.Warn("配置自动补全检查失败，将使用默认值补齐缺失项", "path", path, "err", err)
		return data
	}
	if result == nil {
		return data
	}

	if err := backupConfig(path); err != nil {
		slog.Warn("备份配置文件失败", "path", path, "err", err)
	}
	if err := writeFileAtomic(path, result.Data); err != nil {
		slog.Warn("配置自动补全写回失败，将使用默认值补齐缺失项", "path", path, "err", err)
		return data
	}
	slog.Info("配置已自动补全",
		"path", path,
		"version", fmt.Sprintf("%d -> %d", result.Version, ConfigVersion),
		"added", strings.Join(result.Added, ", "))
	return result.Data
}

func backupConfig(path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	return os.WriteFile(path+".bak", data, fileMode(path))
}

func writeFileAtomic(path string, data []byte) error {
	tmp, err := os.CreateTemp(filepath.Dir(path), ".rill-config-*.tmp")
	if err != nil {
		return err
	}
	tmpName := tmp.Name()
	defer os.Remove(tmpName)

	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := os.Chmod(tmpName, fileMode(path)); err != nil {
		return err
	}
	return os.Rename(tmpName, path)
}

func fileMode(path string) fs.FileMode {
	if info, err := os.Stat(path); err == nil {
		return info.Mode().Perm()
	}
	return 0o644
}