From 2cb9c4a533efa43ef849179e04cd4b6874363d17 Mon Sep 17 00:00:00 2001
From: Extremelyd1 <Extremelyd1@outlook.com>
Date: Sat, 11 May 2024 10:40:17 +0200
Subject: [PATCH] Fix UI softlock and navigation inconsistencies

---
 HKMP/Game/Client/Save/SaveManager.cs |  4 +++
 HKMP/Ui/UiManager.cs                 | 39 +++++++++++++++++++++-------
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/HKMP/Game/Client/Save/SaveManager.cs b/HKMP/Game/Client/Save/SaveManager.cs
index 2fe4279..1ea7100 100644
--- a/HKMP/Game/Client/Save/SaveManager.cs
+++ b/HKMP/Game/Client/Save/SaveManager.cs
@@ -584,6 +584,10 @@ private void UpdateSaveWithData(SaveUpdate saveUpdate) {
         UpdateSaveWithData(index, value);
     }
 
+    /// <summary>
+    /// Set the save data from the given CurrentSave by overriding all values.
+    /// </summary>
+    /// <param name="currentSave">The save data to set.</param>
     public void SetSaveWithData(CurrentSave currentSave) {
         Logger.Info("Received current save, updating...");
 
diff --git a/HKMP/Ui/UiManager.cs b/HKMP/Ui/UiManager.cs
index 920e835..606b8f1 100644
--- a/HKMP/Ui/UiManager.cs
+++ b/HKMP/Ui/UiManager.cs
@@ -312,7 +312,8 @@ public void ReturnToMainMenuFromGame() {
         UM.StartCoroutine(GM.ReturnToMainMenu(
             GameManager.ReturnToMainMenuSaveModes.DontSave,
             _ => {
-                UM.StartCoroutine(UM.HideCurrentMenu());
+                IH.StartUIInput();
+                UM.returnMainMenuPrompt.HighlightDefault();
             }
         ));
     }
@@ -347,8 +348,12 @@ private void TryAddMultiOption() {
             return;
         }
 
-        if (btnParent.FindGameObjectInChildren("StartMultiplayerButton") != null) {
+        var startMultiBtn = btnParent.FindGameObjectInChildren("StartMultiplayerButton");
+        if (startMultiBtn != null) {
             Logger.Info("Multiplayer button is already present");
+            
+            FixMultiplayerButtonNavigation(startMultiBtn);
+            
             return;
         }
 
@@ -358,7 +363,7 @@ private void TryAddMultiOption() {
             return;
         }
 
-        var startMultiBtn = Object.Instantiate(startGameBtn, btnParent.transform);
+        startMultiBtn = Object.Instantiate(startGameBtn, btnParent.transform);
         if (startMultiBtn == null) {
             Logger.Info("startMultiBtn is null");
             return;
@@ -370,9 +375,30 @@ private void TryAddMultiOption() {
         var autoLocalize = startMultiBtn.GetComponent<AutoLocalizeTextUI>();
         autoLocalize.textKey = "StartMultiplayerBtn";
         autoLocalize.RefreshTextFromLocalization();
+
+        var eventTrigger = startMultiBtn.GetComponent<EventTrigger>();
+        eventTrigger.triggers.Clear();
         
+        ChangeBtnTriggers(eventTrigger, () => UM.StartCoroutine(GoToMultiplayerMenu()));
+        
+        // Fix navigation now and when the IH can accept input again
+        FixMultiplayerButtonNavigation(startMultiBtn);
+        
+        UM.StartCoroutine(WaitForInput());
+        IEnumerator WaitForInput() {
+            yield return new WaitUntil(() => IH.acceptingInput);
+            
+            FixMultiplayerButtonNavigation(startMultiBtn);   
+        }
+    }
+    
+    /// <summary>
+    /// Fix the navigation for the multiplayer button that is added to the main menu.
+    /// </summary>
+    /// <param name="multiBtnObject">The game object for the multiplayer button.</param>
+    private void FixMultiplayerButtonNavigation(GameObject multiBtnObject) {
         // Fix navigation for buttons
-        var startMultiBtnMenuBtn = startMultiBtn.GetComponent<MenuButton>();
+        var startMultiBtnMenuBtn = multiBtnObject.GetComponent<MenuButton>();
         if (startMultiBtnMenuBtn != null) {
             var nav = UM.mainMenuButtons.startButton.navigation;
             nav.selectOnDown = startMultiBtnMenuBtn;
@@ -386,11 +412,6 @@ private void TryAddMultiOption() {
             nav.selectOnUp = UM.mainMenuButtons.startButton;
             startMultiBtnMenuBtn.navigation = nav;
         }
-
-        var eventTrigger = startMultiBtn.GetComponent<EventTrigger>();
-        eventTrigger.triggers.Clear();
-        
-        ChangeBtnTriggers(eventTrigger, () => UM.StartCoroutine(GoToMultiplayerMenu()));
     }
 
     /// <summary>