From ad02bf2ee0c4d0bc8437069ada3d19fd0ab156dc Mon Sep 17 00:00:00 2001
From: EliteMasterEric <ericmyllyoja@gmail.com>
Date: Tue, 19 Dec 2023 01:27:58 -0500
Subject: [PATCH] Fix to GameOverSubstate exiting to Freeplay instead of Chart
 Editor

---
 source/funkin/play/GameOverSubState.hx     | 24 ++++++++++++++++++++--
 source/funkin/play/PlayState.hx            |  5 ++++-
 source/funkin/ui/freeplay/FreeplayState.hx |  2 +-
 source/funkin/ui/story/StoryMenuState.hx   |  2 +-
 source/funkin/util/Constants.hx            |  1 +
 5 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/source/funkin/play/GameOverSubState.hx b/source/funkin/play/GameOverSubState.hx
index 6eb53e2d5..18e3e0280 100644
--- a/source/funkin/play/GameOverSubState.hx
+++ b/source/funkin/play/GameOverSubState.hx
@@ -64,9 +64,13 @@ class GameOverSubState extends MusicBeatSubState
    */
   var isEnding:Bool = false;
 
-  public function new()
+  var isChartingMode:Bool = false;
+
+  public function new(?params:GameOverParams)
   {
     super();
+
+    this.isChartingMode = params?.isChartingMode ?? false;
   }
 
   /**
@@ -176,9 +180,20 @@ class GameOverSubState extends MusicBeatSubState
       // PlayState.seenCutscene = false; // old thing...
       gameOverMusic.stop();
 
-      if (PlayStatePlaylist.isStoryMode) FlxG.switchState(new StoryMenuState());
+      if (isChartingMode)
+      {
+        this.close();
+        if (FlxG.sound.music != null) FlxG.sound.music.pause(); // Don't reset song position!
+        PlayState.instance.close(); // This only works because PlayState is a substate!
+      }
+      else if (PlayStatePlaylist.isStoryMode)
+      {
+        FlxG.switchState(new StoryMenuState());
+      }
       else
+      {
         FlxG.switchState(new FreeplayState());
+      }
     }
 
     if (gameOverMusic.playing)
@@ -307,3 +322,8 @@ class GameOverSubState extends MusicBeatSubState
     });
   }
 }
+
+typedef GameOverParams =
+{
+  var isChartingMode:Bool;
+}
diff --git a/source/funkin/play/PlayState.hx b/source/funkin/play/PlayState.hx
index e0932e756..c834e0abe 100644
--- a/source/funkin/play/PlayState.hx
+++ b/source/funkin/play/PlayState.hx
@@ -922,7 +922,10 @@ class PlayState extends MusicBeatSubState
         }
         #end
 
-        var gameOverSubState = new GameOverSubState();
+        var gameOverSubState = new GameOverSubState(
+          {
+            isChartingMode: isChartingMode
+          });
         FlxTransitionableSubState.skipNextTransIn = true;
         FlxTransitionableSubState.skipNextTransOut = true;
         openSubState(gameOverSubState);
diff --git a/source/funkin/ui/freeplay/FreeplayState.hx b/source/funkin/ui/freeplay/FreeplayState.hx
index 7c69804d9..f17c3d91e 100644
--- a/source/funkin/ui/freeplay/FreeplayState.hx
+++ b/source/funkin/ui/freeplay/FreeplayState.hx
@@ -97,7 +97,7 @@ class FreeplayState extends MusicBeatSubState
   var stickerSubState:StickerSubState;
 
   //
-  static var rememberedDifficulty:Null<String> = "normal";
+  static var rememberedDifficulty:Null<String> = Constants.DEFAULT_DIFFICULTY;
   static var rememberedSongId:Null<String> = null;
 
   public function new(?stickers:StickerSubState = null)
diff --git a/source/funkin/ui/story/StoryMenuState.hx b/source/funkin/ui/story/StoryMenuState.hx
index 456988873..6e4cdacaf 100644
--- a/source/funkin/ui/story/StoryMenuState.hx
+++ b/source/funkin/ui/story/StoryMenuState.hx
@@ -106,7 +106,7 @@ class StoryMenuState extends MusicBeatState
   var stickerSubState:StickerSubState;
 
   static var rememberedLevelId:Null<String> = null;
-  static var rememberedDifficulty:Null<String> = "normal";
+  static var rememberedDifficulty:Null<String> = Constants.DEFAULT_DIFFICULTY;
 
   public function new(?stickers:StickerSubState = null)
   {
diff --git a/source/funkin/util/Constants.hx b/source/funkin/util/Constants.hx
index f8749567b..123267a49 100644
--- a/source/funkin/util/Constants.hx
+++ b/source/funkin/util/Constants.hx
@@ -123,6 +123,7 @@ class Constants
 
   /**
    * Default list of difficulties for charts.
+   * Assumes no Erect mode, etc.
    */
   public static final DEFAULT_DIFFICULTY_LIST:Array<String> = ['easy', 'normal', 'hard'];