From 1bb25e99c6637d45da957c960b67455c4ec0c98b Mon Sep 17 00:00:00 2001
From: MihailRis <mexaruc@gmail.com>
Date: Thu, 13 Mar 2025 22:05:14 +0300
Subject: [PATCH] add core.capture_output

---
 src/logic/scripting/lua/libs/libcore.cpp  | 28 +++++++++++++++++++++++
 src/logic/scripting/lua/lua_overrides.cpp |  8 ++++---
 src/logic/scripting/scripting.cpp         |  2 ++
 src/logic/scripting/scripting.hpp         |  2 ++
 4 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/logic/scripting/lua/libs/libcore.cpp b/src/logic/scripting/lua/libs/libcore.cpp
index 7b8c152c..2a56644d 100644
--- a/src/logic/scripting/lua/libs/libcore.cpp
+++ b/src/logic/scripting/lua/libs/libcore.cpp
@@ -273,6 +273,33 @@ static int l_blank(lua::State*) {
     return 0;
 }
 
+static int l_capture_output(lua::State* L) {
+    int argc = lua::gettop(L) - 1;
+    if (!lua::isfunction(L, 1)) {
+        throw std::runtime_error("function expected as argument 1");
+    }
+    for (int i = 0; i < argc; i++) {
+        lua::pushvalue(L, i + 2);
+    }
+    lua::pushvalue(L, 1);
+
+    auto prev_output = output_stream;
+    auto prev_error = error_stream;
+
+    std::stringstream captured_output;
+
+    output_stream = &captured_output;
+    error_stream = &captured_output;
+    
+    lua::call_nothrow(L, argc, 0);
+
+    output_stream = prev_output;
+    error_stream = prev_error;
+    
+    lua::pushstring(L, captured_output.str());
+    return 1;
+}
+
 const luaL_Reg corelib[] = {
     {"blank", lua::wrap<l_blank>},
     {"get_version", lua::wrap<l_get_version>},
@@ -292,6 +319,7 @@ const luaL_Reg corelib[] = {
     {"get_setting_info", lua::wrap<l_get_setting_info>},
     {"open_folder", lua::wrap<l_open_folder>},
     {"quit", lua::wrap<l_quit>},
+    {"capture_output", lua::wrap<l_capture_output>},
     {"__load_texture", lua::wrap<l_load_texture>},
     {NULL, NULL}
 };
diff --git a/src/logic/scripting/lua/lua_overrides.cpp b/src/logic/scripting/lua/lua_overrides.cpp
index 12b0bbff..24746ad7 100644
--- a/src/logic/scripting/lua/lua_overrides.cpp
+++ b/src/logic/scripting/lua/lua_overrides.cpp
@@ -2,6 +2,8 @@
 
 #include "libs/api_lua.hpp"
 
+using namespace scripting;
+
 /// @brief Modified version of luaB_print from lbaselib.c
 int l_print(lua::State* L) {
     int n = lua::gettop(L); /* number of arguments */
@@ -16,10 +18,10 @@ int l_print(lua::State* L) {
                 L,
                 LUA_QL("tostring") " must return a string to " LUA_QL("print")
             );
-        if (i > 1) std::cout << "\t";
-        std::cout << s;
+        if (i > 1) *output_stream << "\t";
+        *output_stream << s;
         lua::pop(L); /* pop result */
     }
-    std::cout << std::endl;
+    *output_stream << std::endl;
     return 0;
 }
diff --git a/src/logic/scripting/scripting.cpp b/src/logic/scripting/scripting.cpp
index 66606776..4e7b37e0 100644
--- a/src/logic/scripting/scripting.cpp
+++ b/src/logic/scripting/scripting.cpp
@@ -34,6 +34,8 @@ static debug::Logger logger("scripting");
 
 static inline const std::string STDCOMP = "stdcomp";
 
+std::ostream* scripting::output_stream = &std::cout;
+std::ostream* scripting::error_stream = &std::cerr;
 Engine* scripting::engine = nullptr;
 Level* scripting::level = nullptr;
 const Content* scripting::content = nullptr;
diff --git a/src/logic/scripting/scripting.hpp b/src/logic/scripting/scripting.hpp
index bdd5d5d2..52de1ccb 100644
--- a/src/logic/scripting/scripting.hpp
+++ b/src/logic/scripting/scripting.hpp
@@ -42,6 +42,8 @@ namespace scripting {
     extern Level* level;
     extern BlocksController* blocks;
     extern LevelController* controller;
+    extern std::ostream* output_stream;
+    extern std::ostream* error_stream;
 
     void initialize(Engine* engine);