From 5bd7e057730bda4c1c2b8b29ac580ba90eee3ab2 Mon Sep 17 00:00:00 2001
From: Roman Trapeznikov <me@rom4ik.su>
Date: Mon, 8 Jan 2024 14:38:04 +0300
Subject: [PATCH] work on replay rewind

---
 .../ba_data/python/bascenev1/__init__.py      |  2 +
 .../python/methods/python_methods_scene.cc    | 29 ++++++++++++++
 .../scene_v1/support/client_session_replay.cc | 38 +++++++++++++++++++
 .../scene_v1/support/client_session_replay.h  | 16 ++++++++
 4 files changed, 85 insertions(+)

diff --git a/src/assets/ba_data/python/bascenev1/__init__.py b/src/assets/ba_data/python/bascenev1/__init__.py
index 25c1f5ef..44d7276d 100644
--- a/src/assets/ba_data/python/bascenev1/__init__.py
+++ b/src/assets/ba_data/python/bascenev1/__init__.py
@@ -120,6 +120,7 @@ from _bascenev1 import (
     release_keyboard_input,
     reset_random_player_names,
     resume_replay,
+    rewind_replay,
     broadcastmessage,
     SessionData,
     SessionPlayer,
@@ -400,6 +401,7 @@ __all__ = [
     'release_keyboard_input',
     'reset_random_player_names',
     'resume_replay',
+    'rewind_replay',
     'safecolor',
     'screenmessage',
     'SceneV1AppMode',
diff --git a/src/ballistica/scene_v1/python/methods/python_methods_scene.cc b/src/ballistica/scene_v1/python/methods/python_methods_scene.cc
index 5e74a112..389bfc70 100644
--- a/src/ballistica/scene_v1/python/methods/python_methods_scene.cc
+++ b/src/ballistica/scene_v1/python/methods/python_methods_scene.cc
@@ -1568,6 +1568,34 @@ static PyMethodDef PyResumeReplayDef = {
     "Resumes replay.",
 };
 
+// ------------------------ rewind_replay --------------------------------------
+
+static auto PyRewindReplay(PyObject* self, PyObject* args) -> PyObject* {
+  BA_PYTHON_TRY;
+  auto* appmode = SceneV1AppMode::GetActiveOrThrow();
+  auto* session =
+      dynamic_cast<ClientSessionReplay*>(appmode->GetForegroundSession());
+  if (session == nullptr) {
+    throw Exception(
+        "Attempting to rewind a replay not in replay session context.");
+  }
+  session->RestoreState();
+  Py_RETURN_NONE;
+  BA_PYTHON_CATCH;
+}
+
+static PyMethodDef PyRewindReplayDef = {
+    "rewind_replay",  // name
+    PyRewindReplay,   // method
+    METH_VARARGS,     // flags
+
+    "rewind_replay() -> None\n"
+    "\n"
+    "(internal)\n"
+    "\n"
+    "Rewinds replay.",
+};
+
 // ----------------------- reset_random_player_names ---------------------------
 
 static auto PyResetRandomPlayerNames(PyObject* self, PyObject* args,
@@ -1846,6 +1874,7 @@ auto PythonMethodsScene::GetMethods() -> std::vector<PyMethodDef> {
       PySetReplaySpeedExponentDef,
       PyGetReplaySpeedExponentDef,
       PyIsReplayPausedDef,
+      PyRewindReplayDef,
       PyPauseReplayDef,
       PyResumeReplayDef,
       PySetDebugSpeedExponentDef,
diff --git a/src/ballistica/scene_v1/support/client_session_replay.cc b/src/ballistica/scene_v1/support/client_session_replay.cc
index f5c63095..6fe35bd5 100644
--- a/src/ballistica/scene_v1/support/client_session_replay.cc
+++ b/src/ballistica/scene_v1/support/client_session_replay.cc
@@ -134,6 +134,17 @@ void ClientSessionReplay::FetchMessages() {
 
   // If we have no messages left, read from the file until we get some.
   while (commands().empty()) {
+    {
+      // Before we read next message, let's save our current state.
+      SessionStream out(nullptr, false);
+      DumpFullState(&out);
+
+      fflush(file_);
+      current_state_.file_position_ = ftell(file_);
+      current_state_.message_ = out.GetOutMessage();
+      SaveState();
+    }
+
     std::vector<uint8_t> buffer;
     uint8_t len8;
     uint32_t len32;
@@ -265,4 +276,31 @@ void ClientSessionReplay::OnReset(bool rewind) {
   }
 }
 
+void ClientSessionReplay::SaveState() { states_.push_back(current_state_); }
+
+void ClientSessionReplay::RestoreState() {
+  ScreenMessage("restoring state " + std::to_string(states_.size()));
+  const int N = 500;
+
+  if (states_.size() <= N) {
+    states_.clear();
+    Reset(true);
+    return;
+  }
+
+  for (int i = 0; i < N; ++i) {
+    states_.pop_back();
+  }
+
+  current_state_ = states_.back();
+  RestoreFromCurrentState();
+}
+
+void ClientSessionReplay::RestoreFromCurrentState() {
+  // what to do with messages_fetch_num_? is it used somewhere at all?
+  Reset(true);
+  fseek(file_, current_state_.file_position_, SEEK_SET);
+  HandleSessionMessage(current_state_.message_);
+}
+
 }  // namespace ballistica::scene_v1
diff --git a/src/ballistica/scene_v1/support/client_session_replay.h b/src/ballistica/scene_v1/support/client_session_replay.h
index 4bd82e82..446b70f3 100644
--- a/src/ballistica/scene_v1/support/client_session_replay.h
+++ b/src/ballistica/scene_v1/support/client_session_replay.h
@@ -28,8 +28,24 @@ class ClientSessionReplay : public ClientSession,
 
   void Error(const std::string& description) override;
   void FetchMessages() override;
+  void SaveState();
+  void RestoreState();
 
  private:
+  struct IntermediateState {
+    // Message containing full scene state at the moment.
+    std::vector<uint8_t> message_;
+
+    // A position in replay file where we should continue from.
+    long file_position_;
+  };
+
+  void RestoreFromCurrentState();
+
+  // List of passed states which we can rewind to.
+  std::vector<IntermediateState> states_;
+  IntermediateState current_state_;
+
   uint32_t message_fetch_num_{};
   bool have_sent_client_message_{};
   std::vector<ConnectionToClient*> connections_to_clients_;