sub2clash/common/yaml_patch.go

package common

import (
	"fmt"
	"strings"

	"gopkg.in/yaml.v3"
)

// BuiltSub 保存最终输出所需的完整 YAML 树。
//
// 这里刻意不再保存整份 typed 配置副本：
// - root 是整个转换流程的最终产物
// - 所有常规输出都直接从 root 序列化
// - nodeList 模式也从 root 中提取 proxies，而不是依赖额外状态
type BuiltSub struct {
	root *yaml.Node
}

// MarshalYAML 让 BuiltSub 在输出时直接复用 patch 后的 YAML 树，
// 从而避免再次经过 struct round-trip 丢失未知字段。
func (b *BuiltSub) MarshalYAML() (any, error) {
	if b == nil || b.root == nil {
		return nil, nil
	}
	if b.root.Kind == yaml.DocumentNode {
		if len(b.root.Content) == 0 {
			return nil, nil
		}
		return b.root.Content[0], nil
	}
	return b.root, nil
}

// MarshalNodeListYAML 从最终 YAML 树中提取 proxies 节点，构造 nodeList 模式输出。
// 这样 nodeList 也直接复用最终 root，而不是依赖额外的 typed struct 副本。
func (b *BuiltSub) MarshalNodeListYAML() ([]byte, error) {
	if b == nil || b.root == nil {
		return yaml.Marshal(&yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"})
	}

	proxiesNode, err := GetYAMLPath(b.root, "proxies")
	if err != nil {
		return nil, err
	}

	root := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
	if proxiesNode != nil && !isNullYAMLNode(proxiesNode) {
		setMappingValue(root, "proxies", cloneYAMLNode(proxiesNode))
	}

	return yaml.Marshal(root)
}

// ParseYAMLDocument 把原始 YAML 解析成 DocumentNode，
// 并确保根内容最终是一个可写入的 mapping 节点。
func ParseYAMLDocument(data []byte) (*yaml.Node, error) {
	var doc yaml.Node
	if err := yaml.Unmarshal(data, &doc); err != nil {
		return nil, err
	}
	if _, err := rootMappingNode(&doc); err != nil {
		return nil, err
	}
	return &doc, nil
}

// HasYAMLPath 判断某个点路径是否存在。
// 这里仅关心“是否找到节点”，不关心节点具体类型。
func HasYAMLPath(doc *yaml.Node, path string) bool {
	current, err := GetYAMLPath(doc, path)
	return err == nil && current != nil
}

// GetYAMLPath 按 a.b.c 这种点路径向下查找节点。
// 当前实现只支持 mapping 之间的逐层下钻，不处理数组索引路径。
func GetYAMLPath(doc *yaml.Node, path string) (*yaml.Node, error) {
	segments := splitYAMLPath(path)
	if len(segments) == 0 {
		return nil, fmt.Errorf("yaml path is empty")
	}

	current, err := rootMappingNode(doc)
	if err != nil {
		return nil, err
	}

	for _, segment := range segments {
		next := findMappingValue(current, segment)
		if next == nil {
			return nil, nil
		}
		current = next
	}

	return current, nil
}

// SetYAMLPath 按点路径写入一个值；不存在的中间层会自动补成 mapping。
// 例如 a.b.c=1 会在缺失时依次创建 a 和 b 两层对象节点。
func SetYAMLPath(doc *yaml.Node, path string, value any) error {
	segments := splitYAMLPath(path)
	if len(segments) == 0 {
		return fmt.Errorf("yaml path is empty")
	}

	current, err := rootMappingNode(doc)
	if err != nil {
		return err
	}

	for idx, segment := range segments[:len(segments)-1] {
		next := findMappingValue(current, segment)
		if next == nil {
			next = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
			setMappingValue(current, segment, next)
		}
		if next.Kind != yaml.MappingNode {
			return fmt.Errorf("yaml path %q segment %q is not a mapping", path, strings.Join(segments[:idx+1], "."))
		}
		current = next
	}

	encoded, err := encodeYAMLNode(value)
	if err != nil {
		return err
	}

	setMappingValue(current, segments[len(segments)-1], encoded)
	return nil
}

// EnsureYAMLSequencePath 确保某个路径最终是 sequence（YAML 数组）节点。
// 不存在时会自动创建，已存在但类型不匹配时返回错误。
func EnsureYAMLSequencePath(doc *yaml.Node, path string) (*yaml.Node, error) {
	return ensureYAMLPathKind(doc, path, yaml.SequenceNode, "!!seq")
}

// EnsureYAMLMappingPath 确保某个路径最终是 mapping（YAML 对象）节点。
func EnsureYAMLMappingPath(doc *yaml.Node, path string) (*yaml.Node, error) {
	return ensureYAMLPathKind(doc, path, yaml.MappingNode, "!!map")
}

// SetYAMLMappingField 在一个 mapping 节点里设置单个字段。
// 它等价于“在当前对象上写 key: value”。
func SetYAMLMappingField(node *yaml.Node, key string, value any) error {
	if node == nil || node.Kind != yaml.MappingNode {
		return fmt.Errorf("yaml node is not a mapping")
	}

	encoded, err := encodeYAMLNode(value)
	if err != nil {
		return err
	}

	setMappingValue(node, key, encoded)
	return nil
}

// AppendYAMLSequenceValue 向 sequence 节点末尾追加一个元素。
func AppendYAMLSequenceValue(node *yaml.Node, value any) error {
	if node == nil || node.Kind != yaml.SequenceNode {
		return fmt.Errorf("yaml node is not a sequence")
	}

	encoded, err := encodeYAMLNode(value)
	if err != nil {
		return err
	}

	node.Content = append(node.Content, encoded)
	return nil
}

// FindYAMLSequenceMappingByStringField 在 YAML 数组中查找一个对象元素，
// 要求该对象存在指定字段且字段值等于目标字符串。
//
// 例如在 proxy-groups 里按 name 查找：
//   - name: 节点选择
//     type: select
func FindYAMLSequenceMappingByStringField(node *yaml.Node, field string, value string) *yaml.Node {
	if node == nil || node.Kind != yaml.SequenceNode {
		return nil
	}

	for _, item := range node.Content {
		if item == nil || item.Kind != yaml.MappingNode {
			continue
		}
		fieldNode := findMappingValue(item, field)
		if fieldNode == nil || fieldNode.Kind != yaml.ScalarNode {
			continue
		}
		if fieldNode.Value == value {
			return item
		}
	}

	return nil
}

// splitYAMLPath 把 a.b.c 这种点路径拆成 [a b c]。
// 空片段会被忽略，避免出现连续点号时产生无意义路径段。
func splitYAMLPath(path string) []string {
	parts := strings.Split(path, ".")
	segments := make([]string, 0, len(parts))
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		segments = append(segments, part)
	}
	return segments
}

// ensureYAMLPathKind 是 EnsureYAMLSequencePath / EnsureYAMLMappingPath 的底层实现。
// 它会：
// 1. 逐层确保中间节点存在且都是 mapping
// 2. 确保最后一个节点存在，且类型符合预期
func ensureYAMLPathKind(doc *yaml.Node, path string, kind yaml.Kind, tag string) (*yaml.Node, error) {
	segments := splitYAMLPath(path)
	if len(segments) == 0 {
		return nil, fmt.Errorf("yaml path is empty")
	}

	current, err := rootMappingNode(doc)
	if err != nil {
		return nil, err
	}

	// 跳过最后一个元素在后面处理
	for idx, segment := range segments[:len(segments)-1] {
		next := findMappingValue(current, segment)
		if next == nil || isNullYAMLNode(next) {
			next = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
			setMappingValue(current, segment, next)
		}
		if next.Kind != yaml.MappingNode {
			return nil, fmt.Errorf("yaml path %q segment %q is not a mapping", path, strings.Join(segments[:idx+1], "."))
		}
		current = next
	}

	lastSegment := segments[len(segments)-1]
	node := findMappingValue(current, lastSegment)
	if node == nil || isNullYAMLNode(node) {
		node = &yaml.Node{Kind: kind, Tag: tag}
		setMappingValue(current, lastSegment, node)
	}
	if node.Kind != kind {
		return nil, fmt.Errorf("yaml path %q is not a %s", path, yamlKindName(kind))
	}

	return node, nil
}

// rootMappingNode 统一把“文档根”整理成一个可操作的 mapping 节点。
//
// yaml.v3 通常把整份 YAML 包在 DocumentNode 下，真正的内容位于 Content[0]。
// 当前项目的 patch 逻辑都假定最外层是 key-value 结构，因此这里会：
// 1. 处理空文档
// 2. 取出 DocumentNode 的实际根内容
// 3. 确保该根内容是 mapping
func rootMappingNode(doc *yaml.Node) (*yaml.Node, error) {
	if doc == nil {
		return nil, fmt.Errorf("yaml document is nil")
	}

	root := doc
	if doc.Kind == 0 {
		doc.Kind = yaml.DocumentNode
	}
	if doc.Kind == yaml.DocumentNode {
		if len(doc.Content) == 0 {
			doc.Content = append(doc.Content, &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"})
		}
		root = doc.Content[0]
	}

	if root.Kind == 0 {
		root.Kind = yaml.MappingNode
		root.Tag = "!!map"
	}
	if root.Kind != yaml.MappingNode {
		return nil, fmt.Errorf("yaml root must be a mapping node")
	}

	return root, nil
}

// isNullYAMLNode 判断一个节点是否为空/未初始化/null。
// 这让我们在“路径不存在”和“路径存在但值为 null”时都能按缺失处理。
func isNullYAMLNode(node *yaml.Node) bool {
	if node == nil {
		return true
	}
	if node.Kind == 0 {
		return true
	}
	return node.Kind == yaml.ScalarNode && node.Tag == "!!null"
}

// yamlKindName 仅用于生成更可读的错误信息。
func yamlKindName(kind yaml.Kind) string {
	switch kind {
	case yaml.MappingNode:
		return "mapping"
	case yaml.SequenceNode:
		return "sequence"
	case yaml.ScalarNode:
		return "scalar"
	case yaml.DocumentNode:
		return "document"
	default:
		return "node"
	}
}

// findMappingValue 在 mapping 节点中按 key 查找对应的 value 节点。
//
// 需要注意：yaml.v3 的 MappingNode.Content 不是 map，而是交替存储：
// [key1, value1, key2, value2, ...]
// 所以这里每次 idx += 2，依次跳过一个完整的 key-value 对。
func findMappingValue(node *yaml.Node, key string) *yaml.Node {
	if node == nil || node.Kind != yaml.MappingNode {
		return nil
	}
	for idx := 0; idx+1 < len(node.Content); idx += 2 {
		if node.Content[idx].Value == key {
			return node.Content[idx+1]
		}
	}
	return nil
}

// setMappingValue 在 mapping 节点中设置 key 对应的 value。
// 如果 key 已存在，就原位替换；否则在末尾追加一组新的 key-value。
func setMappingValue(node *yaml.Node, key string, value *yaml.Node) {
	for idx := 0; idx+1 < len(node.Content); idx += 2 {
		if node.Content[idx].Value == key {
			node.Content[idx+1] = value
			return
		}
	}

	node.Content = append(node.Content,
		&yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key},
		value,
	)
}

// encodeYAMLNode 把普通 Go 值编码成 *yaml.Node，方便统一塞回 YAML 树。
// 如果 Encode 产生的是 DocumentNode，这里会自动取出它的实际内容节点。
func encodeYAMLNode(value any) (*yaml.Node, error) {
	var node yaml.Node
	if err := node.Encode(value); err != nil {
		return nil, err
	}
	if node.Kind == yaml.DocumentNode {
		if len(node.Content) == 0 {
			return &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}, nil
		}
		return node.Content[0], nil
	}
	return &node, nil
}

// cloneYAMLNode 深拷贝一个节点树，避免把同一个子树同时挂到多个输出根下。
func cloneYAMLNode(node *yaml.Node) *yaml.Node {
	if node == nil {
		return nil
	}

	clone := *node
	if len(node.Content) != 0 {
		clone.Content = make([]*yaml.Node, len(node.Content))
		for i := range node.Content {
			clone.Content[i] = cloneYAMLNode(node.Content[i])
		}
	}
	return &clone
}