#include <modules/Screen.h>

#include <modules/Translation.h>

#include <modules/Units.h>

#include <MiscUtils.h>

#include <VTableInterpose.h>

#include "uicommon.h"

#include "df/creature_raw.h"

#include "df/ui_look_list.h"

#include "df/viewscreen_announcelistst.h"

#include "df/viewscreen_petst.h"

#include "df/viewscreen_storesst.h"

#include "df/viewscreen_layer_stockpilest.h"

#include "df/viewscreen_layer_militaryst.h"

#include "df/viewscreen_layer_noblelistst.h"

#include "df/viewscreen_layer_workshop_profilest.h"

#include "df/viewscreen_topicmeeting_fill_land_holder_positionsst.h"

#include "df/viewscreen_tradegoodsst.h"

#include "df/viewscreen_unitlistst.h"

#include "df/viewscreen_buildinglistst.h"

#include "df/viewscreen_joblistst.h"

#include "df/historical_figure.h"

#include "df/viewscreen_locationsst.h"

#include "df/interface_key.h"

#include "df/interfacest.h"

#include "df/layer_object_listst.h"

#include "df/job.h"

#include "df/report.h"

#include "modules/Job.h"

#include "df/global_objects.h"

#include "df/viewscreen_dwarfmodest.h"

#include "modules/Gui.h"

#include "df/unit.h"

#include "df/misc_trait_type.h"

#include "df/unit_misc_trait.h"

using namespace std;

using std::set;

using std::vector;

using std::string;

using namespace DFHack;

using namespace df::enums;

DFHACK_PLUGIN("search");

DFHACK_PLUGIN_IS_ENABLED(is_enabled);

REQUIRE_GLOBAL(gps);

REQUIRE_GLOBAL(gview);

REQUIRE_GLOBAL(ui);

REQUIRE_GLOBAL(ui_building_assign_units);

REQUIRE_GLOBAL(ui_building_in_assign);

REQUIRE_GLOBAL(ui_building_item_cursor);

REQUIRE_GLOBAL(ui_look_cursor);

REQUIRE_GLOBAL(ui_look_list);

/*

Search Plugin

A plugin that adds a "Search" hotkey to some screens (Units, Trade and Stocks)

that allows filtering of the list items by a typed query.

Works by manipulating the vector(s) that the list based viewscreens use to store

their items. When a search is started the plugin saves the original vectors and

with each keystroke creates a new filtered vector off the saves for the screen

to use.

*/

void make_text_dim(int x1, int x2, int y)

{

for (int x = x1; x <= x2; x++)

{

Screen::Pen pen = Screen::readTile(x,y);

if (pen.valid())

{

if (pen.fg != 0)

{

if (pen.fg == 7)

pen.adjust(0,true);

else

pen.bold = 0;

}

Screen::paintTile(pen,x,y);

}

}

}

static bool is_live_screen(const df::viewscreen *screen)

{

for (df::viewscreen *cur = &gview->view; cur; cur = cur->child)

if (cur == screen)

return true;

return false;

}

static string get_unit_description(df::unit *unit)

{

if (!unit)

return "";

string desc;

auto name = Units::getVisibleName(unit);

if (name->has_name)

desc = Translation::TranslateName(name, false);

desc += ", " + Units::getProfessionName(unit); // Check animal type too

return desc;

}

static bool cursor_key_pressed (std::set<df::interface_key> *input)

{

// give text input (e.g. "2") priority over cursor keys

for (auto it = input->begin(); it != input->end(); ++it)

{

if (Screen::keyToChar(*it) != -1)

return false;

}

return

input->count(df::interface_key::CURSOR_UP) ||

input->count(df::interface_key::CURSOR_DOWN) ||

input->count(df::interface_key::CURSOR_LEFT) ||

input->count(df::interface_key::CURSOR_RIGHT) ||

input->count(df::interface_key::CURSOR_UPLEFT) ||

input->count(df::interface_key::CURSOR_UPRIGHT) ||

input->count(df::interface_key::CURSOR_DOWNLEFT) ||

input->count(df::interface_key::CURSOR_DOWNRIGHT) ||

input->count(df::interface_key::CURSOR_UP_FAST) ||

input->count(df::interface_key::CURSOR_DOWN_FAST) ||

input->count(df::interface_key::CURSOR_LEFT_FAST) ||

input->count(df::interface_key::CURSOR_RIGHT_FAST) ||

input->count(df::interface_key::CURSOR_UPLEFT_FAST) ||

input->count(df::interface_key::CURSOR_UPRIGHT_FAST) ||

input->count(df::interface_key::CURSOR_DOWNLEFT_FAST) ||

input->count(df::interface_key::CURSOR_DOWNRIGHT_FAST) ||

input->count(df::interface_key::CURSOR_UP_Z) ||

input->count(df::interface_key::CURSOR_DOWN_Z) ||

input->count(df::interface_key::CURSOR_UP_Z_AUX) ||

input->count(df::interface_key::CURSOR_DOWN_Z_AUX);

}

//

// START: Generic Search functionality

//

template <class S, class T>

class search_generic

{

public:

bool init(S *screen)

{

if (screen != viewscreen && !reset_on_change())

return false;

if (!can_init(screen))

{

if (is_valid())

{

clear_search();

reset_all();

}

return false;

}

if (!is_valid())

{

this->viewscreen = screen;

this->cursor_pos = get_viewscreen_cursor();

this->primary_list = get_primary_list();

this->select_key = get_search_select_key();

select_token = Screen::charToKey(select_key);

shift_select_token = Screen::charToKey(select_key + 'A' - 'a');

valid = true;

do_post_init();

}

return true;

}

// Called each time you enter or leave a searchable screen. Resets everything.

virtual void reset_all()

{

reset_search();

valid = false;

primary_list = NULL;

viewscreen = NULL;

select_key = 's';

}

bool reset_on_change()

{

if (valid && is_live_screen(viewscreen))

return false;

reset_all();

return true;

}

bool is_valid()

{

return valid;

}

// A new keystroke is received in a searchable screen

virtual bool process_input(set<df::interface_key> *input)

{

// If the page has two search options (Trade screen), only allow one to operate

// at a time

if (lock != NULL && lock != this)

return false;

// Allows custom preprocessing for each screen

if (!should_check_input(input))

return false;

bool key_processed = true;

if (entry_mode)

{

// Query typing mode

df::interface_key last_token = get_string_key(input);

int charcode = Screen::keyToChar(last_token);

if (charcode >= 32 && charcode <= 126)

{

// Standard character

search_string += char(charcode);

do_search();

}

else if (last_token == interface_key::STRING_A000)

{

// Backspace

if (search_string.length() > 0)

{

search_string.erase(search_string.length()-1);

do_search();

}

}

else if (input->count(interface_key::SELECT) || input->count(interface_key::LEAVESCREEN))

{

// ENTER or ESC: leave typing mode

end_entry_mode();

}

else if (cursor_key_pressed(input))

{

// Arrow key pressed. Leave entry mode and allow screen to process key

end_entry_mode();

key_processed = false;

}

}

// Not in query typing mode

else if (input->count(select_token))

{

// Hotkey pressed, enter typing mode

start_entry_mode();

}

else if (input->count(shift_select_token))

{

// Shift + Hotkey pressed, clear query

clear_search();

}

else

{

// Not a key for us, pass it on to the screen

key_processed = false;

}

return key_processed || entry_mode; // Only pass unrecognized keys down if not in typing mode

}

// Called after a keystroke has been processed

virtual void do_post_input_feed()

{

}

static search_generic<S, T> *lock;

bool in_entry_mode()

{

return entry_mode;

}

protected:

virtual string get_element_description(T element) const = 0;

virtual void render() const = 0;

virtual int32_t *get_viewscreen_cursor() = 0;

virtual vector<T> *get_primary_list() = 0;

search_generic()

{

reset_all();

}

virtual bool can_init(S *screen)

{

return true;

}

virtual void do_post_init()

{

}

void start_entry_mode()

{

entry_mode = true;

lock = this;

}

void end_entry_mode()

{

entry_mode = false;

lock = NULL;

}

virtual char get_search_select_key()

{

return 's';

}

virtual void reset_search()

{

end_entry_mode();

search_string = "";

saved_list1.clear();

}

// Shortcut to clear the search immediately

virtual void clear_search()

{

if (saved_list1.size() > 0)

{

*primary_list = saved_list1;

saved_list1.clear();

}

search_string = "";

}

virtual void save_original_values()

{

saved_list1 = *primary_list;

}

virtual void do_pre_incremental_search()

{

}

virtual void clear_viewscreen_vectors()

{

primary_list->clear();

}

virtual void add_to_filtered_list(size_t i)

{

primary_list->push_back(saved_list1[i]);

}

virtual void do_post_search()

{

}

virtual bool is_valid_for_search(size_t index)

{

return true;

}

virtual bool force_in_search(size_t index)

{

return false;

}

// The actual sort

virtual void do_search()

{

if (search_string.length() == 0)

{

clear_search();

return;

}

if (saved_list1.size() == 0)

// On first run, save the original list

save_original_values();

else

do_pre_incremental_search();

clear_viewscreen_vectors();

string search_string_l = toLower(search_string);

for (size_t i = 0; i < saved_list1.size(); i++ )

{

if (force_in_search(i))

{

add_to_filtered_list(i);

continue;

}

if (!is_valid_for_search(i))

continue;

T element = saved_list1[i];

string desc = toLower(get_element_description(element));

if (desc.find(search_string_l) != string::npos)

{

add_to_filtered_list(i);

}

}

do_post_search();

if (cursor_pos)

*cursor_pos = 0;

}

virtual bool should_check_input(set<df::interface_key> *input)

{

return true;

}

// Display hotkey message

void print_search_option(int x, int y = -1) const

{

auto dim = Screen::getWindowSize();

if (y == -1)

y = dim.y - 2;

OutputString((entry_mode) ? 4 : 12, x, y, string(1, select_key));

OutputString((entry_mode) ? 10 : 15, x, y, ": Search");

if (search_string.length() > 0 || entry_mode)

OutputString(15, x, y, ": " + search_string);

if (entry_mode)

OutputString(10, x, y, "_");

}

S *viewscreen;

vector <T> saved_list1, reference_list, *primary_list;

//bool redo_search;

string search_string;

protected:

int *cursor_pos;

char select_key;

bool valid;

bool entry_mode;

df::interface_key select_token;

df::interface_key shift_select_token;

};

template <class S, class T> search_generic<S, T> *search_generic<S, T> ::lock = NULL;

// Search class helper for layered screens

template <class S, class T, int LIST_ID>

class layered_search : public search_generic<S, T>

{

protected:

virtual bool can_init(S *screen)

{

auto list = getLayerList(screen);

if (!is_list_valid(screen) || !list || !list->active)

return false;

return true;

}

virtual bool is_list_valid(S*)

{

return true;

}

virtual void do_search()

{

search_generic<S,T>::do_search();

auto list = getLayerList(this->viewscreen);

list->num_entries = this->get_primary_list()->size();

}

int32_t *get_viewscreen_cursor()

{

auto list = getLayerList(this->viewscreen);

return &list->cursor;

}

virtual void clear_search()

{

search_generic<S,T>::clear_search();

if (is_list_valid(this->viewscreen))

{

auto list = getLayerList(this->viewscreen);

list->num_entries = this->get_primary_list()->size();

}

}

private:

static df::layer_object_listst *getLayerList(const df::viewscreen_layer *layer)

{

return virtual_cast<df::layer_object_listst>(vector_get(layer->layer_objects, LIST_ID));

}

};

// Parent class for screens that have more than one primary list to synchronise

template < class S, class T, class PARENT = search_generic<S,T> >

class search_multicolumn_modifiable_generic : public PARENT

{

protected:

vector <T> reference_list;

vector <int> saved_indexes;

bool read_only;

virtual void update_saved_secondary_list_item(size_t i, size_t j) = 0;

virtual void save_secondary_values() = 0;

virtual void clear_secondary_viewscreen_vectors() = 0;

virtual void add_to_filtered_secondary_lists(size_t i) = 0;

virtual void clear_secondary_saved_lists() = 0;

virtual void reset_secondary_viewscreen_vectors() = 0;

virtual void restore_secondary_values() = 0;

virtual void do_post_init()

{

// If true, secondary list isn't modifiable so don't bother synchronising values

read_only = false;

}

void reset_all()

{

PARENT::reset_all();

reference_list.clear();

saved_indexes.clear();

reset_secondary_viewscreen_vectors();

}

void reset_search()

{

PARENT::reset_search();

reference_list.clear();

saved_indexes.clear();

clear_secondary_saved_lists();

}

virtual void clear_search()

{

if (this->saved_list1.size() > 0)

{

do_pre_incremental_search();

restore_secondary_values();

}

clear_secondary_saved_lists();

PARENT::clear_search();

do_post_search();

}

virtual bool is_match(T &a, T &b) = 0;

virtual bool is_match(vector<T> &a, vector<T> &b) = 0;

void do_pre_incremental_search()

{

PARENT::do_pre_incremental_search();

if (read_only)

return;

bool list_has_been_sorted = (this->primary_list->size() == reference_list.size()

&& !is_match(*this->primary_list, reference_list));

for (size_t i = 0; i < saved_indexes.size(); i++)

{

int adjusted_item_index = i;

if (list_has_been_sorted)

{

for (size_t j = 0; j < this->primary_list->size(); j++)

{

if (is_match((*this->primary_list)[j], reference_list[i]))

{

adjusted_item_index = j;

break;

}

}

}

update_saved_secondary_list_item(saved_indexes[i], adjusted_item_index);

}

saved_indexes.clear();

}

void clear_viewscreen_vectors()

{

search_generic<S,T>::clear_viewscreen_vectors();

saved_indexes.clear();

clear_secondary_viewscreen_vectors();

}

void add_to_filtered_list(size_t i)

{

search_generic<S,T>::add_to_filtered_list(i);

add_to_filtered_secondary_lists(i);

if (!read_only)

saved_indexes.push_back(i); // Used to map filtered indexes back to original, if needed

}

virtual void do_post_search()

{

if (!read_only)

reference_list = *this->primary_list;

}

void save_original_values()

{

search_generic<S,T>::save_original_values();

save_secondary_values();

}

};

// This basic match function is separated out from the generic multi column class, because the

// pets screen, which uses a union in its primary list, will cause a compile failure if this

// match function exists in the generic class

template < class S, class T, class PARENT = search_generic<S,T> >

class search_multicolumn_modifiable : public search_multicolumn_modifiable_generic<S, T, PARENT>

{

bool is_match(T &a, T &b)

{

return a == b;

}

bool is_match(vector<T> &a, vector<T> &b)

{

return a == b;

}

};

// General class for screens that have only one secondary list to keep in sync

template < class S, class T, class V, class PARENT = search_generic<S,T> >

class search_twocolumn_modifiable : public search_multicolumn_modifiable<S, T, PARENT>

{

public:

protected:

virtual vector<V> * get_secondary_list() = 0;

virtual void do_post_init()

{

search_multicolumn_modifiable<S, T, PARENT>::do_post_init();

secondary_list = get_secondary_list();

}

void save_secondary_values()

{

saved_secondary_list = *secondary_list;

}

void reset_secondary_viewscreen_vectors()

{

secondary_list = NULL;

}

virtual void update_saved_secondary_list_item(size_t i, size_t j)

{

saved_secondary_list[i] = (*secondary_list)[j];

}

void clear_secondary_viewscreen_vectors()

{

secondary_list->clear();

}

void add_to_filtered_secondary_lists(size_t i)

{

secondary_list->push_back(saved_secondary_list[i]);

}

void clear_secondary_saved_lists()

{

saved_secondary_list.clear();

}

void restore_secondary_values()

{

*secondary_list = saved_secondary_list;

}

vector<V> *secondary_list, saved_secondary_list;

};

// Parent struct for the hooks, use optional param D to generate multiple search classes in the same

// viewscreen but different static modules

template <class T, class V, int D = 0>

struct generic_search_hook : T

{

typedef T interpose_base;

static V module;

DEFINE_VMETHOD_INTERPOSE(void, feed, (set<df::interface_key> *input))

{

if (!module.init(this))

{

INTERPOSE_NEXT(feed)(input);

return;

}

if (!module.process_input(input))

{

INTERPOSE_NEXT(feed)(input);

module.do_post_input_feed();

}

}

DEFINE_VMETHOD_INTERPOSE(void, render, ())

{

bool ok = module.init(this);

INTERPOSE_NEXT(render)();

if (ok)

module.render();

}

DEFINE_VMETHOD_INTERPOSE(bool, key_conflict, (df::interface_key key))

{

if (module.in_entry_mode() && (key == interface_key::MOVIES || key == interface_key::HELP))

return true;

return INTERPOSE_NEXT(key_conflict)(key);

}

};

template <class T, class V, int D> V generic_search_hook<T, V, D> ::module;

// Hook definition helpers

#define IMPLEMENT_HOOKS_WITH_ID(screen, module, id, prio) \

typedef generic_search_hook<screen, module, id> module##_hook; \

template<> IMPLEMENT_VMETHOD_INTERPOSE_PRIO(module##_hook, feed, prio); \

template<> IMPLEMENT_VMETHOD_INTERPOSE_PRIO(module##_hook, render, prio); \

template<> IMPLEMENT_VMETHOD_INTERPOSE_PRIO(module##_hook, key_conflict, prio)

#define IMPLEMENT_HOOKS(screen, module) \

typedef generic_search_hook<screen, module> module##_hook; \

template<> IMPLEMENT_VMETHOD_INTERPOSE(module##_hook, feed); \

template<> IMPLEMENT_VMETHOD_INTERPOSE(module##_hook, render); \

template<> IMPLEMENT_VMETHOD_INTERPOSE(module##_hook, key_conflict)

#define IMPLEMENT_HOOKS_PRIO(screen, module, prio) \

typedef generic_search_hook<screen, module> module##_hook; \

template<> IMPLEMENT_VMETHOD_INTERPOSE_PRIO(module##_hook, feed, prio); \

template<> IMPLEMENT_VMETHOD_INTERPOSE_PRIO(module##_hook, render, prio); \

template<> IMPLEMENT_VMETHOD_INTERPOSE_PRIO(module##_hook, key_conflict, prio);

//

// END: Generic Search functionality

//

//

// START: Animal screen search

//

typedef search_multicolumn_modifiable_generic<df::viewscreen_petst, df::viewscreen_petst::T_animal> pets_search_base;

class pets_search : public pets_search_base

{

typedef df::viewscreen_petst::T_animal T_animal;

typedef df::viewscreen_petst::T_mode T_mode;

public:

void render() const

{

print_search_option(25, 4);

}

private:

bool can_init(df::viewscreen_petst *screen)

{

return pets_search_base::can_init(screen) && screen->mode == T_mode::List;

}

int32_t *get_viewscreen_cursor()

{

return &viewscreen->cursor;

}

vector<T_animal> *get_primary_list()

{

return &viewscreen->animal;

}

virtual void do_post_init()

{

is_vermin = &viewscreen->is_vermin;

is_tame = &viewscreen->is_tame;

is_adopting = &viewscreen->is_adopting;

}

string get_element_description(df::viewscreen_petst::T_animal element) const

{

return get_unit_description(element.unit);

}

bool should_check_input()

{

return viewscreen->mode == T_mode::List;

}

bool is_valid_for_search(size_t i)

{

return is_vermin_s[i] == 0;

}

void save_secondary_values()

{

is_vermin_s = *is_vermin;

is_tame_s = *is_tame;

is_adopting_s = *is_adopting;

}

void reset_secondary_viewscreen_vectors()

{

is_vermin = NULL;

is_tame = NULL;

is_adopting = NULL;

}

void update_saved_secondary_list_item(size_t i, size_t j)

{

is_vermin_s[i] = (*is_vermin)[j];

is_tame_s[i] = (*is_tame)[j];

is_adopting_s[i] = (*is_adopting)[j];

}

void clear_secondary_viewscreen_vectors()

{

is_vermin->clear();

is_tame->clear();

is_adopting->clear();

}

void add_to_filtered_secondary_lists(size_t i)

{

is_vermin->push_back(is_vermin_s[i]);

is_tame->push_back(is_tame_s[i]);

is_adopting->push_back(is_adopting_s[i]);

}

void clear_secondary_saved_lists()

{

is_vermin_s.clear();

is_tame_s.clear();

is_adopting_s.clear();

}

void restore_secondary_values()

{

*is_vermin = is_vermin_s;

*is_tame = is_tame_s;

*is_adopting = is_adopting_s;

}

bool is_match(T_animal &a, T_animal &b)

{

return a.unit == b.unit;

}

bool is_match(vector<T_animal> &a, vector<T_animal> &b)

{

for (size_t i = 0; i < a.size(); i++)

{

if (!is_match(a[i], b[i]))

return false;

}

return true;

}

std::vector<char > *is_vermin, is_vermin_s;

std::vector<char > *is_tame, is_tame_s;

std::vector<char > *is_adopting, is_adopting_s;

};

IMPLEMENT_HOOKS_WITH_ID(df::viewscreen_petst, pets_search, 1, 0);

//

// END: Animal screen search

//

//

// START: Animal knowledge screen search

//

typedef search_generic<df::viewscreen_petst, int32_t> animal_knowledge_search_base;

class animal_knowledge_search : public animal_knowledge_search_base

{

typedef df::viewscreen_petst::T_mode T_mode;

bool can_init(df::viewscreen_petst *screen)

{

return animal_knowledge_search_base::can_init(screen) && screen->mode == T_mode::TrainingKnowledge;

}

public:

void render() const

{

print_search_option(2, 4);

}

private:

int32_t *get_viewscreen_cursor()

{

return NULL;

}

vector<int32_t> *get_primary_list()

{

return &viewscreen->known;

}

string get_element_description(int32_t id) const

{

auto craw = df::creature_raw::find(id);

string out;

if (craw)

{

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

out += craw->name[i] + " ";

}

return out;

}

};

IMPLEMENT_HOOKS_WITH_ID(df::viewscreen_petst, animal_knowledge_search, 2, 0);

//

// END: Animal knowledge screen search

//

//

// START: Animal trainer search

//

typedef search_twocolumn_modifiable<df::viewscreen_petst, df::unit*, df::viewscreen_petst::T_trainer_mode> animal_trainer_search_base;

class animal_trainer_search : public animal_trainer_search_base

{

typedef df::viewscreen_petst::T_mode T_mode;

typedef df::viewscreen_petst::T_trainer_mode T_trainer_mode;

bool can_init(df::viewscreen_petst *screen)

{

return animal_trainer_search_base::can_init(screen) && screen->mode == T_mode::SelectTrainer;

}

public:

void render() const

{

Screen::paintTile(Screen::Pen(186, 8, 0), 14, 2);

Screen::paintTile(Screen::Pen(186, 8, 0), gps->dimx - 14, 2);

Screen::paintTile(Screen::Pen(201, 8, 0), 14, 1);

Screen::paintTile(Screen::Pen(187, 8, 0), gps->dimx - 14, 1);

for (int x = 15; x <= gps->dimx - 15; ++x)

{

Screen::paintTile(Screen::Pen(205, 8, 0), x, 1);

Screen::paintTile(Screen::Pen(0, 0, 0), x, 2);

}

print_search_option(16, 2);

}

private:

int32_t *get_viewscreen_cursor()

{

return &viewscreen->trainer_cursor;

}

vector<df::unit*> *get_primary_list()

{

return &viewscreen->trainer_unit;

}

string get_element_description(df::unit *u) const

{

return get_unit_description(u);

}

std::vector<T_trainer_mode> *get_secondary_list()

{

return &viewscreen->trainer_mode;

}

public:

bool process_input(set<df::interface_key> *input)

{

if (input->count(interface_key::SELECT) && viewscreen->trainer_unit.empty() && !in_entry_mode())

return true;

return animal_trainer_search_base::process_input(input);

}