BigChef/processor/jsonpath/jsonpath.go

package jsonpath

import (
	"fmt"
	"strconv"
)

// JSONStep represents a single step in a JSONPath query
type JSONStep struct {
	Type  StepType
	Key   string // For Child/RecursiveDescent
	Index int    // For Index (use -1 for wildcard "*")
}

// JSONNode represents a value in the JSON data with its path
type JSONNode struct {
	Value interface{} // The value found at the path
	Path  string      // The exact JSONPath where the value was found
}

// StepType defines the types of steps in a JSONPath
type StepType int

const (
	RootStep             StepType = iota // $ - The root element
	ChildStep                            // .key - Direct child access
	RecursiveDescentStep                 // ..key - Recursive search for key
	WildcardStep                         // .* - All children of an object
	IndexStep                            // [n] - Array index access (or [*] for all elements)
)

// TraversalMode determines how the traversal behaves
type TraversalMode int

const (
	CollectMode     TraversalMode = iota // Just collect matched nodes
	ModifyFirstMode                      // Modify first matching node
	ModifyAllMode                        // Modify all matching nodes
)

// ParseJSONPath parses a JSONPath string into a sequence of steps
func ParseJSONPath(path string) ([]JSONStep, error) {
	if len(path) == 0 || path[0] != '$' {
		return nil, fmt.Errorf("path must start with $; received: %q", path)
	}

	steps := []JSONStep{}
	i := 0

	for i < len(path) {
		switch path[i] {
		case '$':
			steps = append(steps, JSONStep{Type: RootStep})
			i++
		case '.':
			i++
			if i < len(path) && path[i] == '.' {
				// Recursive descent
				i++
				key, nextPos := readKey(path, i)
				steps = append(steps, JSONStep{Type: RecursiveDescentStep, Key: key})
				i = nextPos
			} else {
				// Child step or wildcard
				key, nextPos := readKey(path, i)
				if key == "*" {
					steps = append(steps, JSONStep{Type: WildcardStep})
				} else {
					steps = append(steps, JSONStep{Type: ChildStep, Key: key})
				}
				i = nextPos
			}
		case '[':
			// Index step
			i++
			indexStr, nextPos := readIndex(path, i)
			if indexStr == "*" {
				steps = append(steps, JSONStep{Type: IndexStep, Index: -1})
			} else {
				index, err := strconv.Atoi(indexStr)
				if err != nil {
					return nil, fmt.Errorf("invalid index: %s; error: %w", indexStr, err)
				}
				steps = append(steps, JSONStep{Type: IndexStep, Index: index})
			}
			i = nextPos + 1 // Skip closing ]
		default:
			return nil, fmt.Errorf("unexpected character: %c at position %d; path: %q", path[i], i, path)
		}
	}

	return steps, nil
}

// readKey extracts a key name from the path
func readKey(path string, start int) (string, int) {
	i := start
	for ; i < len(path); i++ {
		if path[i] == '.' || path[i] == '[' {
			break
		}
	}
	return path[start:i], i
}

// readIndex extracts an array index or wildcard from the path
func readIndex(path string, start int) (string, int) {
	i := start
	for ; i < len(path); i++ {
		if path[i] == ']' {
			break
		}
	}
	return path[start:i], i
}

// Get retrieves values with their paths from data at the specified JSONPath
// Each returned JSONNode contains both the value and its exact path in the data structure
func Get(data interface{}, path string) ([]JSONNode, error) {
	steps, err := ParseJSONPath(path)
	if err != nil {
		return nil, fmt.Errorf("failed to parse JSONPath %q: %w", path, err)
	}

	results := []JSONNode{}
	err = traverseWithPaths(data, steps, &results, "$")
	if err != nil {
		return nil, fmt.Errorf("failed to traverse JSONPath %q: %w", path, err)
	}
	return results, nil
}

// Set updates the value at the specified JSONPath in the original data structure.
// It only modifies the first matching node.
func Set(data interface{}, path string, value interface{}) error {
	steps, err := ParseJSONPath(path)
	if err != nil {
		return fmt.Errorf("failed to parse JSONPath %q: %w", path, err)
	}

	success := false
	err = setWithPath(data, steps, &success, value, "$", ModifyFirstMode)
	if err != nil {
		return fmt.Errorf("failed to set value at JSONPath %q: %w", path, err)
	}
	return nil
}

// SetAll updates all matching values at the specified JSONPath.
func SetAll(data interface{}, path string, value interface{}) error {
	steps, err := ParseJSONPath(path)
	if err != nil {
		return fmt.Errorf("failed to parse JSONPath %q: %w", path, err)
	}

	success := false
	err = setWithPath(data, steps, &success, value, "$", ModifyAllMode)
	if err != nil {
		return fmt.Errorf("failed to set value at JSONPath %q: %w", path, err)
	}
	return nil
}

// setWithPath modifies values while tracking paths
func setWithPath(node interface{}, steps []JSONStep, success *bool, value interface{}, currentPath string, mode TraversalMode) error {
	if node == nil || *success && mode == ModifyFirstMode {
		return nil
	}

	// Skip root step
	actualSteps := steps
	if len(steps) > 0 && steps[0].Type == RootStep {
		actualSteps = steps[1:]
	}

	// If we have no steps left, we're setting the root value
	if len(actualSteps) == 0 {
		// For the root node, we need to handle it differently depending on what's passed in
		// since we can't directly replace the interface{} variable

		// We'll signal success and let the JSONProcessor handle updating the root
		*success = true
		return nil
	}

	// Process the first step
	step := actualSteps[0]
	remainingSteps := actualSteps[1:]
	isLastStep := len(remainingSteps) == 0

	switch step.Type {
	case ChildStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return fmt.Errorf("node at path %q is not a map; actual type: %T", currentPath, node)
		}

		childPath := currentPath + "." + step.Key

		if isLastStep {
			// We've reached the target, set the value
			m[step.Key] = value
			*success = true
			return nil
		}

		// Create intermediate nodes if necessary
		child, exists := m[step.Key]
		if !exists {
			// Create missing intermediate node
			if len(remainingSteps) > 0 && remainingSteps[0].Type == IndexStep {
				child = []interface{}{}
			} else {
				child = map[string]interface{}{}
			}
			m[step.Key] = child
		}

		err := setWithPath(child, remainingSteps, success, value, childPath, mode)
		if err != nil {
			return fmt.Errorf("failed to set value at JSONPath %q: %w", childPath, err)
		}

	case IndexStep:
		arr, ok := node.([]interface{})
		if !ok {
			return fmt.Errorf("node at path %q is not an array; actual type: %T", currentPath, node)
		}

		// Handle wildcard index
		if step.Index == -1 {
			for i, item := range arr {
				itemPath := fmt.Sprintf("%s[%d]", currentPath, i)
				if isLastStep {
					arr[i] = value
					*success = true
					if mode == ModifyFirstMode {
						return nil
					}
				} else {
					err := setWithPath(item, remainingSteps, success, value, itemPath, mode)
					if err != nil {
						return fmt.Errorf("failed to set value at JSONPath %q: %w", itemPath, err)
					}
					if *success && mode == ModifyFirstMode {
						return nil
					}
				}
			}
			return nil
		}

		// Handle specific index
		if step.Index >= 0 && step.Index < len(arr) {
			item := arr[step.Index]
			itemPath := fmt.Sprintf("%s[%d]", currentPath, step.Index)
			if isLastStep {
				arr[step.Index] = value
				*success = true
			} else {
				err := setWithPath(item, remainingSteps, success, value, itemPath, mode)
				if err != nil {
					return fmt.Errorf("failed to set value at JSONPath %q: %w", itemPath, err)
				}
			}
		}

	case RecursiveDescentStep:
		// For recursive descent, first check direct match at this level
		if m, ok := node.(map[string]interface{}); ok && step.Key != "*" {
			if val, exists := m[step.Key]; exists {
				directPath := currentPath + "." + step.Key
				if isLastStep {
					m[step.Key] = value
					*success = true
					if mode == ModifyFirstMode {
						return nil
					}
				} else {
					err := setWithPath(val, remainingSteps, success, value, directPath, mode)
					if err != nil {
						return fmt.Errorf("failed to set value at JSONPath %q: %w", directPath, err)
					}
					if *success && mode == ModifyFirstMode {
						return nil
					}
				}
			}
		}

		// Then continue recursion to all children
		switch n := node.(type) {
		case map[string]interface{}:
			for k, v := range n {
				childPath := currentPath + "." + k
				// Skip keys we've already processed directly
				if step.Key != "*" && k == step.Key {
					continue
				}
				err := setWithPath(v, steps, success, value, childPath, mode)
				if err != nil {
					return fmt.Errorf("failed to set value at JSONPath %q: %w", childPath, err)
				}
				if *success && mode == ModifyFirstMode {
					return nil
				}
			}
		case []interface{}:
			for i, v := range n {
				childPath := fmt.Sprintf("%s[%d]", currentPath, i)
				err := setWithPath(v, steps, success, value, childPath, mode)
				if err != nil {
					return fmt.Errorf("failed to set value at JSONPath %q: %w", childPath, err)
				}
				if *success && mode == ModifyFirstMode {
					return nil
				}
			}
		}

	case WildcardStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return fmt.Errorf("node at path %q is not a map; actual type: %T", currentPath, node)
		}

		for k, v := range m {
			childPath := currentPath + "." + k
			if isLastStep {
				m[k] = value
				*success = true
				if mode == ModifyFirstMode {
					return nil
				}
			} else {
				err := setWithPath(v, remainingSteps, success, value, childPath, mode)
				if err != nil {
					return fmt.Errorf("failed to set value at JSONPath %q: %w", childPath, err)
				}
				if *success && mode == ModifyFirstMode {
					return nil
				}
			}
		}
	}
	return nil
}

// traverseWithPaths tracks both nodes and their paths during traversal
func traverseWithPaths(node interface{}, steps []JSONStep, results *[]JSONNode, currentPath string) error {
	if len(steps) == 0 || node == nil {
		return fmt.Errorf("cannot traverse with empty steps or nil node; steps length: %d, node: %v", len(steps), node)
	}

	// Skip root step
	actualSteps := steps
	if steps[0].Type == RootStep {
		if len(steps) == 1 {
			*results = append(*results, JSONNode{Value: node, Path: currentPath})
			return nil
		}
		actualSteps = steps[1:]
	}

	// Process the first step
	step := actualSteps[0]
	remainingSteps := actualSteps[1:]
	isLastStep := len(remainingSteps) == 0

	switch step.Type {
	case ChildStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return fmt.Errorf("node is not a map; actual type: %T", node)
		}

		child, exists := m[step.Key]
		if !exists {
			return fmt.Errorf("key not found: %s in node at path: %s", step.Key, currentPath)
		}

		childPath := currentPath + "." + step.Key
		if isLastStep {
			*results = append(*results, JSONNode{Value: child, Path: childPath})
		} else {
			err := traverseWithPaths(child, remainingSteps, results, childPath)
			if err != nil {
				return fmt.Errorf("failed to traverse JSONPath %q: %w", childPath, err)
			}
		}

	case IndexStep:
		arr, ok := node.([]interface{})
		if !ok {
			return fmt.Errorf("node is not an array; actual type: %T", node)
		}

		// Handle wildcard index
		if step.Index == -1 {
			for i, item := range arr {
				itemPath := fmt.Sprintf("%s[%d]", currentPath, i)
				if isLastStep {
					*results = append(*results, JSONNode{Value: item, Path: itemPath})
				} else {
					err := traverseWithPaths(item, remainingSteps, results, itemPath)
					if err != nil {
						return fmt.Errorf("failed to traverse JSONPath %q: %w", itemPath, err)
					}
				}
			}
			return nil
		}

		// Handle specific index
		if step.Index >= 0 && step.Index < len(arr) {
			item := arr[step.Index]
			itemPath := fmt.Sprintf("%s[%d]", currentPath, step.Index)
			if isLastStep {
				*results = append(*results, JSONNode{Value: item, Path: itemPath})
			} else {
				err := traverseWithPaths(item, remainingSteps, results, itemPath)
				if err != nil {
					return fmt.Errorf("failed to traverse JSONPath %q: %w", itemPath, err)
				}
			}
		} else {
			return fmt.Errorf("index %d out of bounds for array at path: %s", step.Index, currentPath)
		}

	case RecursiveDescentStep:
		// For recursive descent, first check direct match at this level
		if m, ok := node.(map[string]interface{}); ok && step.Key != "*" {
			if val, exists := m[step.Key]; exists {
				directPath := currentPath + "." + step.Key
				if isLastStep {
					*results = append(*results, JSONNode{Value: val, Path: directPath})
				} else {
					err := traverseWithPaths(val, remainingSteps, results, directPath)
					if err != nil {
						return fmt.Errorf("failed to traverse JSONPath %q: %w", directPath, err)
					}
				}
			}
		}

		// For wildcard, collect this node
		if step.Key == "*" && isLastStep {
			*results = append(*results, JSONNode{Value: node, Path: currentPath})
		}

		// Then continue recursion to all children
		switch n := node.(type) {
		case map[string]interface{}:
			for k, v := range n {
				childPath := currentPath + "." + k
				err := traverseWithPaths(v, steps, results, childPath) // Use the same steps
				if err != nil {
					return fmt.Errorf("failed to traverse JSONPath %q: %w", childPath, err)
				}
			}
		case []interface{}:
			for i, v := range n {
				childPath := fmt.Sprintf("%s[%d]", currentPath, i)
				err := traverseWithPaths(v, steps, results, childPath) // Use the same steps
				if err != nil {
					return fmt.Errorf("failed to traverse JSONPath %q: %w", childPath, err)
				}
			}
		}

	case WildcardStep:
		m, ok := node.(map[string]interface{})
		if !ok {
			return fmt.Errorf("node is not a map; actual type: %T", node)
		}

		for k, v := range m {
			childPath := currentPath + "." + k
			if isLastStep {
				*results = append(*results, JSONNode{Value: v, Path: childPath})
			} else {
				err := traverseWithPaths(v, remainingSteps, results, childPath)
				if err != nil {
					return fmt.Errorf("failed to traverse JSONPath %q: %w", childPath, err)
				}
			}
		}
	}
	return nil
}