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
var globalDataMutex = &sync.Mutex{}

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
	}

	// matches = matches[:1]

	// --- 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 extractPlayerAchievementsFromFile(filepath.Join(cleanedRoot, match), &wgPass1)
	}
	wgPass1.Wait()
	logger.Info("Finished Pass 1: Extracted data for %d unique players globally from %d files.", len(allPlayerNamesGlobal), len(matches))
	if *debug {
		globalDataMutex.Lock()
		logger.Debug("Total achievements collected 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 updateLuaFileState(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.")
}

// Helper function to count total achievements for debugging
func countTotalAchievements(achMap map[string][]NSQMessage) int {
	count := 0
	for _, achList := range achMap {
		count += len(achList)
	}
	return count
}

// extractPlayerAchievementsFromFile is for Pass 1
func extractPlayerAchievementsFromFile(path string, wg *sync.WaitGroup) {
	log := logger.Default.WithPrefix(path)
	log.Info("Extracting achievements")
	defer wg.Done()
	L := lua.NewState()
	defer L.Close()

	filestat, err := os.Stat(path)
	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("Getting Heimdall_Achievements")
	heimdallAchievements := L.GetGlobal("Heimdall_Achievements")
	if heimdallAchievements.Type() == lua.LTNil {
		log.Warning("Heimdall_Achievements not found in %q. Skipping file.", path)
		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 in %q. No player data to extract.", path)
		return
	}
	log.Info("Casting players table")
	playersTable, ok := playersTableLua.(*lua.LTable)
	if !ok {
		log.Warning("'players' field in Heimdall_Achievements is not a table in %q (type: %s). Skipping.", path, playersTableLua.Type().String())
		return
	}

	var filePlayerAchievements []NSQMessage
	var filePlayerNames = make(map[string]bool)

	log.Info("Iterating over players")
	counter := 0
	playersTable.ForEach(func(playerNameLua lua.LValue, playerAchievementsLua lua.LValue) {
		currentPlayerName := playerNameLua.String()
		filePlayerNames[currentPlayerName] = true // Track name

		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
				filePlayerAchievements = append(filePlayerAchievements, currentAchievement)
			}

			counter++
			if counter%2000 == 0 {
				log.Info("Processed %d achievements", counter)
			}
		})
	})
	log.Info("Processed %d achievements", counter)

	if len(filePlayerAchievements) > 0 || len(filePlayerNames) > 0 {
		globalDataMutex.Lock()
		for _, ach := range filePlayerAchievements {
			allPlayersAchievementsGlobal[ach.Name] = append(allPlayersAchievementsGlobal[ach.Name], ach)
		}
		for name := range filePlayerNames {
			allPlayerNamesGlobal[name] = true
		}
		globalDataMutex.Unlock()
		log.Info("Players in file: %d. Achievements in file: %d.", len(filePlayerNames), len(filePlayerAchievements))
	} else {
		log.Info("No player data or names extracted")
	}
}

// updateLuaFileState is for Pass 2
func updateLuaFileState(path string, wg *sync.WaitGroup, allKnownPlayerNames map[string]bool) {
	log := logger.Default.WithPrefix(filepath.Base(path))
	log.Info("Updating Lua state")
	defer wg.Done()
	L := lua.NewState()
	defer L.Close()

	if err := L.DoFile(path); err != nil {
		log.Error("error executing Lua file %q: %v. Cannot update its state.", path, err)
		return
	}

	heimdallAchievementsVal := L.GetGlobal("Heimdall_Achievements")
	if heimdallAchievementsVal.Type() == lua.LTNil {
		log.Warning("Heimdall_Achievements not found in %q after script execution. Cannot set 'alreadySeen' or clear 'players'.", path)
		return
	}

	heimdallAchievementsTable, ok := heimdallAchievementsVal.(*lua.LTable)
	if !ok {
		log.Warning("Heimdall_Achievements in %q is not a table (type: %s). Cannot update.", path, heimdallAchievementsVal.Type().String())
		return
	}

	luaAlreadySeen := L.NewTable()
	for name := range allKnownPlayerNames {
		luaAlreadySeen.RawSetString(name, lua.LTrue)
	}

	L.SetField(heimdallAchievementsTable, "alreadySeen", luaAlreadySeen)
	log.Debug("Set Heimdall_Achievements.alreadySeen for %q with %d total player names.", path, len(allKnownPlayerNames))

	L.SetField(heimdallAchievementsTable, "players", L.NewTable())
	log.Debug("Cleared Heimdall_Achievements.players for %q.", path)
}

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
}