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
var Info *log.Logger
var Success *log.Logger

// ModificationRecord tracks a single value modification
type ModificationRecord struct {
	File      string
	OldValue  string
	NewValue  string
	Operation string
	Context   string
}

// GlobalStats tracks all modifications across files
type GlobalStats struct {
	TotalMatches       int
	TotalModifications int
	Modifications      []ModificationRecord
	ProcessedFiles     int
	FailedFiles        int
}

var stats GlobalStats

func init() {
	// Configure standard logging to be hidden by default
	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
	log.SetOutput(io.Discard) // Disable default logging to stdout

	// Set up custom loggers with different severity levels
	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)

	Info = log.New(os.Stdout,
		fmt.Sprintf("%sInfo:%s ", "\033[0;94m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)

	Success = log.New(os.Stdout,
		fmt.Sprintf("%s✨ SUCCESS:%s ", "\033[0;92m", "\033[0m"),
		log.Lmicroseconds|log.Lshortfile)

	// Initialize global stats
	stats = GlobalStats{
		Modifications: make([]ModificationRecord, 0),
	}
}

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 - need regex pattern, lua expression, and at least one file")
		flag.Usage()
		return
	}

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

	Info.Printf("Starting modifier with pattern '%s', expression '%s' on %d files", regexPattern, luaExpr, len(files))

	// Prepare the Lua expression
	originalLuaExpr := luaExpr
	luaExpr = buildLuaScript(luaExpr)
	if originalLuaExpr != luaExpr {
		Info.Printf("Transformed Lua expression from '%s' to '%s'", originalLuaExpr, luaExpr)
	}

	// Handle special pattern modifications
	originalPattern := regexPattern
	patternModified := false

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

	// Make sure the regex can match across multiple lines by adding (?s) flag
	if !strings.HasPrefix(regexPattern, "(?s)") {
		regexPattern = "(?s)" + regexPattern
		patternModified = true
	}

	if patternModified {
		Info.Printf("Modified regex pattern from '%s' to '%s'", originalPattern, regexPattern)
	}

	// Compile the pattern for file processing
	pattern, err := regexp.Compile(regexPattern)
	if err != nil {
		Error.Printf("Invalid regex pattern '%s': %v", regexPattern, err)
		return
	}

	// Process each file
	for _, file := range files {
		Info.Printf("🔄 Processing file: %s", file)
		err := processFile(file, pattern, luaExpr, originalLuaExpr)
		if err != nil {
			Error.Printf("❌ Failed to process file %s: %v", file, err)
			stats.FailedFiles++
		} else {
			Info.Printf("✅ Successfully processed file: %s", file)
			stats.ProcessedFiles++
		}
	}

	// Print summary of all modifications
	printSummary(originalLuaExpr)
}

// printSummary outputs a formatted summary of all modifications made
func printSummary(operation string) {
	if stats.TotalModifications == 0 {
		Warning.Printf("No modifications were made in any files")
		return
	}

	Success.Printf("Operation complete! Modified %d values in %d/%d files using '%s'",
		stats.TotalModifications, stats.ProcessedFiles, stats.ProcessedFiles+stats.FailedFiles, operation)

	// Group modifications by file for better readability
	fileGroups := make(map[string][]ModificationRecord)
	for _, mod := range stats.Modifications {
		fileGroups[mod.File] = append(fileGroups[mod.File], mod)
	}

	// Print modifications by file
	for file, mods := range fileGroups {
		Success.Printf("📄 %s: %d modifications", file, len(mods))

		// Show some sample modifications (max 5 per file to avoid overwhelming output)
		maxSamples := 5
		if len(mods) > maxSamples {
			for i := 0; i < maxSamples; i++ {
				mod := mods[i]
				Success.Printf("   %d. '%s' → '%s' %s",
					i+1, limitString(mod.OldValue, 20), limitString(mod.NewValue, 20), mod.Context)
			}
			Success.Printf("   ... and %d more modifications", len(mods)-maxSamples)
		} else {
			for i, mod := range mods {
				Success.Printf("   %d. '%s' → '%s' %s",
					i+1, limitString(mod.OldValue, 20), limitString(mod.NewValue, 20), mod.Context)
			}
		}
	}

	// Print a nice visual indicator of success
	if stats.TotalModifications > 0 {
		Success.Printf("🎉 All done! Modified %d values successfully!", stats.TotalModifications)
	}
}

// buildLuaScript creates a complete Lua script from the expression
func buildLuaScript(luaExpr string) string {
	// Track if we modified the expression
	modified := false
	original := luaExpr

	// Auto-prepend v1 for expressions starting with operators
	if strings.HasPrefix(luaExpr, "*") ||
		strings.HasPrefix(luaExpr, "/") ||
		strings.HasPrefix(luaExpr, "+") ||
		strings.HasPrefix(luaExpr, "-") ||
		strings.HasPrefix(luaExpr, "^") ||
		strings.HasPrefix(luaExpr, "%") {
		luaExpr = "v1 = v1" + luaExpr
		modified = true
	} else if strings.HasPrefix(luaExpr, "=") {
		// Handle direct assignment with = operator
		luaExpr = "v1 " + luaExpr
		modified = true
	}

	// Add assignment if needed
	if !strings.Contains(luaExpr, "=") {
		luaExpr = "v1 = " + luaExpr
		modified = true
	}

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

	if modified {
		Info.Printf("Transformed Lua expression: '%s' → '%s'", original, luaExpr)
	}

	return luaExpr
}

func processFile(filename string, pattern *regexp.Regexp, luaExpr string, originalExpr string) error {
	fullPath := filepath.Join(".", filename)

	// Read file content
	content, err := os.ReadFile(fullPath)
	if err != nil {
		Error.Printf("Cannot read file %s: %v", fullPath, err)
		return fmt.Errorf("error reading file: %v", err)
	}

	fileContent := string(content)
	Info.Printf("File %s loaded: %d bytes", fullPath, len(content))

	// Process the content
	result, modificationCount, matchCount, err := process(fileContent, pattern, luaExpr, filename, originalExpr)
	if err != nil {
		Error.Printf("Processing failed for %s: %v", fullPath, err)
		return err
	}

	// Update global stats
	stats.TotalMatches += matchCount
	stats.TotalModifications += modificationCount

	if modificationCount == 0 {
		Warning.Printf("No modifications made to %s - pattern didn't match any content", fullPath)
		return nil
	}

	// Write the modified content back
	err = os.WriteFile(fullPath, []byte(result), 0644)
	if err != nil {
		Error.Printf("Failed to save changes to %s: %v", fullPath, err)
		return fmt.Errorf("error writing file: %v", err)
	}

	Info.Printf("Made %d modifications to %s and saved (%d bytes)",
		modificationCount, fullPath, len(result))

	return nil
}

func process(data string, pattern *regexp.Regexp, luaExpr string, filename string, originalExpr string) (string, int, int, error) {
	L := lua.NewState()
	defer L.Close()

	// Initialize Lua environment
	modificationCount := 0
	matchCount := 0

	// Load math library
	L.Push(L.GetGlobal("require"))
	L.Push(lua.LString("math"))
	if err := L.PCall(1, 1, nil); err != nil {
		Error.Printf("Failed to load Lua math library: %v", err)
		return data, 0, 0, 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("Failed to load Lua helper functions: %v", err)
		return data, 0, 0, fmt.Errorf("error loading helper functions: %v", err)
	}

	// Process all regex matches
	result := pattern.ReplaceAllStringFunc(data, func(match string) string {
		matchCount++
		captures := pattern.FindStringSubmatch(match)
		if len(captures) <= 1 {
			// No capture groups, return unchanged
			Warning.Printf("Match found but no capture groups: %s", limitString(match, 50))
			return match
		}

		// Set up global variables v1, v2, etc. for the Lua context
		captureValues := make([]string, len(captures)-1)
		for i, capture := range captures[1:] {
			captureValues[i] = capture
			// 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))
			} else {
				L.SetGlobal(fmt.Sprintf("v%d", i+1), lua.LString(capture))
			}
		}

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

		// Get the modified values after Lua execution
		modifications := make(map[int]string)
		for i := 0; i < len(captures)-1 && i < 12; i++ {
			varName := fmt.Sprintf("v%d", i+1)
			luaVal := L.GetGlobal(varName)

			if luaVal == lua.LNil {
				continue // Skip if nil (no modification)
			}

			oldVal := captures[i+1]
			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)
			}

			// Record modification if the value actually changed
			if newVal != oldVal {
				modifications[i] = newVal
			}
		}

		// Apply modifications to the matched text
		if len(modifications) == 0 {
			return match // No changes
		}

		result := match
		for i, newVal := range modifications {
			oldVal := captures[i+1]
			result = strings.Replace(result, oldVal, newVal, 1)

			// Extract a bit of context from the match for better reporting
			contextStart := max(0, strings.Index(match, oldVal)-10)
			contextLength := min(30, len(match)-contextStart)
			if contextStart+contextLength > len(match) {
				contextLength = len(match) - contextStart
			}
			contextStr := "..." + match[contextStart:contextStart+contextLength] + "..."

			// Log the modification
			Info.Printf("Modified value [%d]: '%s' → '%s'", i+1, limitString(oldVal, 30), limitString(newVal, 30))

			// Record the modification for summary
			stats.Modifications = append(stats.Modifications, ModificationRecord{
				File:      filename,
				OldValue:  oldVal,
				NewValue:  newVal,
				Operation: originalExpr,
				Context:   fmt.Sprintf("(in %s)", limitString(contextStr, 30)),
			})
		}

		modificationCount++
		return result
	})

	return result, modificationCount, matchCount, nil
}

// limitString truncates a string to maxLen and adds "..." if truncated
func limitString(s string, maxLen int) string {
	s = strings.ReplaceAll(s, "\n", "\\n")
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen-3] + "..."
}

// max returns the maximum of two integers
func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// min returns the minimum of two integers
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}