Refactor a large part of main to be less retarded

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")
@@ -69,135 +68,153 @@ func main() {
 		return
 	}
 
-	// matches = matches[:1]
+	luaStates := loadLuaStates(matches)
+	achievements := loadAchievements(luaStates)
 
 	// --- Pass 1: Extract all data ---
-	logger.Info("Starting Pass 1: Extracting data from all Heimdall.lua files...")
+	// logger.Info("Starting Pass 1: Extracting data from all Heimdall.lua files...")
-	var wgPass1 sync.WaitGroup
+	// var wgPass1 sync.WaitGroup
-	for _, match := range matches {
+	// for _, match := range matches {
-		wgPass1.Add(1)
+	// 	wgPass1.Add(1)
-		go loadAchievements(filepath.Join(cleanedRoot, match), &wgPass1)
+	// 	go loadAchievements(filepath.Join(cleanedRoot, match), &wgPass1)
-	}
+	// }
-	wgPass1.Wait()
+	// wgPass1.Wait()
-	logger.Info("Finished Pass 1: Loaded %d unique players from %d files.", len(allPlayerNamesGlobal), len(matches))
+	// logger.Info("Finished Pass 1: Loaded %d unique players from %d files.", len(allPlayerNamesGlobal), len(matches))
-	if *debug {
+	// if *debug {
-		globalDataMutex.Lock()
+	// 	globalDataMutex.Lock()
-		logger.Debug("Total achievements loaded globally: %d", countTotalAchievements(allPlayersAchievementsGlobal))
+	// 	logger.Debug("Total achievements loaded globally: %d", countTotalAchievements(allPlayersAchievementsGlobal))
-		globalDataMutex.Unlock()
+	// 	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)
 	// }
 
+	// wgSave := sync.WaitGroup{}
+	// wgSave.Add(1)
 	// go func() {
-	// 	globalDataMutex.Lock()
+	// 	logger.Info("Saving achievements to database...")
-	// 	defer globalDataMutex.Unlock()
 	// 	for playerName, achList := range allPlayersAchievementsGlobal {
+	// 		logger.Debug("Saving %d achievements for player %s", len(achList), playerName)
 	// 		for _, ach := range achList {
-	// 			// ach.Name is already correctly set during extraction
+	// 			Save(&ach, &db)
-	// 			nsqMessagesChan <- ach
-	// 			logger.Debug("Queued NSQ message for Player: %s, AchID: %s", playerName, ach.ID)
 	// 		}
 	// 	}
-	// 	close(nsqMessagesChan) // Close channel when all messages are sent
+	// 	wgSave.Done()
-	// 	logger.Info("All NSQ messages queued.")
 	// }()
 
-	// --- Pass 2: Update Lua file states (in memory) ---
+	// // --- Process and Send to NSQ ---
-	logger.Info("Starting Pass 2: Updating Lua states (setting alreadySeen and clearing players)...")
+	// // logger.Info("Starting NSQ message publishing...")
-	var wgPass2 sync.WaitGroup
+	// // nsqMessagesChan := make(chan NSQMessage, 10000) // Increased buffer size
-	if len(allPlayerNamesGlobal) > 0 { // Only run pass 2 if there are players to report
+	// // var wgNsqWorkers sync.WaitGroup
-		for _, match := range matches {
+	// // for i := 0; i < nsqWorkers; i++ {
-			wgPass2.Add(1)
+	// // 	wgNsqWorkers.Add(1)
-			go updateLuaFileState(filepath.Join(cleanedRoot, match), &wgPass2, allPlayerNamesGlobal)
+	// // 	go NsqWorker(&wgNsqWorkers, nsqMessagesChan)
-		}
+	// // }
-		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
+	// // go func() {
-	wgSave.Wait()
+	// // 	globalDataMutex.Lock()
-	logger.Info("All NSQ workers finished. Program complete.")
+	// // 	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.")
 }
 
-// Helper function to count total achievements for debugging
+func loadLuaStates(matches []string) *sync.Map {
-func countTotalAchievements(achMap map[string][]NSQMessage) int {
+	wg := sync.WaitGroup{}
-	count := 0
+	fileLuaStates := &sync.Map{}
-	for _, achList := range achMap {
+	for _, match := range matches {
-		count += len(achList)
+		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)
 	}
-	return count
+	wg.Wait()
+	return fileLuaStates
 }
 
-func loadAchievements(path string, wg *sync.WaitGroup) {
+func loadAchievements(luaStates *sync.Map) *sync.Map {
-	log := logger.Default.WithPrefix(path)
+	achievements := &sync.Map{}
-	log.Info("Extracting achievements")
+	wg := sync.WaitGroup{}
-	defer wg.Done()
+	luaStates.Range(func(path, state any) bool {
-	L := lua.NewState()
+		wg.Add(1)
-	defer L.Close()
+		go func(path string, state *lua.LState) {
-
+			log := logger.Default.WithPrefix(path)
-	filestat, err := os.Stat(path)
+			defer wg.Done()
-	if err != nil {
+			// We directly mutate achievements to avoid reducing and mapping later on
-		log.Error("error getting file stats: %v", err)
+			// Removing 1 off of the x of the O(xn)
-		return
+			loadStateAchievements(state, log, achievements)
-	}
+		}(path.(string), state.(*lua.LState))
-	log.Info("File size: %.2f MB", float64(filestat.Size())/1024/1024)
+		return true
-
+	})
-	log.Info("Running Lua file")
+	wg.Wait()
-	if err := L.DoFile(path); err != nil {
+	return achievements
-		log.Error("error executing Lua file %q: %v", path, err)
+}
-		return
+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 +255,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 +266,13 @@ func loadAchievements(path string, wg *sync.WaitGroup) {
 		})
 	})
 	log.Info("Processed %d achievements", counter)
-
+	achievements.Range(func(key, value any) bool {
-	if len(filePlayerAchievements) > 0 || len(filePlayerNames) > 0 {
+		log.Trace("Player: %s, Achievements: %d", key, len(*value.(*[]NSQMessage)))
-		globalDataMutex.Lock()
+		return true
-		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 {
@@ -336,3 +308,98 @@ func Publish(msg NSQMessage) error {
 	}
 	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
+}