#include <algorithm>

#include <map>

#include <string>

#include <vector>

#include "DataDefs.h"

#include "Export.h"

#include "PluginManager.h"

#include "modules/Units.h"

#include "df/emotion_type.h"

#include "df/ui.h"

#include "df/unit.h"

#include "df/unit_personality.h"

#include "df/unit_soul.h"

#include "df/unit_thought_type.h"

#include "df/world.h"

using namespace std;

using namespace DFHack;

DFHACK_PLUGIN("misery");

DFHACK_PLUGIN_IS_ENABLED(is_enabled);

REQUIRE_GLOBAL(world);

REQUIRE_GLOBAL(ui);

REQUIRE_GLOBAL(cur_year);

REQUIRE_GLOBAL(cur_year_tick);

typedef df::unit_personality::T_emotions Emotion;

static int factor = 1;

static int tick = 0;

const int INTERVAL = 1000;

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

void add_misery(df::unit *unit);

void clear_misery(df::unit *unit);

const int FAKE_EMOTION_FLAG = (1 << 30);

const int STRENGTH_MULTIPLIER = 100;

bool is_valid_unit (df::unit *unit) {

if (!Units::isOwnRace(unit) || !Units::isOwnCiv(unit))

return false;

if (!Units::isActive(unit))

return false;

return true;

}

inline bool is_fake_emotion (Emotion *e) {

return e->flags.whole & FAKE_EMOTION_FLAG;

}

void add_misery (df::unit *unit) {

// Add a fake miserable thought

// Remove any fake ones that already exist

if (!unit || !unit->status.current_soul)

return;

clear_misery(unit);

auto &emotions = unit->status.current_soul->personality.emotions;

Emotion *e = new Emotion;

e->type = df::emotion_type::MISERY;

e->thought = df::unit_thought_type::SoapyBath;

e->flags.whole |= FAKE_EMOTION_FLAG;

emotions.push_back(e);

for (Emotion *e : emotions) {

if (is_fake_emotion(e)) {

e->year = *cur_year;

e->year_tick = *cur_year_tick;

e->strength = STRENGTH_MULTIPLIER * factor;

e->severity = STRENGTH_MULTIPLIER * factor;

}

}

}

void clear_misery (df::unit *unit) {

if (!unit || !unit->status.current_soul)

return;

auto &emotions = unit->status.current_soul->personality.emotions;

auto it = remove_if(emotions.begin(), emotions.end(), [](Emotion *e) {

if (is_fake_emotion(e)) {

delete e;

return true;

}

return false;

});

emotions.erase(it, emotions.end());

}

DFhackCExport command_result plugin_shutdown(color_ostream& out) {

factor = 0;

return CR_OK;

}

DFhackCExport command_result plugin_onupdate(color_ostream& out) {

static bool wasLoaded = false;

if ( factor == 0 || !world || !world->map.block_index ) {

if ( wasLoaded ) {

//we just unloaded the game: clear all data

factor = 0;

is_enabled = false;

wasLoaded = false;

}

return CR_OK;

}

if ( !wasLoaded ) {

wasLoaded = true;

}

if ( tick < INTERVAL ) {

tick++;

return CR_OK;

}

tick = 0;

//TODO: consider units.active

for (df::unit *unit : world->units.all) {

if (is_valid_unit(unit)) {

add_misery(unit);

}

}

return CR_OK;

}

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

commands.push_back(PluginCommand(

"misery",

"Increase the intensity of negative dwarven thoughts.",

misery));

return CR_OK;

}

DFhackCExport command_result plugin_enable(color_ostream &out, bool enable)

{

if (enable != is_enabled)

{

is_enabled = enable;

factor = enable ? 1 : 0;

tick = INTERVAL;

}

return CR_OK;

}

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

if ( !world || !world->map.block_index ) {

out.printerr("misery can only be enabled in fortress mode with a fully-loaded game.\n");

return CR_FAILURE;

}

if ( parameters.size() < 1 || parameters.size() > 2 ) {

return CR_WRONG_USAGE;

}

if ( parameters[0] == "disable" ) {

if ( parameters.size() > 1 ) {

return CR_WRONG_USAGE;

}

factor = 0;

is_enabled = false;

return CR_OK;

} else if ( parameters[0] == "enable" ) {

is_enabled = true;

factor = 1;

if ( parameters.size() == 2 ) {

int a = atoi(parameters[1].c_str());

if ( a < 1 ) {

out.printerr("Second argument must be a positive integer.\n");

return CR_WRONG_USAGE;

}

factor = a;

}

tick = INTERVAL;

} else if ( parameters[0] == "clear" ) {

for (df::unit *unit : world->units.all) {

if (is_valid_unit(unit)) {

clear_misery(unit);

}

}

} else {

int a = atoi(parameters[0].c_str());

if ( a < 0 ) {

return CR_WRONG_USAGE;

}

factor = a;

is_enabled = factor > 0;

}

return CR_OK;

}