wow-AchievementSniffer/main.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

	_ "embed"

	logger "git.site.quack-lab.dev/dave/cylogger"
	"github.com/bmatcuk/doublestar/v4"
	lua "github.com/yuin/gopher-lua"
)

const nsqEndpoint = "https://nsq.site.quack-lab.dev/pub?topic=wowspy"

var debug *bool
var nsqWorkers = 32

var allPlayersAchievementsGlobal = make(map[string][]NSQMessage) // PlayerName -> list of all their achievements
var allPlayerNamesGlobal = make(map[string]bool)                 // Set of all player names

func main() {
	root := flag.String("root", ".", "Root workdir")
	debug = flag.Bool("d", false, "Debug")
	flag.Parse()
	logger.InitFlag()
	if *debug {
		logger.SetLevel(logger.LevelDebug) // Assuming LevelDebug is the correct constant for cylogger
	}

	db := DB{
		path: "service/data/db.db",
	}
	err := db.Open()
	if err != nil {
		logger.Error("error opening database: %v", err)
		return
	}
	defer db.Close()

	logger.Info("Root: %q", *root)
	cleanedRoot := strings.Replace(*root, "~", os.Getenv("HOME"), 1)
	cleanedRoot, err = filepath.Abs(cleanedRoot)
	if err != nil {
		logger.Error("error getting absolute path: %v", err)
		return
	}
	cleanedRoot = filepath.Clean(cleanedRoot)
	cleanedRoot = strings.TrimSuffix(cleanedRoot, "/")

	logger.Info("Looking for Heimdall.lua in %q", cleanedRoot)
	matches, err := doublestar.Glob(os.DirFS(cleanedRoot), "**/Heimdall.lua")
	if err != nil {
		logger.Error("error matching Heimdall.lua: %v", err)
		return
	}
	logger.Info("Found %d Heimdall.lua files.", len(matches))
	if len(matches) == 0 {
		logger.Info("No Heimdall.lua files found. Exiting.")
		return
	}

	luaStates := loadLuaStates(matches)
	achievements := loadAchievements(luaStates)
	wg := sync.WaitGroup{}
	wg.Add(1)
	// We can save the achievements to the database while doing something else unrelated
	go saveAchievementsToDB(&db, achievements)
	saveAchievementsToSourceFiles(luaStates, achievements)
	wg.Wait()

	// --- Pass 1: Extract all data ---
	// logger.Info("Starting Pass 1: Extracting data from all Heimdall.lua files...")
	// var wgPass1 sync.WaitGroup
	// for _, match := range matches {
	// 	wgPass1.Add(1)
	// 	go loadAchievements(filepath.Join(cleanedRoot, match), &wgPass1)
	// }
	// wgPass1.Wait()
	// logger.Info("Finished Pass 1: Loaded %d unique players from %d files.", len(allPlayerNamesGlobal), len(matches))
	// if *debug {
	// 	globalDataMutex.Lock()
	// 	logger.Debug("Total achievements loaded globally: %d", countTotalAchievements(allPlayersAchievementsGlobal))
	// 	globalDataMutex.Unlock()
	// }

	// wgSave := sync.WaitGroup{}
	// wgSave.Add(1)
	// go func() {
	// 	logger.Info("Saving achievements to database...")
	// 	for playerName, achList := range allPlayersAchievementsGlobal {
	// 		logger.Debug("Saving %d achievements for player %s", len(achList), playerName)
	// 		for _, ach := range achList {
	// 			Save(&ach, &db)
	// 		}
	// 	}
	// 	wgSave.Done()
	// }()

	// // --- Process and Send to NSQ ---
	// // logger.Info("Starting NSQ message publishing...")
	// // nsqMessagesChan := make(chan NSQMessage, 10000) // Increased buffer size
	// // var wgNsqWorkers sync.WaitGroup
	// // for i := 0; i < nsqWorkers; i++ {
	// // 	wgNsqWorkers.Add(1)
	// // 	go NsqWorker(&wgNsqWorkers, nsqMessagesChan)
	// // }

	// // go func() {
	// // 	globalDataMutex.Lock()
	// // 	defer globalDataMutex.Unlock()
	// // 	for playerName, achList := range allPlayersAchievementsGlobal {
	// // 		for _, ach := range achList {
	// // 			// ach.Name is already correctly set during extraction
	// // 			nsqMessagesChan <- ach
	// // 			logger.Debug("Queued NSQ message for Player: %s, AchID: %s", playerName, ach.ID)
	// // 		}
	// // 	}
	// // 	close(nsqMessagesChan) // Close channel when all messages are sent
	// // 	logger.Info("All NSQ messages queued.")
	// // }()

	// // --- Pass 2: Update Lua file states (in memory) ---
	// logger.Info("Starting Pass 2: Updating Lua states (setting alreadySeen and clearing players)...")
	// var wgPass2 sync.WaitGroup
	// if len(allPlayerNamesGlobal) > 0 { // Only run pass 2 if there are players to report
	// 	for _, match := range matches {
	// 		wgPass2.Add(1)
	// 		go saveLuaFileState(filepath.Join(cleanedRoot, match), &wgPass2, allPlayerNamesGlobal)
	// 	}
	// 	wgPass2.Wait()
	// 	logger.Info("Finished Pass 2: Lua states updated where applicable.")
	// } else {
	// 	logger.Info("Skipping Pass 2 as no players were found globally.")
	// }

	// // wgNsqWorkers.Wait() // Wait for all NSQ messages to be processed
	// wgSave.Wait()
	// logger.Info("All NSQ workers finished. Program complete.")
}

func saveAchievementsToSourceFiles(luaStates *sync.Map, achievements *sync.Map) {
	wg := sync.WaitGroup{}
	luaStates.Range(func(k, v any) bool {
		path := k.(string)
		state := v.(*lua.LState)
		log := logger.Default.WithPrefix(path)

		log.Info("Clearing existing achievements")
		achievementTable := state.GetGlobal("Heimdall_Achievements")
		if achievementTable.Type() != lua.LTTable {
			achievementTable = &lua.LTable{}
			state.SetGlobal("Heimdall_Achievements", achievementTable)
		}
		log.Info("Clearing existing players table")
		emptyTable := &lua.LTable{}
		state.SetField(achievementTable, "players", emptyTable)

		log.Info("Updating seen table")
		seenTable := state.GetField(achievementTable, "alreadySeen")
		if seenTable.Type() != lua.LTTable {
			seenTable = &lua.LTable{}
			state.SetField(achievementTable, "alreadySeen", seenTable)
		}

		wg.Add(1)
		go func() {
			defer wg.Done()
			updateSeenTable(state, seenTable, achievements, log)
			// State itself should now be updated
		}()

		writeToSource(path, state)

		return true
	})
	wg.Wait()
}

func writeToSource(path string, state *lua.LState) {
	
}

func updateSeenTable(state *lua.LState, seenTable lua.LValue, achievements *sync.Map, log *logger.Logger) {
	achievements.Range(func(k, v any) bool {
		playerName := k.(string)
		state.SetField(seenTable, playerName, lua.LBool(true))
		return true
	})

	// Debug confirmation
	// seenTableLua, ok := seenTable.(*lua.LTable)
	// if !ok {
	// 	log.Error("seenTable is not a table")
	// 	return
	// }
	// state.ForEach(seenTableLua, func(k, v lua.LValue) {
	// 	log.Info("Seen: %s", k)
	// })
}

func saveAchievementsToDB(db *DB, achievements *sync.Map) {
	achievements.Range(func(k, v any) bool {
		playerName := k.(string)
		playerAchievements := v.(*[]NSQMessage)
		logger.Info("Saving %d achievements for player %s", len(*playerAchievements), playerName)
		for _, ach := range *playerAchievements {
			Save(ach, db)
		}
		return true
	})
}

func loadLuaStates(matches []string) *sync.Map {
	wg := sync.WaitGroup{}
	fileLuaStates := &sync.Map{}
	for _, match := range matches {
		wg.Add(1)
		go func(path string) {
			defer wg.Done()
			log := logger.Default.WithPrefix(path)
			L := lua.NewState()

			filestat, err := os.Stat(match)
			if err != nil {
				log.Error("error getting file stats: %v", err)
				return
			}
			log.Info("File size: %.2f MB", float64(filestat.Size())/1024/1024)

			log.Info("Running Lua file")
			if err := L.DoFile(path); err != nil {
				log.Error("error executing Lua file %q: %v", path, err)
				return
			}
			log.Info("Lua file loaded")
			fileLuaStates.Store(match, L)
		}(match)
	}
	wg.Wait()
	return fileLuaStates
}

func loadAchievements(luaStates *sync.Map) *sync.Map {
	achievements := &sync.Map{}
	wg := sync.WaitGroup{}
	luaStates.Range(func(path, state any) bool {
		wg.Add(1)
		go func(path string, state *lua.LState) {
			log := logger.Default.WithPrefix(path)
			defer wg.Done()
			// We directly mutate achievements to avoid reducing and mapping later on
			// Removing 1 off of the x of the O(xn)
			loadStateAchievements(state, log, achievements)
		}(path.(string), state.(*lua.LState))
		return true
	})
	wg.Wait()
	return achievements
}
func loadStateAchievements(L *lua.LState, log *logger.Logger, achievements *sync.Map) {
	log.Info("Getting Heimdall_Achievements")
	heimdallAchievements := L.GetGlobal("Heimdall_Achievements")
	if heimdallAchievements.Type() == lua.LTNil {
		log.Warning("Heimdall_Achievements not found. Skipping file.")
		return
	}

	log.Info("Getting players table")
	playersTableLua := L.GetField(heimdallAchievements, "players")
	if playersTableLua.Type() == lua.LTNil {
		log.Info("'players' table is nil in Heimdall_Achievements. No player data to extract.")
		return
	}
	log.Info("Casting players table")
	playersTable, ok := playersTableLua.(*lua.LTable)
	if !ok {
		log.Warning("'players' field in Heimdall_Achievements is not a table. Skipping.")
		return
	}

	log.Info("Iterating over players")
	counter := 0
	playersTable.ForEach(func(playerNameLua lua.LValue, playerAchievementsLua lua.LValue) {
		currentPlayerName := playerNameLua.String()
		playerAchievements, _ := achievements.LoadOrStore(currentPlayerName, &[]NSQMessage{})
		playerAchievementsSlice := playerAchievements.(*[]NSQMessage)

		achievementsTableLua, ok := playerAchievementsLua.(*lua.LTable)
		if !ok {
			log.Error("Achievements for player %s is not a table. Skipping achievements for this player.", currentPlayerName)
			return
		}

		achievementsTableLua.ForEach(func(_ lua.LValue, achievementDataLua lua.LValue) {
			achievementTable, ok := achievementDataLua.(*lua.LTable)
			if !ok {
				log.Error("Achievement data for player %s is not a table. Skipping this achievement.", currentPlayerName)
				return
			}

			currentAchievement := NSQMessage{Name: currentPlayerName}

			idVal := achievementTable.RawGetString("id")
			if idVal.Type() == lua.LTNumber {
				currentAchievement.ID = lua.LVAsString(idVal)
			} else if idVal.Type() == lua.LTString {
				currentAchievement.ID = idVal.String()
			} else {
				log.Warning("Missing or invalid 'id' (expected number or string) for achievement for player %s.", currentPlayerName)
			}

			dateVal := achievementTable.RawGetString("date")
			if dateVal.Type() == lua.LTString {
				currentAchievement.Date = dateVal.String()
			} else {
				log.Warning("Missing or invalid 'date' (expected string) for achievement for player %s.", currentPlayerName)
			}

			completedVal := achievementTable.RawGetString("completed")
			if completedVal.Type() == lua.LTBool {
				currentAchievement.Completed = lua.LVAsBool(completedVal)
			} else {
				log.Warning("Missing or invalid 'completed' (expected boolean) for achievement for player %s.", currentPlayerName)
			}

			if currentAchievement.ID != "" { // Ensure we have at least an ID before adding
				// Will this change be reflected in the map...?
				*playerAchievementsSlice = append(*playerAchievementsSlice, currentAchievement)
			}

			counter++
			if counter%2000 == 0 {
				log.Info("Processed %d achievements", counter)
			}
		})
	})
	log.Info("Processed %d achievements", counter)
	achievements.Range(func(key, value any) bool {
		log.Trace("Player: %s, Achievements: %d", key, len(*value.(*[]NSQMessage)))
		return true
	})
}

// updateLuaFileState is for Pass 2
func NsqWorker(wg *sync.WaitGroup, messages <-chan NSQMessage) { // Changed to read-only channel
	defer wg.Done()
	for msg := range messages {
		err := Publish(msg)
		if err != nil {
			logger.Warning("error publishing message for player %s, achievement ID %s: %v", msg.Name, msg.ID, err)
			// Optionally, add retry logic or dead-letter queue here
			continue
		}
		logger.Debug("Successfully published achievement for %s: ID %s", msg.Name, msg.ID)
	}
}

func Publish(msg NSQMessage) error {
	data := bytes.Buffer{}
	err := json.NewEncoder(&data).Encode(msg)
	if err != nil {
		return err // Error encoding JSON
	}
	resp, err := http.Post(nsqEndpoint, "application/json", &data)
	if err != nil {
		return err // Error making HTTP POST request
	}
	defer resp.Body.Close() // Ensure body is closed

	if resp.StatusCode != http.StatusOK {
		// Read body for more details if not OK
		bodyBytes, readErr := io.ReadAll(resp.Body)
		if readErr != nil {
			logger.Warning("Failed to read error response body from NSQ: %v", readErr)
		}
		return fmt.Errorf("nsq publish failed: status code %d, body: %s", resp.StatusCode, string(bodyBytes))
	}
	return nil
}

func saveLuaFileState(path string, wg *sync.WaitGroup, allKnownPlayerNames map[string]bool) {
	log := logger.Default.WithPrefix(path)
	log.Info("Saving Lua state")
	defer wg.Done()
}

// writeLuaTable writes a Lua table back to disk using Lua's own serialization
func writeLuaTable(path string, table *lua.LTable) error {
	L := lua.NewState()
	defer L.Close()

	// Create a Lua function that will write our table
	script := `
		local function writeTable(tbl)
			local file = io.open("%s", "w")
			if not file then
				return false, "Could not open file for writing"
			end

			-- Write the original file content up to our table
			local original = io.open("%s", "r")
			if original then
				local content = original:read("*all")
				original:close()

				-- Find where our table starts
				local startPos = content:find("Heimdall_Achievements%s*=%s*{")
				if startPos then
					file:write(content:sub(1, startPos-1))
				end
			end

			-- Write our table
			file:write("Heimdall_Achievements = ")

			-- Use Lua's built-in table serialization
			local function serialize(tbl, indent)
				indent = indent or 0
				local str = "{\n"
				for k, v in pairs(tbl) do
					str = str .. string.rep("\t", indent + 1)
					if type(k) == "string" then
						str = str .. string.format('["%s"] = ', k)
					else
						str = str .. string.format("[%s] = ", k)
					end

					if type(v) == "table" then
						str = str .. serialize(v, indent + 1)
					elseif type(v) == "string" then
						str = str .. string.format('"%s"', v)
					elseif type(v) == "boolean" then
						str = str .. tostring(v)
					else
						str = str .. tostring(v)
					end
					str = str .. ",\n"
				end
				str = str .. string.rep("\t", indent) .. "}"
				return str
			end

			file:write(serialize(tbl))
			file:write("\n")

			-- Write the rest of the file
			if original then
				local content = original:read("*all")
				original:close()

				-- Find where our table ends
				local endPos = content:find("}%s*$")
				if endPos then
					file:write(content:sub(endPos))
				end
			end

			file:close()
			return true
		end

		return writeTable(Heimdall_Achievements)
	`

	// Set our table in the Lua state
	L.SetGlobal("Heimdall_Achievements", table)

	// Execute the script
	if err := L.DoString(script); err != nil {
		return fmt.Errorf("failed to write Lua table: %v", err)
	}

	return nil
}