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")
	}

	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:]

	// Generate the Lua script
	luaScript := buildLuaScript(luaExpr)

	if strings.Contains(regexPattern, "!num") {
		regexPattern = strings.ReplaceAll(regexPattern, "!num", "(\\d*\\.?\\d+)")
	}

	// Make sure the regex can match across multiple lines by adding (?s) flag
	if !strings.HasPrefix(regexPattern, "(?s)") {
		regexPattern = "(?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 {
		err := processFile(file, pattern, luaScript)
		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 {
	log.Printf("Building Lua script from expression: %s", luaExpr)

	// 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 = "v1" + luaExpr
		log.Printf("Prepended 'v1' to expression, new expression: %s", luaExpr)
	} else if strings.HasPrefix(luaExpr, "=") {
		// Handle direct assignment with = operator
		luaExpr = "v1 " + luaExpr
		log.Printf("Handled direct assignment, new expression: %s", luaExpr)
	}

	// Replace shorthand v1, v2, etc. with their direct variable names (no need for array notation)
	// Note: we're keeping this regex in case the user writes v[1] style, but we'll convert it to v1
	shorthandRegex := regexp.MustCompile(`\bv\[(\d+)\]\b`)
	luaExpr = shorthandRegex.ReplaceAllString(luaExpr, "v$1")
	log.Printf("Converted shorthand variables, new expression: %s", luaExpr)

	// Add custom script header with helper functions
	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
`

	// Generate the full script
	var fullScript string
	if strings.Contains(luaExpr, "v1 =") || strings.Contains(luaExpr, "v2 =") ||
		strings.Contains(luaExpr, "v3 =") || strings.Contains(luaExpr, "return") {
		// Already has assignments, use as is
		fullScript = fmt.Sprintf(`%s
function process(...)
    %s
    return v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12
end
`, scriptHeader, luaExpr)
		log.Println("Generated Lua script with assignments.")
	} else {
		// Simple expression, assign to v1
		fullScript = fmt.Sprintf(`%s
function process(...)
    local result = %s
    v1 = result
    return v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12
end
`, scriptHeader, luaExpr)
		log.Println("Generated Lua script for simple expression.")
	}

	log.Printf("Final Lua script: %s", fullScript)
	return fullScript
}

func processFile(filename string, pattern *regexp.Regexp, luaScript 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", fullPath)

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

	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", fullPath)

	return nil
}

func process(data string, pattern *regexp.Regexp, luaScript 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)
	}
	log.Println("Lua math library loaded successfully.")

	// Load the Lua script
	if err := L.DoString(luaScript); err != nil {
		Error.Printf("Error in Lua script: %v", err)
		return data, fmt.Errorf("error in Lua script: %v", err)
	}
	log.Println("Lua script executed successfully.")

	// 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 v%d to %f", i+1, floatVal)
			} else {
				L.SetGlobal(fmt.Sprintf("v%d", i+1), lua.LString(capture))
				log.Printf("Set v%d to %s", i+1, capture)
			}
		}

		// Call the Lua function
		L.Push(L.GetGlobal("process"))
		if err := L.PCall(0, 12, nil); err != nil { // We're returning up to 12 values
			Error.Printf("Error executing Lua script: %v", err)
			return match // Return unchanged on error
		}
		log.Println("Lua function 'process' called successfully.")

		// Get the return values (up to 12)
		returnValues := make([]lua.LValue, 12)
		for i := 0; i < 12; i++ {
			// Get the value from the stack (if exists)
			if L.GetTop() >= i+1 {
				returnValues[i] = L.Get(-12 + i)
			} else {
				returnValues[i] = lua.LNil
			}
		}
		// Pop all return values
		L.Pop(L.GetTop())

		// 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("Return value v%d is nil, skipping replacement.", i+1)
				continue // Skip if nil return
			}

			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", oldVal, newVal)
				result = strings.Replace(result, oldVal, newVal, 1)
			}
		}

		return result
	})

	log.Println("Processing completed successfully.")
	return result, nil
}