From 97cf9c509307d2eac747fe24127f0a1984b16b81 Mon Sep 17 00:00:00 2001
From: Martin Michelsen <fuzziqersoftware+g@gmail.com>
Date: Sun, 4 Feb 2024 17:06:58 -0800
Subject: [PATCH] fix challenge overlays after retry on non-BB versions

---
 src/CommandFormats.hh     |  6 +++---
 src/ReceiveCommands.cc    |  4 ++++
 src/ReceiveSubcommands.cc | 28 +++++++++++++++++++++++++++-
 3 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/src/CommandFormats.hh b/src/CommandFormats.hh
index aefc3b6e..3927dc59 100644
--- a/src/CommandFormats.hh
+++ b/src/CommandFormats.hh
@@ -5229,12 +5229,12 @@ struct G_Unknown_6x95 {
 // 6x96: Unknown (not valid on Episode 3)
 // This command has a handler, but it does nothing.
 
-// 6x97: Unknown (not valid on Episode 3)
+// 6x97: Select Challenge Mode failure option (not valid on Episode 3)
 
-struct G_Unknown_6x97 {
+struct G_SelectChallengeModeFailureOption_6x97 {
   G_UnusedHeader header;
   le_uint32_t unused1 = 0;
-  le_uint32_t unknown_a1 = 0; // Must be 0 or 1
+  le_uint32_t is_retry = 0;
   le_uint32_t unused2 = 0;
   le_uint32_t unused3 = 0;
 } __packed__;
diff --git a/src/ReceiveCommands.cc b/src/ReceiveCommands.cc
index cb3c3df5..2576f836 100644
--- a/src/ReceiveCommands.cc
+++ b/src/ReceiveCommands.cc
@@ -3672,6 +3672,10 @@ static void on_DF_BB(shared_ptr<Client> c, uint16_t command, uint32_t, string& d
         throw runtime_error("challenge template index in quest metadata does not match index sent by client");
       }
 
+      for (auto& m : l->floor_item_managers) {
+        m.clear();
+      }
+
       for (auto lc : l->clients) {
         if (lc) {
           lc->use_default_bank();
diff --git a/src/ReceiveSubcommands.cc b/src/ReceiveSubcommands.cc
index 3e3b25ac..fe953bf8 100644
--- a/src/ReceiveSubcommands.cc
+++ b/src/ReceiveSubcommands.cc
@@ -3175,6 +3175,32 @@ static void on_request_challenge_grave_recovery_item_bb(shared_ptr<Client> c, ui
   }
 }
 
+static void on_challenge_mode_retry_or_quit(shared_ptr<Client> c, uint8_t command, uint8_t flag, void* data, size_t size) {
+  const auto& cmd = check_size_t<G_SelectChallengeModeFailureOption_6x97>(data, size);
+
+  auto l = c->require_lobby();
+  if ((cmd.is_retry == 1) && l->is_game() && l->quest && (l->quest->challenge_template_index >= 0)) {
+    auto s = l->require_server_state();
+
+    for (auto& m : l->floor_item_managers) {
+      m.clear();
+    }
+
+    for (auto lc : l->clients) {
+      if (lc) {
+        lc->use_default_bank();
+        lc->create_challenge_overlay(lc->version(), l->quest->challenge_template_index, s->level_table);
+        lc->log.info("Created challenge overlay");
+        l->assign_inventory_and_bank_item_ids(lc, true);
+      }
+    }
+
+    l->load_maps();
+  }
+
+  forward_subcommand(c, command, flag, data, size);
+}
+
 static void on_quest_exchange_item_bb(shared_ptr<Client> c, uint8_t, uint8_t, void* data, size_t size) {
   auto l = c->require_lobby();
   if (l->is_game() &&
@@ -3730,7 +3756,7 @@ const SubcommandDefinition subcommand_definitions[0x100] = {
     /* 6x94 */ {0x00, 0x00, 0x94, on_warp},
     /* 6x95 */ {0x00, 0x00, 0x95, on_forward_check_game},
     /* 6x96 */ {0x00, 0x00, 0x96, on_forward_check_game},
-    /* 6x97 */ {0x00, 0x00, 0x97, on_forward_check_game},
+    /* 6x97 */ {0x00, 0x00, 0x97, on_challenge_mode_retry_or_quit},
     /* 6x98 */ {0x00, 0x00, 0x98, on_forward_check_game},
     /* 6x99 */ {0x00, 0x00, 0x99, on_forward_check_game},
     /* 6x9A */ {0x00, 0x00, 0x9A, on_forward_check_game_client},