Update file paths for matches and change logging level for achievement saving

dbwriter.go

@@ -65,7 +65,7 @@ var whitelistedAchievements = map[string]bool{
 	"12448": true,
 }
 
-func Save(message *NSQMessage, db *DB) error {
+func Save(message NSQMessage, db *DB) error {
 	_, ok := whitelistedAchievements[message.ID]
 	if !ok {
 		logger.Debug("Received message for non-whitelisted achievement %s", message.ID)

main.go

@@ -26,7 +26,6 @@ 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")
@@ -68,136 +67,215 @@ func main() {
 		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 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()
+	for i, match := range matches {
+		matches[i] = filepath.Join(cleanedRoot, match)
 	}
 
-	wgSave := sync.WaitGroup{}
-	wgSave.Add(1)
+	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() {
-		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()
+		defer wg.Done()
+		saveAchievementsToDB(&db, achievements)
 	}()
+	saveAchievementsToSourceFiles(luaStates, achievements)
+	wg.Wait()
+}
 
-	// --- 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)
-	// }
+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)
 
-	// 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)
+		log.Info("Clearing existing achievements")
+		achievementTable := state.GetGlobal("Heimdall_Achievements")
+		if achievementTable.Type() != lua.LTTable {
+			achievementTable = &lua.LTable{}
+			state.SetGlobal("Heimdall_Achievements", achievementTable)
 		}
-		wgPass2.Wait()
-		logger.Info("Finished Pass 2: Lua states updated where applicable.")
-	} else {
-		logger.Info("Skipping Pass 2 as no players were found globally.")
-	}
+		log.Info("Clearing existing players table")
+		emptyTable := &lua.LTable{}
+		state.SetField(achievementTable, "players", emptyTable)
 
-	// wgNsqWorkers.Wait() // Wait for all NSQ messages to be processed
-	wgSave.Wait()
-	logger.Info("All NSQ workers finished. Program complete.")
+		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()
 }
 
-// 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
-}
-
-func loadAchievements(path string, wg *sync.WaitGroup) {
+func fixSource(path string, achievements *sync.Map) {
 	log := logger.Default.WithPrefix(path)
-	log.Info("Extracting achievements")
-	defer wg.Done()
-	L := lua.NewState()
-	defer L.Close()
-
-	filestat, err := os.Stat(path)
+	log.Info("Reading source file")
+	fileContent, err := os.ReadFile(path)
 	if err != nil {
-		log.Error("error getting file stats: %v", err)
+		logger.Error("Failed to read file: %v", err)
 		return
 	}
-	log.Info("File size: %.2f MB", float64(filestat.Size())/1024/1024)
+	strContent := string(fileContent)
+	log.Info("Read %d bytes", len(strContent))
 
-	log.Info("Running Lua file")
-	if err := L.DoFile(path); err != nil {
-		log.Error("error executing Lua file %q: %v", path, err)
+	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 in %q. Skipping file.", path)
+		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 in %q. No player data to extract.", path)
+		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 in %q (type: %s). Skipping.", path, playersTableLua.Type().String())
+		log.Warning("'players' field in Heimdall_Achievements is not a table. Skipping.")
 		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
+		playerAchievements, _ := achievements.LoadOrStore(currentPlayerName, &[]NSQMessage{})
+		playerAchievementsSlice := playerAchievements.(*[]NSQMessage)
 
 		achievementsTableLua, ok := playerAchievementsLua.(*lua.LTable)
 		if !ok {
@@ -238,7 +316,8 @@ func loadAchievements(path string, wg *sync.WaitGroup) {
 			}
 
 			if currentAchievement.ID != "" { // Ensure we have at least an ID before adding
-				filePlayerAchievements = append(filePlayerAchievements, currentAchievement)
+				// Will this change be reflected in the map...?
+				*playerAchievementsSlice = append(*playerAchievementsSlice, currentAchievement)
 			}
 
 			counter++
@@ -248,59 +327,13 @@ func loadAchievements(path string, wg *sync.WaitGroup) {
 		})
 	})
 	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")
-	}
+	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 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 {