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
	}
	for i, match := range matches {
		matches[i] = filepath.Join(cleanedRoot, match)
	}

	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 func() {
		defer wg.Done()
		saveAchievementsToDB(&db, achievements)
	}()
	saveAchievementsToSourceFiles(luaStates, achievements)
	wg.Wait()
}

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

		fixSource(path, achievements)
		return true
	})
	wg.Wait()
}

func fixSource(path string, achievements *sync.Map) {
	log := logger.Default.WithPrefix(path)
	log.Info("Reading source file")
	fileContent, err := os.ReadFile(path)
	if err != nil {
		logger.Error("Failed to read file: %v", err)
		return
	}
	strContent := string(fileContent)
	log.Info("Read %d bytes", len(strContent))

	strContent = removeAchievements(strContent)
	log.Info("Removed achievements, now %d bytes", len(strContent))

	strContent = addAlreadySeen(strContent, achievements)
	log.Info("Added alreadySeen, now %d bytes", len(strContent))

	log.Info("Writing file")
	err = os.WriteFile(path, []byte(strContent), 0644)
	if err != nil {
		logger.Error("Failed to write file: %v", err)
		return
	}
	log.Info("Done")
}

func removeAchievements(sourceContent string) string {
	lines := strings.Split(sourceContent, "\n")
	writeIndex := 0
	isInPlayers := false
	for _, line := range lines {
		if strings.HasPrefix(line, "\t[\"players\"] = {") {
			isInPlayers = true
			lines[writeIndex] = line
			writeIndex++
			continue
		}
		if isInPlayers && strings.HasPrefix(line, "\t}") {
			isInPlayers = false
			lines[writeIndex] = line
			writeIndex++
			continue
		}
		if !isInPlayers {
			lines[writeIndex] = line
			writeIndex++
		}
	}
	return strings.Join(lines[:writeIndex], "\n")
}

func addAlreadySeen(strContent string, achievements *sync.Map) string {
	lines := strings.Split(strContent, "\n")
	modifiedLines := make([]string, 0, len(lines))
	for _, line := range lines {
		if strings.HasPrefix(line, "\t[\"alreadySeen\"] = {") {
			modifiedLines = append(modifiedLines, line)
			achievements.Range(func(k, v any) bool {
				logger.Trace("Adding alreadySeen for %s", k)
				playerName := k.(string)
				modifiedLines = append(modifiedLines, fmt.Sprintf("\t\t[\"%s\"] = true,", playerName))
				return true
			})
			continue
		}
		modifiedLines = append(modifiedLines, line)
	}
	return strings.Join(modifiedLines, "\n")
}

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

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
}