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+)\" \"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) < 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)
		if err != nil {
			Error.Printf("error reading file: %v", err)
			continue
		}

		fileContent := string(content)
		result, err := process(fileContent, pattern, fullScript)
		if err != nil {
			Error.Printf("Error processing file %s: %v", file, err)
		}

		err = os.WriteFile(fullPath, []byte(result), 0644)
		if err != nil {
			Error.Printf("error writing file: %v", err)
		}
	}
}

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 {
		return data, fmt.Errorf("error loading Lua math library: %v", err)
	}

	// Load the Lua script
	if err := L.DoString(luaScript); err != nil {
		return data, fmt.Errorf("error in Lua script: %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
			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)
			}
			return result
		}

		// If something went wrong, return the original match
		return match
	})

	return result, nil
}