From 8531a67519e11d2e3cfc6f7b4e274ad083f20fcf Mon Sep 17 00:00:00 2001
From: WerWolv <werwolv98@gmail.com>
Date: Fri, 7 Jun 2024 21:27:01 +0200
Subject: [PATCH] fix: Don't unload background .NET scripts

---
 .../include/loaders/dotnet/dotnet_loader.hpp  |  1 +
 .../script_loader/include/loaders/loader.hpp  | 11 ++++---
 .../source/loaders/dotnet/dotnet_loader.cpp   | 32 +++++++++++++++----
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/plugins/script_loader/include/loaders/dotnet/dotnet_loader.hpp b/plugins/script_loader/include/loaders/dotnet/dotnet_loader.hpp
index f0cb3467f..6debde53c 100644
--- a/plugins/script_loader/include/loaders/dotnet/dotnet_loader.hpp
+++ b/plugins/script_loader/include/loaders/dotnet/dotnet_loader.hpp
@@ -15,6 +15,7 @@ namespace hex::script::loader {
 
         bool initialize() override;
         bool loadAll() override;
+        void clearScripts() override;
 
     private:
         std::function<int(const std::string &, bool, const std::fs::path&)>  m_runMethod;
diff --git a/plugins/script_loader/include/loaders/loader.hpp b/plugins/script_loader/include/loaders/loader.hpp
index 77688230f..0d047db3a 100644
--- a/plugins/script_loader/include/loaders/loader.hpp
+++ b/plugins/script_loader/include/loaders/loader.hpp
@@ -17,6 +17,8 @@ namespace hex::script::loader {
 
     struct Script {
         std::string name;
+        std::fs::path path;
+
         bool background;
         std::function<void()> entryPoint;
         const ScriptLoader *loader;
@@ -29,9 +31,10 @@ namespace hex::script::loader {
 
         virtual bool initialize() = 0;
         virtual bool loadAll() = 0;
+        virtual void clearScripts() = 0;
 
-        void addScript(std::string name, bool background, std::function<void()> entryPoint) {
-            m_scripts.emplace_back(std::move(name), background, std::move(entryPoint), this);
+        void addScript(std::string name, std::fs::path path, bool background, std::function<void()> entryPoint) {
+            m_scripts.emplace_back(std::move(name), std::move(path), background, std::move(entryPoint), this);
         }
 
         const auto& getScripts() const {
@@ -43,8 +46,8 @@ namespace hex::script::loader {
         }
 
     protected:
-        void clearScripts() {
-            m_scripts.clear();
+        auto& getScripts() {
+            return m_scripts;
         }
 
     private:
diff --git a/plugins/script_loader/source/loaders/dotnet/dotnet_loader.cpp b/plugins/script_loader/source/loaders/dotnet/dotnet_loader.cpp
index 6e282930c..f948ad61f 100644
--- a/plugins/script_loader/source/loaders/dotnet/dotnet_loader.cpp
+++ b/plugins/script_loader/source/loaders/dotnet/dotnet_loader.cpp
@@ -235,22 +235,36 @@ namespace hex::script::loader {
                 if (!std::fs::exists(scriptPath))
                     continue;
 
+                bool skip = false;
+                for (const auto &existingScript : getScripts()) {
+                    if (existingScript.path == scriptPath) {
+                        skip = true;
+                    }
+                }
+                if (skip)
+                    continue;
+
                 const bool hasMain = m_methodExists("Main", scriptPath);
                 const bool hasOnLoad = m_methodExists("OnLoad", scriptPath);
                 const auto scriptName = entry.path().stem().string();
 
+                if (hasMain && hasOnLoad) {
+                    log::error("Script '{}' has both a Main() and a OnLoad() function. Only one is allowed per script.", scriptName);
+                    continue;
+                } else if (!hasMain && !hasOnLoad) {
+                    log::error("Script '{}' has neither a Main() nor a OnLoad() function.", scriptName);
+                    continue;
+                }
+
                 if (hasMain) {
-                    this->addScript(scriptName, false, [this, scriptPath] {
+                    this->addScript(scriptName, scriptPath, false, [this, scriptPath] {
                         auto result = m_runMethod("Main", false, scriptPath);
                         if (result != 0) {
                             ui::ToastError::open(hex::format("Script '{}' running failed with code {}", result));
                         }
                     });
                 } else if (hasOnLoad) {
-                    this->addScript(scriptName, true, [] {});
-                }
-
-                if (hasOnLoad) {
+                    this->addScript(scriptName, scriptPath, true, [] {});
                     auto result = m_runMethod("OnLoad", true, scriptPath);
                     if (result != 0) {
                         TaskManager::doLater([=] {
@@ -258,11 +272,17 @@ namespace hex::script::loader {
                         });
                     }
                 }
-
             }
         }
 
         return true;
     }
 
+    void DotNetLoader::clearScripts() {
+        std::erase_if(getScripts(), [](const Script *script) {
+            return !script->background;
+        });
+    }
+
+
 }