package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

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

var Error *log.Logger
var Warning *log.Logger

func init() {
	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
	logger := io.MultiWriter(os.Stdout)
	log.SetOutput(logger)

	Error = log.New(io.MultiWriter(os.Stderr, os.Stdout),
		fmt.Sprintf("%sERROR:%s ", "\033[0;101m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
	Warning = log.New(os.Stdout,
		fmt.Sprintf("%sWarning:%s ", "\033[0;93m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)
}

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+)\" \"v1 * 1.5 * v2\" 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, "   or direct assignment: %s \"<value>(\\d+)</value>\" \"=0\" data.xml\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "\nNote: v1, v2, etc. are used to refer to capture groups.\n")
		fmt.Fprintf(os.Stderr, "      If expression starts with an operator like *, /, +, -, =, etc., v1 is automatically prepended\n")
		fmt.Fprintf(os.Stderr, "      You can use any valid Lua code, including if statements, loops, etc.\n")
	}

	flag.Parse()
	args := flag.Args()
	if len(args) < 3 {
		Error.Println("Insufficient arguments")
		flag.Usage()
		return
	}

	regexPattern := args[0]
	luaExpr := args[1]
	files := args[2:]

	log.Printf("Parsed arguments: regexPattern=%s, luaExpr=%s, files=%v", regexPattern, luaExpr, files)

	// Prepare the Lua expression
	luaExpr = buildLuaScript(luaExpr)
	log.Printf("Built Lua expression: %s", luaExpr)

	if strings.Contains(regexPattern, "!num") {
		regexPattern = strings.ReplaceAll(regexPattern, "!num", "(\\d*\\.?\\d+)")
		log.Printf("Updated regexPattern after replacing !num: %s", regexPattern)
	}

	// Make sure the regex can match across multiple lines by adding (?s) flag
	if !strings.HasPrefix(regexPattern, "(?s)") {
		regexPattern = "(?s)" + regexPattern
		log.Printf("Updated regexPattern to match across multiple lines: %s", regexPattern)
	}

	// 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)
		err := processFile(file, pattern, luaExpr)
		if err != nil {
			Error.Printf("Error processing file %s: %v", file, err)
		}
	}
}

// buildLuaScript creates a complete Lua script from the expression
func buildLuaScript(luaExpr string) string {
	if strings.HasPrefix(luaExpr, "*") {
		luaExpr = "v1 = v1" + luaExpr
	} else if strings.HasPrefix(luaExpr, "/") {
		luaExpr = "v1 = v1" + luaExpr
	} else if strings.HasPrefix(luaExpr, "+") {
		luaExpr = "v1 = v1" + luaExpr
	} else if strings.HasPrefix(luaExpr, "-") {
		luaExpr = "v1 = v1" + luaExpr
	} else if strings.HasPrefix(luaExpr, "^") {
		luaExpr = "v1 = v1" + luaExpr
	} else if strings.HasPrefix(luaExpr, "%") {
		luaExpr = "v1 = v1" + luaExpr
	} else if strings.HasPrefix(luaExpr, "=") {
		// Handle direct assignment with = operator
		luaExpr = "v1 " + luaExpr
	}
	// Check if the expression starts with an operator that needs special handling
	if !strings.Contains(luaExpr, "=") {
		luaExpr = "v1 = " + luaExpr
	}

	// Replace shorthand v1, v2, etc. with their direct variable names
	shorthandRegex := regexp.MustCompile(`\bv\[(\d+)\]\b`)
	luaExpr = shorthandRegex.ReplaceAllString(luaExpr, "v$1")

	log.Printf("Final Lua expression after processing: %s", luaExpr)
	return luaExpr
}

func processFile(filename string, pattern *regexp.Regexp, luaExpr string) error {
	fullPath := filepath.Join(".", filename)
	log.Printf("Processing file: %s", fullPath)

	content, err := os.ReadFile(fullPath)
	if err != nil {
		Error.Printf("Error reading file %s: %v", fullPath, err)
		return fmt.Errorf("error reading file: %v", err)
	}
	log.Printf("Successfully read file: %s, content length: %d bytes", fullPath, len(content))

	fileContent := string(content)
	result, err := process(fileContent, pattern, luaExpr)
	if err != nil {
		Error.Printf("Error processing content of file %s: %v", fullPath, err)
		return err
	}
	log.Printf("Successfully processed content of file: %s, result length: %d bytes", fullPath, len(result))

	err = os.WriteFile(fullPath, []byte(result), 0644)
	if err != nil {
		Error.Printf("Error writing file %s: %v", fullPath, err)
		return fmt.Errorf("error writing file: %v", err)
	}
	log.Printf("Successfully wrote modified content to file: %s, new content length: %d bytes", fullPath, len(result))

	return nil
}

func process(data string, pattern *regexp.Regexp, luaExpr string) (string, error) {
	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 {
		Error.Printf("Error loading Lua math library: %v", err)
		return data, fmt.Errorf("error loading Lua math library: %v", err)
	}

	// Initialize helper functions
	helperScript := `
-- 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 err := L.DoString(helperScript); err != nil {
		Error.Printf("Error loading helper functions: %v", err)
		return data, fmt.Errorf("error loading helper functions: %v", err)
	}

	// Process all regex matches
	result := pattern.ReplaceAllStringFunc(data, func(match string) string {
		captures := pattern.FindStringSubmatch(match)
		if len(captures) <= 1 {
			// No capture groups, return unchanged
			log.Printf("No capture groups found for match: %s", match)
			return match
		}

		// Set up global variables v1, v2, etc.
		for i, capture := range captures[1:] {
			// Convert each capture to float if possible
			floatVal, err := strconv.ParseFloat(capture, 64)
			if err == nil {
				L.SetGlobal(fmt.Sprintf("v%d", i+1), lua.LNumber(floatVal))
				log.Printf("Set global v%d to float value: %f", i+1, floatVal)
			} else {
				L.SetGlobal(fmt.Sprintf("v%d", i+1), lua.LString(capture))
				log.Printf("Set global v%d to string value: %s", i+1, capture)
			}
		}

		// Execute the user's Lua code
		if err := L.DoString(luaExpr); err != nil {
			Error.Printf("Error executing Lua script: %v", err)
			return match // Return unchanged on error
		}

		// Get the values of v1-v12 after code execution
		returnValues := make([]lua.LValue, 12)
		for i := 0; i < 12; i++ {
			varName := fmt.Sprintf("v%d", i+1)
			returnValues[i] = L.GetGlobal(varName)
			log.Printf("Retrieved value for %s: %v", varName, returnValues[i])
		}

		// Replace each capture group with its new value (if available)
		result := match
		for i := 0; i < len(captures)-1 && i < 12; i++ {
			if returnValues[i] == lua.LNil {
				log.Printf("Skipping replacement for %s as it is nil", captures[i+1])
				continue // Skip if nil
			}

			oldVal := captures[i+1]
			var newVal string

			switch v := returnValues[i].(type) {
			case lua.LNumber:
				newVal = strconv.FormatFloat(float64(v), 'f', -1, 64)
			case lua.LString:
				newVal = string(v)
			default:
				newVal = fmt.Sprintf("%v", v)
			}

			// If we have a value, replace it
			if newVal != "" {
				log.Printf("Replacing %s with %s in match: %s", oldVal, newVal, match)
				result = strings.Replace(result, oldVal, newVal, 1)
			}
		}

		return result
	})

	return result, nil
}