From fc440bea2abd467c2275089c813139277436c5b5 Mon Sep 17 00:00:00 2001
From: Lion Kortlepel <development@kortlepel.com>
Date: Mon, 29 Nov 2021 00:37:00 +0100
Subject: [PATCH] Add ability to switch into other lua states

---
 include/TConsole.h   |  5 +++--
 include/TLuaEngine.h |  1 +
 src/TConsole.cpp     | 33 +++++++++++++++++++++++++++------
 src/TLuaEngine.cpp   |  5 +++++
 4 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/include/TConsole.h b/include/TConsole.h
index ae34001..ab45eb0 100644
--- a/include/TConsole.h
+++ b/include/TConsole.h
@@ -18,7 +18,7 @@ public:
     Commandline& Internal() { return mCommandline; }
 
 private:
-    void ChangeToLuaConsole();
+    void ChangeToLuaConsole(const std::string& LuaStateId);
     void ChangeToRegularConsole();
 
     Commandline mCommandline;
@@ -27,5 +27,6 @@ private:
     TLuaEngine* mLuaEngine { nullptr };
     bool mIsLuaConsole { false };
     bool mFirstTime { true };
-    const std::string mStateId = "BEAMMP_SERVER_CONSOLE";
+    std::string mStateId;
+    const std::string mDefaultStateId = "BEAMMP_SERVER_CONSOLE";
 };
diff --git a/include/TLuaEngine.h b/include/TLuaEngine.h
index c38afc8..87e0c0d 100644
--- a/include/TLuaEngine.h
+++ b/include/TLuaEngine.h
@@ -84,6 +84,7 @@ public:
 
     static void WaitForAll(std::vector<std::shared_ptr<TLuaResult>>& Results);
     void ReportErrors(const std::vector<std::shared_ptr<TLuaResult> >& Results);
+    bool HasState(TLuaStateId StateId);
     [[nodiscard]] std::shared_ptr<TLuaResult> EnqueueScript(TLuaStateId StateID, const TLuaChunk& Script);
     [[nodiscard]] std::shared_ptr<TLuaResult> EnqueueFunctionCall(TLuaStateId StateID, const std::string& FunctionName, const std::vector<TLuaArgTypes>& Args);
     void EnsureStateExists(TLuaStateId StateId, const std::string& Name, bool DontCallOnInit = false);
diff --git a/src/TConsole.cpp b/src/TConsole.cpp
index da71a78..10bd03d 100644
--- a/src/TConsole.cpp
+++ b/src/TConsole.cpp
@@ -2,6 +2,7 @@
 #include "Common.h"
 #include "Compat.h"
 
+#include "CustomAssert.h"
 #include "LuaAPI.h"
 #include "TLuaEngine.h"
 
@@ -91,10 +92,15 @@ void TConsole::BackupOldLog() {
     }
 }
 
-void TConsole::ChangeToLuaConsole() {
+void TConsole::ChangeToLuaConsole(const std::string& LuaStateId) {
     if (!mIsLuaConsole) {
+        mStateId = LuaStateId;
         mIsLuaConsole = true;
-        Application::Console().WriteRaw("Entered Lua console. To exit, type `exit()`");
+        if (mStateId != mDefaultStateId) {
+            Application::Console().WriteRaw("Entered Lua console for state '" + mStateId + "'. To exit, type `exit()`");
+        } else {
+            Application::Console().WriteRaw("Entered Lua console. To exit, type `exit()`");
+        }
         mCachedRegularHistory = mCommandline.history();
         mCommandline.set_history(mCachedLuaHistory);
         mCommandline.set_prompt("lua> ");
@@ -104,10 +110,15 @@ void TConsole::ChangeToLuaConsole() {
 void TConsole::ChangeToRegularConsole() {
     if (mIsLuaConsole) {
         mIsLuaConsole = false;
-        Application::Console().WriteRaw("Left Lua console.");
+        if (mStateId != mDefaultStateId) {
+            Application::Console().WriteRaw("Left Lua console for state '" + mStateId + "'.");
+        } else {
+            Application::Console().WriteRaw("Left Lua console.");
+        }
         mCachedLuaHistory = mCommandline.history();
         mCommandline.set_history(mCachedRegularHistory);
         mCommandline.set_prompt("> ");
+        mStateId = mDefaultStateId;
     }
 }
 
@@ -144,8 +155,18 @@ TConsole::TConsole() {
                     if (cmd == "exit") {
                         beammp_info("gracefully shutting down");
                         Application::GracefullyShutdown();
-                    } else if (cmd == "lua") {
-                        ChangeToLuaConsole();
+                    } else if (cmd.size() >= 3 && cmd.substr(0, 3) == "lua") {
+                        if (cmd.size() > 3) {
+                            auto NewStateId = cmd.substr(4);
+                            beammp_assert(!NewStateId.empty());
+                            if (mLuaEngine->HasState(NewStateId)) {
+                                ChangeToLuaConsole(NewStateId);
+                            } else {
+                                Application::Console().WriteRaw("Lua state '" + NewStateId + "' is not a known state. Didn't switch to Lua.");
+                            }
+                        } else {
+                            ChangeToLuaConsole(mDefaultStateId);
+                        }
                     } else if (!cmd.empty()) {
                         auto FutureIsNonNil =
                             [](const std::shared_ptr<TLuaResult>& Future) {
@@ -208,5 +229,5 @@ void TConsole::WriteRaw(const std::string& str) {
 
 void TConsole::InitializeLuaConsole(TLuaEngine& Engine) {
     mLuaEngine = &Engine;
-    Engine.EnsureStateExists(mStateId, "Console");
+    Engine.EnsureStateExists(mDefaultStateId, "Console");
 }
diff --git a/src/TLuaEngine.cpp b/src/TLuaEngine.cpp
index 72e0b94..7c1a5d2 100644
--- a/src/TLuaEngine.cpp
+++ b/src/TLuaEngine.cpp
@@ -162,6 +162,11 @@ void TLuaEngine::ReportErrors(const std::vector<std::shared_ptr<TLuaResult>>& Re
     }
 }
 
+bool TLuaEngine::HasState(TLuaStateId StateId) {
+    std::unique_lock Lock(mLuaStatesMutex);
+    return mLuaStates.find(StateId) != mLuaStates.end();
+}
+
 std::shared_ptr<TLuaResult> TLuaEngine::EnqueueScript(TLuaStateId StateID, const TLuaChunk& Script) {
     std::unique_lock Lock(mLuaStatesMutex);
     return mLuaStates.at(StateID)->EnqueueScript(Script);