From b6c97214ca97586728261e8bdcf9b28dd9de6d46 Mon Sep 17 00:00:00 2001
From: Josh Cooper <cooper.s.josh@outlook.com>
Date: Sat, 3 Sep 2022 11:18:46 -0700
Subject: [PATCH] Adds more features to spectate

---
 docs/plugins/spectate.rst |  36 +++++---
 plugins/spectate.cpp      | 175 ++++++++++++++++++++++++--------------
 2 files changed, 134 insertions(+), 77 deletions(-)

diff --git a/docs/plugins/spectate.rst b/docs/plugins/spectate.rst
index d6ea84581..e49cb81d1 100644
--- a/docs/plugins/spectate.rst
+++ b/docs/plugins/spectate.rst
@@ -2,39 +2,47 @@ spectate
 ========
 
 .. dfhack-tool::
-    :summary: Automatically follow exciting dwarves.
+    :summary: Automatically follow productive dwarves.
     :tags: fort interface
-    :no-command:
 
 Usage
 -----
 
 ::
 
-    spectate
-    spectate <option>
+    enable spectate
+    disable spectate
+    spectate <option> <value>
 
 
-The plugin will automatically switch which dwarf is being followed periodically,
-preferring dwarves on z-levels with the highest job activity.
+When enabled, the plugin will automatically switch which dwarf is being
+followed periodically, preferring dwarves on z-levels with the highest
+job activity.
+
+To set features that toggle between two states, use {0,1} to specify
+which state the feature should be in. Anything else will take any positive
+value.
 
 Examples
 --------
 
 ``spectate``
-    See the status of spectate, what is enabled or disabled.
+    The plugin reports its feature status.
 
-``spectate enable``
-    Enable spectate plugin to pseudo-randomly follow dwarves around.
 
-``spectate auto-unpause``
+``spectate auto-unpause 1``
     Enable the spectate plugin to automatically dismiss pause events caused
     by the game. Siege events are one example of such a game event.
 
+``spectate tick-interval 50``
+    Set the tick interval the followed dwarf can be changed at back to its
+    default value.
+
 Options
 -------
 
-:no option:     Show plugin status.
-:enable:        Enable plugin.
-:disable:       Disable plugin.
-:auto-unpause:  Toggle auto-dismissal of game pause events.
+:no option:      Show plugin status.
+:focus-jobs:     Toggle whether the plugin should always be following a job.
+:auto-unpause:   Toggle auto-dismissal of game pause events.
+:auto-disengage: Toggle auto-disengagement of plugin through player intervention.
+:tick-interval:  Set the plugin's tick interval for changing the followed dwarf.
diff --git a/plugins/spectate.cpp b/plugins/spectate.cpp
index e278ec6e3..cb8ce06b0 100644
--- a/plugins/spectate.cpp
+++ b/plugins/spectate.cpp
@@ -33,21 +33,21 @@ using namespace df::enums;
 
 DFHACK_PLUGIN("spectate");
 DFHACK_PLUGIN_IS_ENABLED(enabled);
-bool dismiss_pause_events = false;
+Pausing::AnnouncementLock* pause_lock = nullptr;
+bool lock_collision = false;
+bool unpause_enabled = false;
+bool disengage_enabled = false;
+bool focus_jobs_enabled = false;
 bool following_dwarf = false;
 df::unit* our_dorf = nullptr;
-void* job_watched = nullptr;
+df::job* job_watched = nullptr;
 int32_t timestamp = -1;
+uint64_t tick_span = 50;
 REQUIRE_GLOBAL(world);
 REQUIRE_GLOBAL(ui);
 REQUIRE_GLOBAL(pause_state);
 REQUIRE_GLOBAL(d_init);
-
-// todo: implement as user configurable variables
-#define tick_span 50
 #define base 0.99
-Pausing::AnnouncementLock* pause_lock = nullptr;
-bool lock_collision = false;
 
 command_result spectate (color_ostream &out, std::vector <std::string> & parameters);
 
@@ -81,33 +81,50 @@ DFhackCExport command_result plugin_shutdown (color_ostream &out) {
     return CR_OK;
 }
 
-DFhackCExport command_result plugin_onupdate(color_ostream &out) {
-    if (lock_collision) {
-        if (dismiss_pause_events) {
-            // player asked for auto-unpause enabled
-            World::SaveAnnouncementSettings();
-            if (World::DisableAnnouncementPausing()){
-                // now that we've got what we want, we can lock it down
-                lock_collision = false;
-                pause_lock->lock();
-            }
+void onTick(color_ostream& out, void* tick);
+void onJobStart(color_ostream &out, void* job);
+void onJobCompletion(color_ostream &out, void* job);
+
+void enable_auto_unpause(color_ostream &out, bool state){
+    if(unpause_enabled != state && lock_collision) {
+        // when enabled, lock collision means announcements haven't been disabled
+        // when disabled, lock collision means announcement are still disabled
+        // the only state left to consider here is what the lock should be set to
+        lock_collision = false;
+        unpause_enabled = state;
+        if (unpause_enabled) {
+            pause_lock->lock();
         } else {
-            if (World::RestoreAnnouncementSettings()) {
-                lock_collision = false;
-            }
+            // this one should be redundant, the lock should already be unlocked right now
+            pause_lock->unlock();
         }
+        out.print(unpause_enabled ? "auto-unpause: on\n" : "auto-unpause: off\n");
+        return;
     }
-    while (dismiss_pause_events && !world->status.popups.empty()) {
-        // dismiss announcement popup(s)
-        Gui::getCurViewscreen(true)->feed_key(interface_key::CLOSE_MEGA_ANNOUNCEMENT);
+    unpause_enabled = state;
+    // update the announcement settings if we can
+    if (unpause_enabled) {
+        if (World::SaveAnnouncementSettings()) {
+            World::DisableAnnouncementPausing();
+            pause_lock->lock();
+        } else {
+            lock_collision = true;
+        }
+    } else {
+        pause_lock->unlock();
+        if (!World::RestoreAnnouncementSettings()) {
+            // this in theory shouldn't happen, if others use the lock like we do in spectate
+            lock_collision = true;
+        }
+    }
+    // report to the user how things went
+    if (!lock_collision){
+        out.print(unpause_enabled ? "auto-unpause: on\n" : "auto-unpause: off\n");
+    } else {
+        out.print("auto-unpause: must wait for another Pausing::AnnouncementLock to be lifted. This setting will complete when the lock lifts.\n");
     }
-    return DFHack::CR_OK;
 }
 
-void onTick(color_ostream& out, void* tick);
-void onJobStart(color_ostream &out, void* job);
-void onJobCompletion(color_ostream &out, void* job);
-
 DFhackCExport command_result plugin_enable(color_ostream &out, bool enable) {
     namespace EM = EventManager;
     if (enable && !enabled) {
@@ -130,48 +147,78 @@ DFhackCExport command_result plugin_enable(color_ostream &out, bool enable) {
     return DFHack::CR_OK;
 }
 
+DFhackCExport command_result plugin_onupdate(color_ostream &out) {
+    if (lock_collision) {
+        if (unpause_enabled) {
+            // player asked for auto-unpause enabled
+            World::SaveAnnouncementSettings();
+            if (World::DisableAnnouncementPausing()){
+                // now that we've got what we want, we can lock it down
+                lock_collision = false;
+                pause_lock->lock();
+            }
+        } else {
+            if (World::RestoreAnnouncementSettings()) {
+                lock_collision = false;
+            }
+        }
+    }
+    while (unpause_enabled && !world->status.popups.empty()) {
+        // dismiss announcement popup(s)
+        Gui::getCurViewscreen(true)->feed_key(interface_key::CLOSE_MEGA_ANNOUNCEMENT);
+    }
+    if (disengage_enabled) {
+        if (our_dorf && our_dorf->id != df::global::ui->follow_unit){
+            plugin_enable(out, false);
+        }
+    }
+    return DFHack::CR_OK;
+}
+
 command_result spectate (color_ostream &out, std::vector <std::string> & parameters) {
     if(!parameters.empty()) {
-        if (parameters[0] == "auto-unpause") {
-            dismiss_pause_events = !dismiss_pause_events;
-
-            // update the announcement settings if we can
-            if (dismiss_pause_events) {
-                if (World::SaveAnnouncementSettings()) {
-                    World::DisableAnnouncementPausing();
-                    pause_lock->lock();
+        if (parameters.size() % 2 != 0) {
+            return DFHack::CR_WRONG_USAGE;
+        }
+        for (int i = 0; i+1 < parameters.size(); i += 2) {
+            if (parameters[i] == "auto-unpause") {
+                if (parameters[i+1] == "0") {
+                    enable_auto_unpause(out, false);
+                } else if (parameters[i+1] == "1") {
+                    enable_auto_unpause(out, true);
                 } else {
-                    lock_collision = true;
+                    return DFHack::CR_WRONG_USAGE;
                 }
-            } else {
-                pause_lock->unlock();
-                if (!World::RestoreAnnouncementSettings()) {
-                    // this in theory shouldn't happen, if others use the lock like we do in spectate
-                    lock_collision = true;
+            } else if (parameters[i] == "auto-disengage") {
+                if (parameters[i+1] == "0") {
+                    disengage_enabled = false;
+                } else if (parameters[i+1] == "1") {
+                    disengage_enabled = true;
+                } else {
+                    return DFHack::CR_WRONG_USAGE;
+                }
+            } else if (parameters[i] == "focus-jobs") {
+                if (parameters[i+1] == "0") {
+                    focus_jobs_enabled = false;
+                } else if (parameters[i+1] == "1") {
+                    focus_jobs_enabled = true;
+                } else {
+                    return DFHack::CR_WRONG_USAGE;
+                }
+            } else if (parameters[i] == "tick-interval") {
+                try {
+                    tick_span = std::stol(parameters[i + 1]);
+                } catch (const std::exception &e) {
+                    out.printerr("%s\n", e.what());
                 }
-            }
-
-            // report to the user how things went
-            if (!lock_collision){
-                out.print(dismiss_pause_events ? "auto-unpause: on\n" : "auto-unpause: off\n");
             } else {
-                out.print("auto-unpause: must wait for another Pausing::AnnouncementLock to be lifted. This setting will complete when the lock lifts.\n");
-            }
-
-            // probably a typo
-            if (parameters.size() == 2) {
-                out.print("If you want additional options open an issue on github, or mention it on discord.\n\n");
                 return DFHack::CR_WRONG_USAGE;
             }
-        } else if (parameters[0] == "disengage") {
-            //todo: cannibalize follow
-        } else {
-            return DFHack::CR_WRONG_USAGE;
         }
     } else {
         out.print(enabled ? "Spectate is enabled.\n" : "Spectate is disabled.\n");
         if(enabled) {
-            out.print(dismiss_pause_events ? "auto-unpause: on.\n" : "auto-unpause: off.\n");
+            out.print(unpause_enabled ? "auto-unpause: on.\n" : "auto-unpause: off.\n");
         }
     }
     return DFHack::CR_OK;
@@ -186,7 +233,7 @@ void onTick(color_ostream& out, void* ptr) {
             df::global::ui->follow_unit = -1;
         }
     }
-    if (!following_dwarf || (tick - timestamp) > tick_span || job_watched == nullptr) {
+    if (!following_dwarf || (focus_jobs_enabled && !job_watched) || (tick - timestamp) > tick_span) {
         std::vector<df::unit*> dwarves;
         for (auto unit: df::global::world->units.active) {
             if (!Units::isCitizen(unit)) {
@@ -196,6 +243,7 @@ void onTick(color_ostream& out, void* ptr) {
         }
         std::uniform_int_distribution<uint64_t> follow_any(0, dwarves.size() - 1);
         if (df::global::ui) {
+            // if you're looking at a warning about a local address escaping, it means the unit* from dwarves (which aren't local)
             our_dorf = dwarves[follow_any(RNG)];
             df::global::ui->follow_unit = our_dorf->id;
             job_watched = our_dorf->job.current_job;
@@ -216,14 +264,14 @@ void onJobStart(color_ostream& out, void* job_ptr) {
     int zcount = ++freq[job->pos.z];
     job_tracker.emplace(job->id);
     // if we're not doing anything~ then let's pick something
-    if (job_watched == nullptr || (tick - timestamp) > tick_span) {
+    if ((focus_jobs_enabled && !job_watched) || (tick - timestamp) > tick_span) {
         following_dwarf = true;
         // todo: allow the user to configure b, and also revise the math
         const double b = base;
         double p = b * ((double) zcount / job_tracker.size());
         std::bernoulli_distribution follow_job(p);
         if (!job->flags.bits.special && follow_job(RNG)) {
-            job_watched = job_ptr;
+            job_watched = job;
             df::unit* unit = Job::getWorker(job);
             if (df::global::ui && unit) {
                 our_dorf = unit;
@@ -252,9 +300,10 @@ void onJobCompletion(color_ostream &out, void* job_ptr) {
     // forget about it
     freq[job->pos.z]--;
     freq[job->pos.z] = freq[job->pos.z] < 0 ? 0 : freq[job->pos.z];
-    job_tracker.erase(job->id);
     // the job doesn't exist, so we definitely need to get rid of that
-    if (job_watched == job_ptr) {
+    job_tracker.erase(job->id);
+    // the event manager clones jobs and returns those clones for completed jobs. So the pointers won't match without a refactor of EM passing clones to both events
+    if (job_watched->id == job->id) {
         job_watched = nullptr;
     }
 }