BigChef/processor/json.go

package processor

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/PaesslerAG/jsonpath"
	lua "github.com/yuin/gopher-lua"
)

// JSONProcessor implements the Processor interface for JSON documents
type JSONProcessor struct{}

// Process implements the Processor interface for JSONProcessor
func (p *JSONProcessor) Process(filename string, pattern string, luaExpr string) (int, int, error) {
	// Read file content
	fullPath := filepath.Join(".", filename)
	content, err := os.ReadFile(fullPath)
	if err != nil {
		return 0, 0, fmt.Errorf("error reading file: %v", err)
	}

	fileContent := string(content)

	// Process the content
	modifiedContent, modCount, matchCount, err := p.ProcessContent(fileContent, pattern, luaExpr)
	if err != nil {
		return 0, 0, err
	}

	// If we made modifications, save the file
	if modCount > 0 {
		err = os.WriteFile(fullPath, []byte(modifiedContent), 0644)
		if err != nil {
			return 0, 0, fmt.Errorf("error writing file: %v", err)
		}
	}

	return modCount, matchCount, nil
}

// ProcessContent implements the Processor interface for JSONProcessor
func (p *JSONProcessor) ProcessContent(content string, pattern string, luaExpr string) (string, int, int, error) {
	// Parse JSON document
	var jsonData interface{}
	err := json.Unmarshal([]byte(content), &jsonData)
	if err != nil {
		return content, 0, 0, fmt.Errorf("error parsing JSON: %v", err)
	}

	// Find nodes matching the JSONPath pattern
	paths, values, err := p.findJSONPaths(jsonData, pattern)
	if err != nil {
		return content, 0, 0, fmt.Errorf("error executing JSONPath: %v", err)
	}

	matchCount := len(paths)
	if matchCount == 0 {
		return content, 0, 0, nil
	}

	// Initialize Lua
	L := lua.NewState()
	defer L.Close()

	// Load math library
	L.Push(L.GetGlobal("require"))
	L.Push(lua.LString("math"))
	if err := L.PCall(1, 1, nil); err != nil {
		return content, 0, 0, fmt.Errorf("error loading Lua math library: %v", err)
	}

	// Load helper functions
	if err := InitLuaHelpers(L); err != nil {
		return content, 0, 0, err
	}

	// Apply modifications to each node
	modCount := 0
	for i, value := range values {
		// Reset Lua state for each node
		L.SetGlobal("v1", lua.LNil)
		L.SetGlobal("s1", lua.LNil)

		// Convert to Lua variables
		err = p.ToLua(L, value)
		if err != nil {
			return content, modCount, matchCount, fmt.Errorf("error converting to Lua: %v", err)
		}

		// Execute Lua script
		if err := L.DoString(luaExpr); err != nil {
			return content, modCount, matchCount, fmt.Errorf("error executing Lua: %v", err)
		}

		// Get modified value
		result, err := p.FromLua(L)
		if err != nil {
			return content, modCount, matchCount, fmt.Errorf("error getting result from Lua: %v", err)
		}

		// Skip if value didn't change
		if fmt.Sprintf("%v", value) == fmt.Sprintf("%v", result) {
			continue
		}

		// Apply the modification to the JSON data
		err = p.updateJSONValue(jsonData, paths[i], result)
		if err != nil {
			return content, modCount, matchCount, fmt.Errorf("error updating JSON: %v", err)
		}

		modCount++
	}

	// Convert the modified JSON back to a string
	jsonBytes, err := json.MarshalIndent(jsonData, "", "  ")
	if err != nil {
		return content, modCount, matchCount, fmt.Errorf("error serializing JSON: %v", err)
	}

	return string(jsonBytes), modCount, matchCount, nil
}

// findJSONPaths finds all JSON paths and their values that match the given JSONPath expression
func (p *JSONProcessor) findJSONPaths(jsonData interface{}, pattern string) ([]string, []interface{}, error) {
	// Extract all matching values using JSONPath
	values, err := jsonpath.Get(pattern, jsonData)
	if err != nil {
		return nil, nil, err
	}

	// Convert values to a slice if it's not already
	valuesSlice := []interface{}{}
	paths := []string{}

	switch v := values.(type) {
	case []interface{}:
		valuesSlice = v
		// Generate paths for array elements
		// This is simplified - for complex JSONPath expressions you might
		// need a more robust approach to generate the exact path
		basePath := pattern
		if strings.Contains(pattern, "[*]") || strings.HasSuffix(pattern, ".*") {
			basePath = strings.Replace(pattern, "[*]", "", -1)
			basePath = strings.Replace(basePath, ".*", "", -1)
			for i := 0; i < len(v); i++ {
				paths = append(paths, fmt.Sprintf("%s[%d]", basePath, i))
			}
		} else {
			for range v {
				paths = append(paths, pattern)
			}
		}
	default:
		valuesSlice = append(valuesSlice, v)
		paths = append(paths, pattern)
	}

	return paths, valuesSlice, nil
}

// updateJSONValue updates a value in the JSON data structure at the given path
func (p *JSONProcessor) updateJSONValue(jsonData interface{}, path string, newValue interface{}) error {
	// This is a simplified approach - for a production system you'd need a more robust solution
	// that can handle all JSONPath expressions
	parts := strings.Split(path, ".")
	current := jsonData

	// Traverse the JSON structure
	for i, part := range parts {
		if i == len(parts)-1 {
			// Last part, set the value
			if strings.HasSuffix(part, "]") {
				// Handle array access
				arrayPart := part[:strings.Index(part, "[")]
				indexPart := part[strings.Index(part, "[")+1 : strings.Index(part, "]")]
				index, err := strconv.Atoi(indexPart)
				if err != nil {
					return fmt.Errorf("invalid array index: %s", indexPart)
				}

				// Get the array
				var array []interface{}
				if arrayPart == "" {
					// Direct array access
					array, _ = current.([]interface{})
				} else {
					// Access array property
					obj, _ := current.(map[string]interface{})
					array, _ = obj[arrayPart].([]interface{})
				}

				// Set the value
				if index >= 0 && index < len(array) {
					array[index] = newValue
				}
			} else {
				// Handle object property
				obj, _ := current.(map[string]interface{})
				obj[part] = newValue
			}
			break
		}

		// Not the last part, continue traversing
		if strings.HasSuffix(part, "]") {
			// Handle array access
			arrayPart := part[:strings.Index(part, "[")]
			indexPart := part[strings.Index(part, "[")+1 : strings.Index(part, "]")]
			index, err := strconv.Atoi(indexPart)
			if err != nil {
				return fmt.Errorf("invalid array index: %s", indexPart)
			}

			// Get the array
			var array []interface{}
			if arrayPart == "" {
				// Direct array access
				array, _ = current.([]interface{})
			} else {
				// Access array property
				obj, _ := current.(map[string]interface{})
				array, _ = obj[arrayPart].([]interface{})
			}

			// Continue with the array element
			if index >= 0 && index < len(array) {
				current = array[index]
			}
		} else {
			// Handle object property
			obj, _ := current.(map[string]interface{})
			current = obj[part]
		}
	}

	return nil
}

// ToLua converts JSON values to Lua variables
func (p *JSONProcessor) ToLua(L *lua.LState, data interface{}) error {
	switch v := data.(type) {
	case float64:
		L.SetGlobal("v1", lua.LNumber(v))
		L.SetGlobal("s1", lua.LString(fmt.Sprintf("%v", v)))
	case int:
		L.SetGlobal("v1", lua.LNumber(v))
		L.SetGlobal("s1", lua.LString(fmt.Sprintf("%d", v)))
	case string:
		L.SetGlobal("s1", lua.LString(v))
		// Try to convert to number if possible
		if val, err := strconv.ParseFloat(v, 64); err == nil {
			L.SetGlobal("v1", lua.LNumber(val))
		} else {
			L.SetGlobal("v1", lua.LNumber(0))
		}
	case bool:
		if v {
			L.SetGlobal("v1", lua.LNumber(1))
		} else {
			L.SetGlobal("v1", lua.LNumber(0))
		}
		L.SetGlobal("s1", lua.LString(fmt.Sprintf("%v", v)))
	default:
		// For complex types, convert to string
		L.SetGlobal("s1", lua.LString(fmt.Sprintf("%v", v)))
		L.SetGlobal("v1", lua.LNumber(0))
	}
	return nil
}

// FromLua retrieves values from Lua
func (p *JSONProcessor) FromLua(L *lua.LState) (interface{}, error) {
	// Check if string variable was modified
	s1 := L.GetGlobal("s1")
	if s1 != lua.LNil {
		if s1Str, ok := s1.(lua.LString); ok {
			// Try to convert to number if it's numeric
			if val, err := strconv.ParseFloat(string(s1Str), 64); err == nil {
				return val, nil
			}
			// If it's "true" or "false", convert to boolean
			if string(s1Str) == "true" {
				return true, nil
			}
			if string(s1Str) == "false" {
				return false, nil
			}
			return string(s1Str), nil
		}
	}

	// Check if numeric variable was modified
	v1 := L.GetGlobal("v1")
	if v1 != lua.LNil {
		if v1Num, ok := v1.(lua.LNumber); ok {
			return float64(v1Num), nil
		}
	}

	// Default return nil
	return nil, nil
}