Implement lua engine for the evaluations

2025-03-22 01:43:58 +01:00
parent fa1e0caf06
commit 7e85497ec4
4 changed files with 188 additions and 55 deletions
--- a/main.go
+++ b/main.go
@@ -8,9 +8,10 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strconv"
 	"strings"

-	"github.com/Knetic/govaluate"
+	lua "github.com/yuin/gopher-lua"
 )

 var Error *log.Logger
@@ -30,67 +31,193 @@ func init() {
 }

 func main() {
+	// Define flags
+	flag.Usage = func() {
+		fmt.Fprintf(os.Stderr, "Usage: %s <regex_with_capture_groups> <lua_expression> <...files>\n", os.Args[0])
+		fmt.Fprintf(os.Stderr, "Example: %s \"<value>(\\d+)</value>,(\\d+)\" \"v[1] * 1.5 * v[2]\" data.xml\n", os.Args[0])
+		fmt.Fprintf(os.Stderr, "   or simplified: %s \"<value>(\\d+)</value>,(\\d+)\" \"v1 * 1.5 * v2\" data.xml\n", os.Args[0])
+		fmt.Fprintf(os.Stderr, "   or even simpler: %s \"<value>(\\d+)</value>\" \"*1.5\" data.xml\n", os.Args[0])
+		fmt.Fprintf(os.Stderr, "\nNote: v1, v2, etc. are shortcuts for v[1], v[2], etc.\n")
+		fmt.Fprintf(os.Stderr, "      If expression starts with an operator like *, /, +, -, etc., v1 is automatically prepended\n")
+	}
+
 	flag.Parse()
 	args := flag.Args()
-	if len(args) == 0 {
-		Error.Println("No files provided")
-		log.Printf("Usage: %s <field> <operation> <...files>", os.Args[0])
-		log.Printf("Example: %s go run . \"<ticksBetweenBursts>(\\d+)</ticksBetweenBursts>\" \"/10\" Vehicle_JLTV.xml", os.Args[0])
-		log.Printf("Note: Field must include valid regex with one capture group")
-		return
-	}
-	field := args[0]
-	operation := args[1]
-	files := args[2:]
-	if field == "" || operation == "" || len(files) == 0 {
-		Error.Println("Invalid arguments")
-		log.Printf("Usage: %s <field> <operation> <...files>", os.Args[0])
-		return
-	}
-	log.Printf("Field: %s", field)
-	log.Printf("Operation: %s", operation)
-	log.Printf("Processing files: %v", files)
-	fieldRegex, err := regexp.Compile(field)
-	if err != nil {
-		Error.Printf("Error compiling field regex: %v", err)
+	if len(args) < 3 {
+		Error.Println("Insufficient arguments")
+		flag.Usage()
 		return
 	}

+	regexPattern := args[0]
+	luaExpr := args[1]
+	files := args[2:]
+
+	// Check if the expression needs v1 to be prepended
+	if strings.HasPrefix(luaExpr, "*") || strings.HasPrefix(luaExpr, "/") ||
+		strings.HasPrefix(luaExpr, "+") || strings.HasPrefix(luaExpr, "-") ||
+		strings.HasPrefix(luaExpr, "^") || strings.HasPrefix(luaExpr, "%") {
+		luaExpr = "v[1]" + luaExpr
+		log.Printf("Expression modified to: %s", luaExpr)
+	}
+
+	// Replace shorthand v1, v2, etc. with v[1], v[2]
+	shorthandRegex := regexp.MustCompile(`\bv(\d+)\b`)
+	luaExpr = shorthandRegex.ReplaceAllString(luaExpr, "v[$1]")
+	log.Printf("Final expression: %s", luaExpr)
+
+	// Check if the expression is a simple expression or explicitly manipulates v table
+	var fullScript string
+	// Add custom script header
+	scriptHeader := `
+-- Custom Lua helpers for math operations
+function min(a, b) return math.min(a, b) end
+function max(a, b) return math.max(a, b) end
+function round(x) return math.floor(x + 0.5) end
+function floor(x) return math.floor(x) end
+function ceil(x) return math.ceil(x) end
+`
+
+	if strings.Contains(luaExpr, "v[1] =") || strings.Contains(luaExpr, "v[2] =") ||
+		strings.Contains(luaExpr, "return") {
+		// Already has assignments, use as is
+		fullScript = fmt.Sprintf(`%s
+function process(v)
+    %s
+    return v
+end
+`, scriptHeader, luaExpr)
+	} else {
+		// Simple expression, capture the result and assign to v[1]
+		fullScript = fmt.Sprintf(`%s
+function process(v)
+    local result = %s
+    v[1] = result
+    return v
+end
+`, scriptHeader, luaExpr)
+	}
+
+	log.Printf("Regex pattern: %s", regexPattern)
+	log.Printf("Lua script: %s", fullScript)
+	log.Printf("Processing files: %v", files)
+
+	// Compile the pattern for file processing
+	pattern, err := regexp.Compile(regexPattern)
+	if err != nil {
+		Error.Printf("Invalid regex pattern: %v", err)
+		return
+	}
+
+	// Process each file
 	for _, file := range files {
-		log.Printf("Processing file: %s", file)
-		fullPath := filepath.Join(".", file)
-		content, err := os.ReadFile(fullPath)
+		err := processFile(file, pattern, fullScript)
 		if err != nil {
-			Error.Printf("Error reading file: %v", err)
-			continue
+			Error.Printf("Error processing file %s: %v", file, err)
 		}
-		lines := strings.Split(string(content), "\n")
-		for i, line := range lines {
-			matches := fieldRegex.FindStringSubmatch(line)
-			if len(matches) > 1 {
-				log.Printf("Line: %s", line)
-				oldValue := matches[1]
-				expression, err := govaluate.NewEvaluableExpression(oldValue + operation)
-				if err != nil {
-					Error.Printf("Error creating expression: %v", err)
-					continue
-				}
-				parameters := map[string]interface{}{"oldValue": oldValue}
-				newValue, err := expression.Evaluate(parameters)
-				if err != nil {
-					Error.Printf("Error evaluating expression: %v", err)
-					continue
-				}
-				updatedLine := strings.Replace(line, oldValue, fmt.Sprintf("%v", newValue), 1)
-				log.Printf("Updated line: %s", updatedLine)
-				lines[i] = updatedLine
-			}
-		}
-		err = os.WriteFile(fullPath, []byte(strings.Join(lines, "\n")), 0644)
-		if err != nil {
-			Error.Printf("Error writing file: %v", err)
-			continue
-		}
-		log.Printf("File %s updated successfully", file)
 	}
 }
+
+func processFile(filename string, pattern *regexp.Regexp, luaScript string) error {
+	log.Printf("Processing file: %s", filename)
+	fullPath := filepath.Join(".", filename)
+
+	content, err := os.ReadFile(fullPath)
+	if err != nil {
+		return fmt.Errorf("error reading file: %v", err)
+	}
+
+	fileContent := string(content)
+	modified := false
+
+	// Initialize Lua state once per file
+	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 fmt.Errorf("error loading Lua math library: %v", err)
+	}
+
+	// Load the Lua script
+	if err := L.DoString(luaScript); err != nil {
+		return fmt.Errorf("error in Lua script: %v", err)
+	}
+
+	// Process all regex matches
+	result := pattern.ReplaceAllStringFunc(fileContent, func(match string) string {
+		captures := pattern.FindStringSubmatch(match)
+		if len(captures) <= 1 {
+			// No capture groups, return unchanged
+			return match
+		}
+
+		// Create a Lua table with the captured values
+		vTable := L.NewTable()
+		for i, capture := range captures[1:] {
+			// Convert each capture to float if possible
+			floatVal, err := strconv.ParseFloat(capture, 64)
+			if err == nil {
+				L.RawSetInt(vTable, i+1, lua.LNumber(floatVal))
+			} else {
+				L.RawSetInt(vTable, i+1, lua.LString(capture))
+			}
+		}
+
+		// Call the Lua function with the table
+		L.Push(L.GetGlobal("process"))
+		L.Push(vTable)
+		if err := L.PCall(1, 1, nil); err != nil {
+			Error.Printf("Error executing Lua script: %v", err)
+			return match // Return unchanged on error
+		}
+
+		// Get the result table
+		resultTable := L.Get(-1)
+		L.Pop(1)
+
+		if tbl, ok := resultTable.(*lua.LTable); ok {
+			// Prepare to replace each capture group
+			result := match
+			for i := 1; i <= len(captures)-1; i++ {
+				oldVal := captures[i]
+
+				// Get new value from Lua
+				luaVal := tbl.RawGetInt(i)
+				var newVal string
+
+				switch v := luaVal.(type) {
+				case lua.LNumber:
+					newVal = strconv.FormatFloat(float64(v), 'f', -1, 64)
+				case lua.LString:
+					newVal = string(v)
+				default:
+					newVal = fmt.Sprintf("%v", v)
+				}
+
+				// Replace old value with new value
+				result = strings.Replace(result, oldVal, newVal, 1)
+			}
+			modified = true
+			return result
+		}
+
+		// If something went wrong, return the original match
+		return match
+	})
+
+	// Only write if changes were made
+	if modified {
+		err = os.WriteFile(fullPath, []byte(result), 0644)
+		if err != nil {
+			return fmt.Errorf("error writing file: %v", err)
+		}
+		log.Printf("File %s updated successfully", filename)
+	} else {
+		log.Printf("No changes made to %s", filename)
+	}
+
+	return nil
+}