From 026ff26ab362dd282324b10a364de7734cabc69f Mon Sep 17 00:00:00 2001
From: Serhii Snitsaruk <limbonaut@pm.me>
Date: Mon, 5 Sep 2022 16:39:40 +0200
Subject: [PATCH] Use single function to get task icon in the editor

---
 bt/bt_task.cpp                    | 11 ++++----
 bt/bt_task.h                      |  2 +-
 editor/limbo_ai_editor_plugin.cpp | 46 ++++++++++++++++++++++++-------
 editor/limbo_ai_editor_plugin.h   |  3 ++
 4 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/bt/bt_task.cpp b/bt/bt_task.cpp
index 4c06a54..3897b13 100644
--- a/bt/bt_task.cpp
+++ b/bt/bt_task.cpp
@@ -215,9 +215,10 @@ String BTTask::get_configuration_warning() const {
 	return String();
 }
 
-Ref<Texture> BTTask::get_icon() const {
-	return EditorNode::get_singleton()->get_class_icon(get_class(), "Object");
-}
+// Ref<Texture> BTTask::get_icon() const {
+// 	// FIXME: Potential problem if tools are not built.
+// 	return EditorNode::get_singleton()->get_class_icon(get_class(), "Object");
+// }
 
 void BTTask::print_tree(int p_initial_tabs) const {
 	String tabs = "--";
@@ -260,8 +261,8 @@ void BTTask::_bind_methods() {
 	BIND_VMETHOD(MethodInfo(PropertyInfo(Variant::STRING, ""), "_generate_name"));
 	ClassDB::bind_method(D_METHOD("_get_configuration_warning"), &BTTask::get_configuration_warning);
 	BIND_VMETHOD(MethodInfo(PropertyInfo(Variant::STRING, ""), "_get_configuration_warning"));
-	ClassDB::bind_method(D_METHOD("_get_icon"), &BTTask::get_icon);
-	BIND_VMETHOD(MethodInfo(PropertyInfo(Variant::OBJECT, "", PROPERTY_HINT_RESOURCE_TYPE, "Texture"), "_get_icon"));
+	// ClassDB::bind_method(D_METHOD("_get_icon"), &BTTask::get_icon);
+	// BIND_VMETHOD(MethodInfo(PropertyInfo(Variant::OBJECT, "", PROPERTY_HINT_RESOURCE_TYPE, "Texture"), "_get_icon"));
 
 	// Public Methods.
 	ClassDB::bind_method(D_METHOD("is_root"), &BTTask::is_root);
diff --git a/bt/bt_task.h b/bt/bt_task.h
index 1bc3cd8..9a2af23 100644
--- a/bt/bt_task.h
+++ b/bt/bt_task.h
@@ -70,7 +70,7 @@ public:
 	int get_child_index(const Ref<BTTask> &p_child) const;
 	Ref<BTTask> next_sibling() const;
 	virtual String get_configuration_warning() const;
-	virtual Ref<Texture> get_icon() const;
+	// virtual Ref<Texture> get_icon() const;
 	void print_tree(int p_initial_tabs = 0) const;
 
 	BTTask();
diff --git a/editor/limbo_ai_editor_plugin.cpp b/editor/limbo_ai_editor_plugin.cpp
index 5f48eca..82e8e54 100644
--- a/editor/limbo_ai_editor_plugin.cpp
+++ b/editor/limbo_ai_editor_plugin.cpp
@@ -8,6 +8,7 @@
 #include "core/dictionary.h"
 #include "core/error_list.h"
 #include "core/error_macros.h"
+#include "core/io/image_loader.h"
 #include "core/io/resource_loader.h"
 #include "core/io/resource_saver.h"
 #include "core/list.h"
@@ -25,6 +26,7 @@
 #include "core/vector.h"
 #include "editor/editor_node.h"
 #include "editor/editor_plugin.h"
+#include "editor/editor_scale.h"
 #include "modules/limboai/bt/actions/bt_action.h"
 #include "modules/limboai/bt/behavior_tree.h"
 #include "modules/limboai/bt/bt_task.h"
@@ -64,7 +66,11 @@ void TaskTree::_update_item(TreeItem *p_item) {
 	Ref<BTTask> task = p_item->get_metadata(0);
 	ERR_FAIL_COND_MSG(!task.is_valid(), "Invalid task reference in metadata.");
 	p_item->set_text(0, task->get_task_name());
-	p_item->set_icon(0, task->get_icon());
+	if (task->get_script_instance() && !task->get_script_instance()->get_script()->get_path().empty()) {
+		p_item->set_icon(0, LimboAIEditor::get_task_icon(task->get_script_instance()->get_script()->get_path()));
+	} else {
+		p_item->set_icon(0, LimboAIEditor::get_task_icon(task->get_class()));
+	}
 	p_item->set_editable(0, false);
 
 	for (int i = 0; i < p_item->get_button_count(0); i++) {
@@ -375,15 +381,8 @@ void TaskPanel::_init() {
 			String meta = E->get();
 			String tname;
 			Ref<Texture> icon;
-			if (meta.begins_with("res:")) {
-				// Scripted class
-				tname = meta.get_file().get_basename();
-				icon = editor->get_object_icon(ResourceLoader::load(meta).ptr());
-			} else {
-				// Core class
-				tname = meta;
-				icon = get_icon(tname, "EditorIcons");
-			}
+			icon = LimboAIEditor::get_task_icon(meta);
+			tname = meta.begins_with("res:") ? meta.get_file().get_basename().trim_prefix("BT") : meta.trim_prefix("BT");
 			sec->add_task_button(tname, icon, meta);
 		}
 		sec->set_filter("");
@@ -719,6 +718,33 @@ void LimboAIEditor::apply_changes() {
 	}
 }
 
+Ref<Texture> LimboAIEditor::get_task_icon(String p_script_path_or_class) {
+	// TODO: Implement caching.
+	String base_type = p_script_path_or_class;
+	if (p_script_path_or_class.begins_with("res:")) {
+		Ref<Script> script = ResourceLoader::load(p_script_path_or_class, "Script");
+		Ref<Script> base_script = script;
+		while (base_script.is_valid()) {
+			StringName name = EditorNode::get_editor_data().script_class_get_name(base_script->get_path());
+			String icon_path = EditorNode::get_editor_data().script_class_get_icon_path(name);
+			if (!icon_path.empty()) {
+				Ref<Image> img = memnew(Image);
+				Error err = ImageLoader::load_image(icon_path, img);
+				if (err == OK) {
+					Ref<ImageTexture> icon = memnew(ImageTexture);
+					img->resize(16 * EDSCALE, 16 * EDSCALE, Image::INTERPOLATE_LANCZOS);
+					icon->create_from_image(img);
+					return icon;
+				}
+			}
+			base_script = base_script->get_base_script();
+		}
+		base_type = script->get_instance_base_type();
+	}
+
+	return EditorNode::get_singleton()->get_class_icon(base_type, "BTTask");
+}
+
 void LimboAIEditor::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("_add_task", "p_task"), &LimboAIEditor::_add_task);
 	ClassDB::bind_method(D_METHOD("_add_task_with_prototype", "p_prototype"), &LimboAIEditor::_add_task_with_prototype);
diff --git a/editor/limbo_ai_editor_plugin.h b/editor/limbo_ai_editor_plugin.h
index 7e6fa55..ce9d4d2 100644
--- a/editor/limbo_ai_editor_plugin.h
+++ b/editor/limbo_ai_editor_plugin.h
@@ -15,6 +15,7 @@
 #include "scene/gui/line_edit.h"
 #include "scene/gui/popup_menu.h"
 #include "scene/gui/tree.h"
+#include "scene/resources/texture.h"
 
 class TaskTree : public Control {
 	GDCLASS(TaskTree, Control);
@@ -146,6 +147,8 @@ protected:
 	static void _bind_methods();
 
 public:
+	static Ref<Texture> get_task_icon(String p_script_path_or_class);
+
 	void apply_changes();
 
 	LimboAIEditor(EditorNode *p_editor);