#include "Debug.h"

#include "LuaTools.h"

#include "PluginLua.h"

#include "PluginManager.h"

#include "VTableInterpose.h"

#include "modules/EventManager.h"

#include "modules/Gui.h"

#include "modules/Units.h"

#include "modules/World.h"

#include "df/announcements.h"

#include "df/d_init.h"

#include "df/plotinfost.h"

#include "df/unit.h"

#include "df/activity_entry.h"

#include "df/viewscreen_dwarfmodest.h"

#include "df/world.h"

#include <random>

using namespace DFHack;

using std::string;

using std::vector;

DFHACK_PLUGIN("spectate");

DFHACK_PLUGIN_IS_ENABLED(is_enabled);

REQUIRE_GLOBAL(d_init);

REQUIRE_GLOBAL(plotinfo);

REQUIRE_GLOBAL(world);

namespace DFHack {

DBG_DECLARE(spectate, control, DebugCategory::LINFO);

DBG_DECLARE(spectate, cycle, DebugCategory::LINFO);

DBG_DECLARE(spectate, event, DebugCategory::LINFO);

}

static std::default_random_engine rng;

static uint32_t next_cycle_unpaused_ms = 0; // threshold for the next cycle

static const size_t MAX_HISTORY = 200;

static const float CITIZEN_COMBAT_PREFERRED_WEIGHT = 25.0f;

static const float NICKNAMED_CITIZEN_PREFERRED_WEIGHT = 15.0f;

static const float OTHER_COMBAT_PREFERRED_WEIGHT = 10.0f;

static const float JOB_WEIGHT = 5.0f;

static const float OTHER_WEIGHT = 1.0f;

static const int32_t RECENT_UNITS_SCAN_CYCLE = 51;

static const float RECENT_UNIT_MULTIPLIER = 2.0f; // weight multiplier for recent units

static const int32_t RECENT_UNIT_MS = 15 * 60 * 1000; // 15 minutes

// jobs that get "other" weight instad of "job" weight

static const std::unordered_set<int32_t> boring_jobs = {

df::job_type::Eat,

df::job_type::Drink,

df::job_type::Sleep,

};

/////////////////////////////////////////////////////

// Configuration

static struct Configuration {

bool auto_unpause;

bool cinematic_action;

bool include_animals;

bool include_hostiles;

bool include_visitors;

bool include_wildlife;

bool prefer_conflict;

bool prefer_new_arrivals;

bool prefer_nicknamed;

int32_t follow_ms;

void reset() {

auto_unpause = false;

cinematic_action = true;

include_animals = false;

include_hostiles = false;

include_visitors = false;

include_wildlife = false;

prefer_conflict = true;

prefer_new_arrivals = true;

prefer_nicknamed = true;

follow_ms = 10000;

}

} config;

/////////////////////////////////////////////////////

// AnnouncementSettings

static class AnnouncementSettings {

bool was_in_settings = false; // whether we were in the vanilla settings screen last update

const size_t announcement_flag_arr_size = sizeof(decltype(df::announcements::flags)) / sizeof(df::announcement_flags);

std::unique_ptr<uint32_t *> saved;

void save_settings(color_ostream &out) {

if (!saved)

saved = std::make_unique<uint32_t *>(new uint32_t[announcement_flag_arr_size]);

DEBUG(control,out).print("saving announcement settings\n");

for (size_t i = 0; i < announcement_flag_arr_size; ++i)

(*saved)[i] = d_init->announcements.flags[i].whole;

}

public:

void reset(color_ostream &out, bool skip_restore) {

was_in_settings = false;

if (saved) {

if (!skip_restore)

restore_settings(out);

delete[] *saved;

saved.reset();

}

}

void on_update(color_ostream &out) {

if (Gui::matchFocusString("dwarfmode/Settings")) {

if (!was_in_settings) {

DEBUG(cycle,out).print("settings screen active; restoring announcement settings\n");

restore_settings(out);

was_in_settings = true;

}

} else if (was_in_settings) {

was_in_settings = false;

if (config.auto_unpause) {

DEBUG(cycle,out).print("settings screen now inactive; disabling announcement pausing\n");

save_and_scrub_settings(out);

}

}

}

void restore_settings(color_ostream &out) {

if (!saved || was_in_settings)

return;

DEBUG(control,out).print("restoring saved announcement settings\n");

for (size_t i = 0; i < announcement_flag_arr_size; ++i)

d_init->announcements.flags[i].whole = (*saved)[i];

}

// remove pausing, popups, and recentering from all announcements

// saves first so the original settings can be restored

void save_and_scrub_settings(color_ostream &out) {

if (Gui::matchFocusString("dwarfmode/Settings")) {

DEBUG(control,out).print("not modifying announcement settings; vanilla settings screen is active\n");

return;

}

save_settings(out);

DEBUG(control,out).print("scrubbing announcement settings\n");

for (auto& flag : d_init->announcements.flags) {

flag.bits.DO_MEGA = false;

flag.bits.PAUSE = false;

flag.bits.RECENTER = false;

}

}

} announcement_settings;

/////////////////////////////////////////////////////

// UnitHistory

static void follow_a_dwarf(color_ostream &out);

static class UnitHistory {

std::deque<int32_t> history;

size_t offset = 0;

public:

void reset() {

history.clear();

offset = 0;

}

void add_to_history(color_ostream &out, int32_t unit_id) {

if (offset > 0) {

DEBUG(cycle,out).print("trimming history forward of offset {}\n", offset);

history.resize(history.size() - offset);

offset = 0;

}

if (history.size() && history.back() == unit_id) {

DEBUG(cycle,out).print("unit {} is already current unit; not adding to history\n", unit_id);

} else {

history.push_back(unit_id);

if (history.size() > MAX_HISTORY) {

DEBUG(cycle,out).print("history full, truncating\n");

history.pop_front();

}

}

DEBUG(cycle,out).print("history now has {} entries\n", history.size());

}

void add_and_follow(color_ostream &out, df::unit *unit) {

// if we're currently following a unit, add it to the history if it's not already there

if (plotinfo->follow_unit > -1 && plotinfo->follow_unit != get_cur_unit_id()) {

DEBUG(cycle,out).print("currently following unit {} that is not in history; adding\n", plotinfo->follow_unit);

add_to_history(out, plotinfo->follow_unit);

}

int32_t id = unit->id;

add_to_history(out, id);

DEBUG(cycle,out).print("now following unit {}: {}\n", id, DF2CONSOLE(Units::getReadableName(unit)));

Gui::revealInDwarfmodeMap(Units::getPosition(unit), false, World::ReadPauseState());

plotinfo->follow_item = -1;

plotinfo->follow_unit = id;

}

void scan_back(color_ostream &out) {

if (history.empty() || offset >= history.size()-1) {

DEBUG(cycle,out).print("already at beginning of history\n");

return;

}

++offset;

int unit_id = get_cur_unit_id();

DEBUG(cycle,out).print("scanning back to unit {} at offset {}\n", unit_id, offset);

if (auto unit = df::unit::find(unit_id))

Gui::revealInDwarfmodeMap(Units::getPosition(unit), false, World::ReadPauseState());

plotinfo->follow_item = -1;

plotinfo->follow_unit = unit_id;

}

void scan_forward(color_ostream &out) {

if (history.empty() || offset == 0) {

DEBUG(cycle,out).print("already at most recent unit; following new unit\n");

follow_a_dwarf(out);

return;

}

--offset;

int unit_id = get_cur_unit_id();

DEBUG(cycle,out).print("scanning forward to unit {} at offset {}\n", unit_id, offset);

if (auto unit = df::unit::find(unit_id))

Gui::revealInDwarfmodeMap(Units::getPosition(unit), false, World::ReadPauseState());

plotinfo->follow_item = -1;

plotinfo->follow_unit = unit_id;

}

int32_t get_cur_unit_id() {

if (offset >= history.size())

return -1;

return history[history.size() - (1 + offset)];

}

} unit_history;

/////////////////////////////////////////////////////

// RecentUnits

static class RecentUnits {

std::unordered_map<int32_t, uint32_t> units; // unit id -> time seen

public:

void reset() {

units.clear();

}

void add(int32_t unit_id) {

units[unit_id] = Core::getInstance().getUnpausedMs();

}

bool contains(int32_t unit_id) {

return units.contains(unit_id);

}

void trim() {

uint32_t unpaused_ms = Core::getInstance().getUnpausedMs();

if (unpaused_ms < RECENT_UNIT_MS)

return;

uint32_t cutoff = unpaused_ms - RECENT_UNIT_MS;

for (auto it = units.begin(); it != units.end();) {

if (it->second < cutoff)

it = units.erase(it);

else

++it;

}

}

} recent_units;

static void on_new_active_unit(color_ostream& out, void* data) {

int32_t unit_id = reinterpret_cast<std::intptr_t>(data);

DEBUG(event,out).print("unit {} has arrived on map\n", unit_id);

recent_units.add(unit_id);

}

static EventManager::EventHandler new_unit_handler(plugin_self, on_new_active_unit, RECENT_UNITS_SCAN_CYCLE);

/////////////////////////////////////////////////////

// plugin API

static command_result do_command(color_ostream &out, vector<string> &parameters);

static void follow_a_dwarf(color_ostream &out);

DFhackCExport command_result plugin_init(color_ostream &out, std::vector <PluginCommand> &commands) {

DEBUG(control,out).print("initializing {}\n", plugin_name);

commands.push_back(PluginCommand(

plugin_name,

"Automated spectator mode.",

do_command));

return CR_OK;

}

static void on_disable(color_ostream &out, bool skip_restore_settings = false) {

EventManager::unregisterAll(plugin_self);

announcement_settings.reset(out, skip_restore_settings);

}

static bool is_squads_open() {

return Gui::matchFocusString("dwarfmode/Squads", Gui::getDFViewscreen());

}

static void set_next_cycle_unpaused_ms(color_ostream &out, bool has_active_combat = false) {

int32_t delay_ms = config.follow_ms;

if (config.cinematic_action && has_active_combat) {

std::normal_distribution<float> distribution(config.follow_ms / 2, config.follow_ms / 6);

delay_ms = distribution(rng);

delay_ms = std::min(config.follow_ms, std::max(1, delay_ms));

}

DEBUG(cycle,out).print("next cycle in {} ms\n", delay_ms);

next_cycle_unpaused_ms = Core::getInstance().getUnpausedMs() + delay_ms;

}

DFhackCExport command_result plugin_enable(color_ostream &out, bool enable) {

if (!Core::getInstance().isMapLoaded() || !World::isFortressMode()) {

out.printerr("Cannot enable {} without a loaded fort.\n", plugin_name);

return CR_FAILURE;

}

if (enable != is_enabled) {

is_enabled = enable;

DEBUG(control,out).print("{} from the API; persisting\n",

is_enabled ? "enabled" : "disabled");

if (enable) {

config.reset();

if (!Lua::CallLuaModuleFunction(out, "plugins.spectate", "refresh_cpp_config")) {

WARN(control,out).print("Failed to refresh config\n");

}

if (is_squads_open()) {

out.printerr("Cannot enable {} while the squads screen is open.\n", plugin_name);

Lua::CallLuaModuleFunction(out, "plugins.spectate", "show_squads_warning");

is_enabled = false;

return CR_FAILURE;

}

INFO(control,out).print("Spectate mode enabled!\n");

EventManager::registerListener(EventManager::EventType::UNIT_NEW_ACTIVE, new_unit_handler);

if (plotinfo->follow_unit > -1)

set_next_cycle_unpaused_ms(out);

else

follow_a_dwarf(out);

} else {

INFO(control,out).print("Spectate mode disabled!\n");

plotinfo->follow_unit = -1;

on_disable(out);

// don't reset the unit history since we may want to re-enable

}

} else {

DEBUG(control,out).print("{} from the API, but already {}; no action\n",

is_enabled ? "enabled" : "disabled",

is_enabled ? "enabled" : "disabled");

}

return CR_OK;

}

DFhackCExport command_result plugin_shutdown (color_ostream &out) {

DEBUG(control,out).print("shutting down {}\n", plugin_name);

on_disable(out);

return CR_OK;

}

DFhackCExport command_result plugin_onstatechange(color_ostream &out, state_change_event event) {

switch (event) {

case SC_WORLD_LOADED:

next_cycle_unpaused_ms = 0;

break;

case SC_WORLD_UNLOADED:

if (is_enabled) {

DEBUG(control,out).print("world unloaded; disabling {}\n",

plugin_name);

is_enabled = false;

on_disable(out, true);

unit_history.reset();

recent_units.reset();

}

break;

default:

break;

}

return CR_OK;

}

DFhackCExport command_result plugin_onupdate(color_ostream &out) {

announcement_settings.on_update(out);

if (plotinfo->follow_unit < 0 || plotinfo->follow_item > -1 || is_squads_open()) {

DEBUG(cycle,out).print("auto-disengage triggered\n");

is_enabled = false;

plotinfo->follow_unit = -1;

on_disable(out);

return CR_OK;

}

if (Core::getInstance().getUnpausedMs() >= next_cycle_unpaused_ms) {

recent_units.trim();

follow_a_dwarf(out);

}

return CR_OK;

}

static command_result do_command(color_ostream &out, vector<string> &parameters) {

bool show_help = false;

if (!Lua::CallLuaModuleFunction(out, "plugins.spectate", "parse_commandline", std::make_tuple(parameters),

1, [&](lua_State *L) {

show_help = !lua_toboolean(L, -1);

})) {

return CR_FAILURE;

}

return show_help ? CR_WRONG_USAGE : CR_OK;

}

/////////////////////////////////////////////////////

// cycle logic

static bool is_in_combat(df::unit *unit) {

if (Units::isCrazed(unit) || unit->mood == df::mood_type::Berserk)

return true;

for (auto activity_id : unit->activities) {

auto activity = df::activity_entry::find(activity_id);

if (activity && activity->type == df::activity_entry_type::Conflict)

return true;

}

return false;

}

static void get_dwarf_buckets(color_ostream &out,

vector<df::unit*> &citizen_combat_units,

vector<df::unit*> &other_combat_units,

vector<df::unit*> &nicknamed_units,

vector<df::unit*> &job_units,

vector<df::unit*> &other_units)

{

for (auto unit : world->units.active) {

if (Units::isDead(unit) || !Units::isActive(unit) || unit->flags1.bits.caged || unit->flags1.bits.chained || Units::isHidden(unit))

continue;

if (!config.include_animals && Units::isAnimal(unit))

continue;

if (!config.include_hostiles && Units::isDanger(unit))

continue;

if (!config.include_visitors && Units::isVisitor(unit))

continue;

if (!config.include_wildlife && Units::isWildlife(unit))

continue;

if (is_in_combat(unit)) {

TRACE(cycle, out).print("unit {} is in combat: {}\n", unit->id, DF2CONSOLE(Units::getReadableName(unit)));

if (Units::isCitizen(unit, true) || Units::isResident(unit, true))

citizen_combat_units.push_back(unit);

else

other_combat_units.push_back(unit);

} else if (config.prefer_nicknamed && !unit->name.nickname.empty()) {

nicknamed_units.push_back(unit);

} else if (unit->job.current_job && !boring_jobs.contains(unit->job.current_job->job_type)) {

job_units.push_back(unit);

} else {

other_units.push_back(unit);

}

}

}

static void add_bucket_to_vectors(const vector<df::unit*> &bucket, vector<df::unit*> &units, vector<float> &intervals, vector<float> &weights, float weight) {

if (bucket.empty())

return;

intervals.push_back(units.size() + bucket.size());

weights.push_back(weight);

units.insert(units.end(), bucket.begin(), bucket.end());

}

static void add_bucket(const vector<df::unit*> &bucket, vector<df::unit*> &units, vector<float> &intervals, vector<float> &weights, float weight) {

if (bucket.empty())

return;

if (config.prefer_new_arrivals) {

vector<df::unit*> new_bucket, old_bucket;

for (auto unit : bucket) {

if (recent_units.contains(unit->id))

new_bucket.push_back(unit);

else

old_bucket.push_back(unit);

}

add_bucket_to_vectors(new_bucket, units, intervals, weights, weight * RECENT_UNIT_MULTIPLIER);

add_bucket_to_vectors(old_bucket, units, intervals, weights, weight);

} else {

add_bucket_to_vectors(bucket, units, intervals, weights, weight);

}

}

#define DUMP_BUCKET(name) \

DEBUG(cycle,out).print("bucket: " #name ", size: {}\n", name.size()); \

if (debug_cycle.isEnabled(DebugCategory::LTRACE)) { \

for (auto u : name) { \

DEBUG(cycle,out).print(" unit {}: {}\n", u->id, DF2CONSOLE(Units::getReadableName(u))); \

} \

}

#define DUMP_FLOAT_VECTOR(name) \

DEBUG(cycle,out).print(#name ":\n"); \

for (float f : name) { \

DEBUG(cycle,out).print(" {}\n", (int)f); \

}

static void follow_a_dwarf(color_ostream &out) {

DEBUG(cycle,out).print("choosing a unit to follow\n");

vector<df::unit*> citizen_combat_units;

vector<df::unit*> other_combat_units;

vector<df::unit*> nicknamed_units;

vector<df::unit*> job_units;

vector<df::unit*> other_units;

get_dwarf_buckets(out, citizen_combat_units, other_combat_units, nicknamed_units, job_units, other_units);

set_next_cycle_unpaused_ms(out, !citizen_combat_units.empty());

// coalesce the buckets and add weights

vector<df::unit*> units;

vector<float> intervals;

vector<float> weights;

intervals.push_back(0);

add_bucket(citizen_combat_units, units, intervals, weights, config.prefer_conflict ? CITIZEN_COMBAT_PREFERRED_WEIGHT : JOB_WEIGHT);

add_bucket(other_combat_units, units, intervals, weights, config.prefer_conflict ? OTHER_COMBAT_PREFERRED_WEIGHT : JOB_WEIGHT);

add_bucket(nicknamed_units, units, intervals, weights, NICKNAMED_CITIZEN_PREFERRED_WEIGHT);

add_bucket(job_units, units, intervals, weights, JOB_WEIGHT);

add_bucket(other_units, units, intervals, weights, OTHER_WEIGHT);

if (units.empty()) {

DEBUG(cycle,out).print("no units to follow\n");

return;

}

std::piecewise_constant_distribution<float> distribution(intervals.begin(), intervals.end(), weights.begin());

int unit_idx = distribution(rng);

df::unit *unit = units[unit_idx];

if (debug_cycle.isEnabled(DebugCategory::LDEBUG)) {

DUMP_BUCKET(citizen_combat_units);

DUMP_BUCKET(other_combat_units);

DUMP_BUCKET(nicknamed_units);

DUMP_BUCKET(job_units);

DUMP_BUCKET(other_units);

DUMP_FLOAT_VECTOR(intervals);

DUMP_FLOAT_VECTOR(weights);

DEBUG(cycle,out).print("selected unit idx {}\n", unit_idx);

}

unit_history.add_and_follow(out, unit);

}

/////////////////////////////////////////////////////

// Lua API

static void spectate_setSetting(color_ostream &out, string name, int val) {

DEBUG(control,out).print("entering spectate_setSetting {} = {}\n", name, val);

if (name == "auto-unpause") {

if (val && !config.auto_unpause) {

announcement_settings.save_and_scrub_settings(out);

} else if (!val && config.auto_unpause) {

announcement_settings.restore_settings(out);

}

config.auto_unpause = val;

} else if (name == "cinematic-action") {

config.cinematic_action = val;

} else if (name == "include-animals") {

config.include_animals = val;

} else if (name == "include-hostiles") {

config.include_hostiles = val;

} else if (name == "include-visitors") {

config.include_visitors = val;

} else if (name == "include-wildlife") {

config.include_wildlife = val;

} else if (name == "prefer-conflict") {

config.prefer_conflict = val;

} else if (name == "prefer-new-arrivals") {

config.prefer_new_arrivals = val;

} else if (name == "prefer-nicknamed") {

config.prefer_nicknamed = val;

} else if (name == "follow-seconds") {

if (val <= 0) {

WARN(control,out).print("follow-seconds must be a positive integer\n");

return;

}

config.follow_ms = val * 1000;

} else {

WARN(control,out).print("Unknown setting: {}\n", name);

}

}

static void spectate_followPrev(color_ostream &out) {

DEBUG(control,out).print("entering spectate_followPrev\n");

unit_history.scan_back(out);

set_next_cycle_unpaused_ms(out);

};

static void spectate_followNext(color_ostream &out) {

DEBUG(control,out).print("entering spectate_followNext\n");

unit_history.scan_forward(out);

set_next_cycle_unpaused_ms(out);

};

static void spectate_addToHistory(color_ostream &out, int32_t unit_id) {

DEBUG(control,out).print("entering spectate_addToHistory; unit_id={}\n", unit_id);

if (!df::unit::find(unit_id)) {

WARN(control,out).print("unit with id {} not found; not adding to history\n", unit_id);

return;

}

unit_history.add_to_history(out, unit_id);

}

DFHACK_PLUGIN_LUA_FUNCTIONS {

DFHACK_LUA_FUNCTION(spectate_setSetting),

DFHACK_LUA_FUNCTION(spectate_followPrev),

DFHACK_LUA_FUNCTION(spectate_followNext),

DFHACK_LUA_FUNCTION(spectate_addToHistory),

DFHACK_LUA_END

};