#include "BehaviorTreeEditor.hpp"
#include "../lua/LuaBehaviorTreeApi.hpp"
#include <imgui.h>

static ImU32 nodeTypeColorU32(const Ogre::String &type)
{
	if (type == "sequence")
		return IM_COL32(0x88, 0xBB, 0xFF, 0xFF);
	if (type == "selector")
		return IM_COL32(0xFF, 0x88, 0x88, 0xFF);
	if (type == "invert")
		return IM_COL32(0xFF, 0xEE, 0x88, 0xFF);
	if (type == "task")
		return IM_COL32(0x88, 0xFF, 0x88, 0xFF);
	if (type == "check")
		return IM_COL32(0xFF, 0xAA, 0x66, 0xFF);
	if (type == "debugPrint")
		return IM_COL32(0xCC, 0xCC, 0xCC, 0xFF);
	if (type == "setAnimationState")
		return IM_COL32(0x88, 0xDD, 0xFF, 0xFF);
	if (type == "isAnimationEnded")
		return IM_COL32(0xDD, 0x88, 0xFF, 0xFF);
	if (type == "delay")
		return IM_COL32(0xFF, 0xD7, 0x00, 0xFF);
	if (type == "teleportToChild")
		return IM_COL32(0xFF, 0x69, 0xB4, 0xFF);
	if (type == "disablePhysics")
		return IM_COL32(0x80, 0x80, 0x80, 0xFF);
	if (type == "enablePhysics")
		return IM_COL32(0x4C, 0xCC, 0x4C, 0xFF);
	if (type == "sendEvent")
		return IM_COL32(0xFF, 0x99, 0x00, 0xFF);
	if (type == "luaTask")
		return IM_COL32(0xAA, 0x66, 0xFF, 0xFF);
	return IM_COL32(0xFF, 0xFF, 0xFF, 0xFF);
}

static void makeLabel(const BehaviorTreeNode &node, char *buf, size_t bufSize)
{
	if (node.isLeaf()) {
		if (node.type == "debugPrint")
			snprintf(buf, bufSize, "[print] \"%s\"",
				 node.name.empty() ? "" : node.name.c_str());
		else if (node.type == "setAnimationState")
			snprintf(buf, bufSize, "[anim] %s",
				 node.name.empty() ? "(none)" :
						     node.name.c_str());
		else if (node.type == "isAnimationEnded")
			snprintf(buf, bufSize, "[ended?] %s",
				 node.name.empty() ? "(none)" :
						     node.name.c_str());
		else
			snprintf(buf, bufSize, "[%s] %s", node.type.c_str(),
				 node.name.empty() ? "(none)" :
						     node.name.c_str());
	} else {
		snprintf(buf, bufSize, "[%s]", node.type.c_str());
	}
}

static size_t findChildIndex(const BehaviorTreeNode &parent,
			     const BehaviorTreeNode *child)
{
	for (size_t i = 0; i < parent.children.size(); i++)
		if (&parent.children[i] == child)
			return i;
	return (size_t)-1;
}

void BehaviorTreeEditor::queueAddChild(BehaviorTreeNode *parent,
				       const Ogre::String &type)
{
	TreeEditOp op;
	op.type = TreeEditOp::AddChild;
	op.targetNode = parent;
	op.childType = type;
	m_editOps.push_back(op);
}

void BehaviorTreeEditor::queueRemoveNode(BehaviorTreeNode *parent,
					 size_t childIndex)
{
	TreeEditOp op;
	op.type = TreeEditOp::RemoveNode;
	op.targetNode = parent;
	op.childIndex = childIndex;
	m_editOps.push_back(op);
}

void BehaviorTreeEditor::queueMoveNode(BehaviorTreeNode *parent,
				       size_t childIndex, int delta)
{
	TreeEditOp op;
	op.type = TreeEditOp::MoveNode;
	op.targetNode = parent;
	op.childIndex = childIndex;
	op.moveDelta = delta;
	m_editOps.push_back(op);
}

void BehaviorTreeEditor::applyEditOps(BehaviorTreeComponent &bt)
{
	bool changed = false;
	for (auto &op : m_editOps) {
		if (!op.targetNode)
			continue;
		auto &vec = op.targetNode->children;
		switch (op.type) {
		case TreeEditOp::AddChild: {
			BehaviorTreeNode child;
			child.type = op.childType;
			if (op.childType == "task" || op.childType == "check" ||
			    op.childType == "luaTask")
				child.name = "unnamed";
			vec.push_back(child);
			changed = true;
			break;
		}
		case TreeEditOp::AddLuaNode: {
			BehaviorTreeNode child;
			child.type = "luaTask";
			child.name = op.luaNodeName;
			vec.push_back(child);
			changed = true;
			break;
		}
		case TreeEditOp::RemoveNode: {
			if (op.childIndex < vec.size()) {
				vec.erase(vec.begin() + op.childIndex);
				changed = true;
			}
			break;
		}
		case TreeEditOp::MoveNode: {
			if (op.childIndex < vec.size()) {
				int newIdx = (int)op.childIndex + op.moveDelta;
				if (newIdx >= 0 && newIdx < (int)vec.size()) {
					std::swap(vec[op.childIndex],
						  vec[newIdx]);
					changed = true;
				}
			}
			break;
		}
		}
	}
	m_editOps.clear();
	if (changed)
		bt.markDirty();
}

// Helper: render a submenu of registered Lua node names.
// Returns true if a node was selected (via queueAddLuaNode).
static bool
renderLuaNodeSubmenu(BehaviorTreeNode *parent,
		     std::vector<BehaviorTreeEditor::TreeEditOp> &editOps)
{
	auto names = editScene::getRegisteredLuaNodeNames();
	if (names.empty()) {
		ImGui::MenuItem("(no Lua nodes registered)", nullptr, false,
				false);
		return false;
	}

	bool selected = false;
	for (const auto &name : names) {
		if (ImGui::MenuItem(name.c_str())) {
			BehaviorTreeEditor::TreeEditOp op;
			op.type = BehaviorTreeEditor::TreeEditOp::AddLuaNode;
			op.targetNode = parent;
			op.luaNodeName = name;
			editOps.push_back(op);
			selected = true;
		}
	}
	return selected;
}

void BehaviorTreeEditor::renderTree(BehaviorTreeNode &node,
				    BehaviorTreeNode *parent,
				    BehaviorTreeNode *&selectedNode,
				    BehaviorTreeNode *&selectedParent, int &id)
{
	int myId = id++;
	bool isLeaf = node.isLeaf();
	ImGuiTreeNodeFlags flags = ImGuiTreeNodeFlags_OpenOnArrow |
				   ImGuiTreeNodeFlags_OpenOnDoubleClick;
	if (isLeaf)
		flags |= ImGuiTreeNodeFlags_Leaf |
			 ImGuiTreeNodeFlags_NoTreePushOnOpen;
	if (selectedNode == &node)
		flags |= ImGuiTreeNodeFlags_Selected;

	char label[256];
	makeLabel(node, label, sizeof(label));

	ImGui::PushStyleColor(
		ImGuiCol_Text,
		ImGui::ColorConvertU32ToFloat4(nodeTypeColorU32(node.type)));
	bool open =
		ImGui::TreeNodeEx(("##btnode" + std::to_string(myId)).c_str(),
				  flags, "%s", label);
	ImGui::PopStyleColor();

	if (ImGui::IsItemClicked()) {
		selectedNode = &node;
		selectedParent = parent;
	}

	if (ImGui::BeginPopupContextItem()) {
		if (node.canHaveChildren()) {
			if (ImGui::MenuItem("Add Sequence"))
				queueAddChild(&node, "sequence");
			if (ImGui::MenuItem("Add Selector"))
				queueAddChild(&node, "selector");
			if (ImGui::MenuItem("Add Invert"))
				queueAddChild(&node, "invert");
		}
		if (node.canHaveChildren() || node.isLeaf()) {
			if (ImGui::MenuItem("Add Task"))
				queueAddChild(&node, "task");
			if (ImGui::MenuItem("Add Check"))
				queueAddChild(&node, "check");
			if (ImGui::MenuItem("Add Debug Print"))
				queueAddChild(&node, "debugPrint");
			if (ImGui::MenuItem("Add Set Animation State"))
				queueAddChild(&node, "setAnimationState");
			if (ImGui::MenuItem("Add Is Animation Ended"))
				queueAddChild(&node, "isAnimationEnded");
			if (ImGui::MenuItem("Add Set Bit"))
				queueAddChild(&node, "setBit");
			if (ImGui::MenuItem("Add Check Bit"))
				queueAddChild(&node, "checkBit");
			if (ImGui::MenuItem("Add Set Value"))
				queueAddChild(&node, "setValue");
			if (ImGui::MenuItem("Add Check Value"))
				queueAddChild(&node, "checkValue");
			if (ImGui::MenuItem("Add Blackboard Dump"))
				queueAddChild(&node, "blackboardDump");
			if (ImGui::MenuItem("Add Delay"))
				queueAddChild(&node, "delay");
			if (ImGui::MenuItem("Add Teleport To Child"))
				queueAddChild(&node, "teleportToChild");
			if (ImGui::MenuItem("Add Disable Physics"))
				queueAddChild(&node, "disablePhysics");
			if (ImGui::MenuItem("Add Enable Physics"))
				queueAddChild(&node, "enablePhysics");
			if (ImGui::MenuItem("Add Send Event"))
				queueAddChild(&node, "sendEvent");
			if (ImGui::BeginMenu("Add Lua Task")) {
				renderLuaNodeSubmenu(&node, m_editOps);
				ImGui::EndMenu();
			}
		}
		if (parent) {
			size_t idx = findChildIndex(*parent, &node);
			if (ImGui::MenuItem("Remove"))
				queueRemoveNode(parent, idx);
			if (ImGui::MenuItem("Move Up"))
				queueMoveNode(parent, idx, -1);
			if (ImGui::MenuItem("Move Down"))
				queueMoveNode(parent, idx, +1);
		}
		ImGui::EndPopup();
	}

	ImGui::PushID(myId);
	if (node.canHaveChildren()) {
		ImGui::SameLine();
		if (ImGui::SmallButton("+"))
			ImGui::OpenPopup("AddChild");
		if (ImGui::BeginPopup("AddChild")) {
			if (ImGui::MenuItem("Sequence"))
				queueAddChild(&node, "sequence");
			if (ImGui::MenuItem("Selector"))
				queueAddChild(&node, "selector");
			if (ImGui::MenuItem("Invert"))
				queueAddChild(&node, "invert");
			if (ImGui::MenuItem("Task"))
				queueAddChild(&node, "task");
			if (ImGui::MenuItem("Check"))
				queueAddChild(&node, "check");
			if (ImGui::MenuItem("Debug Print"))
				queueAddChild(&node, "debugPrint");
			if (ImGui::MenuItem("Set Animation State"))
				queueAddChild(&node, "setAnimationState");
			if (ImGui::MenuItem("Is Animation Ended"))
				queueAddChild(&node, "isAnimationEnded");
			if (ImGui::MenuItem("Set Bit"))
				queueAddChild(&node, "setBit");
			if (ImGui::MenuItem("Check Bit"))
				queueAddChild(&node, "checkBit");
			if (ImGui::MenuItem("Set Value"))
				queueAddChild(&node, "setValue");
			if (ImGui::MenuItem("Check Value"))
				queueAddChild(&node, "checkValue");
			if (ImGui::MenuItem("Blackboard Dump"))
				queueAddChild(&node, "blackboardDump");
			if (ImGui::MenuItem("Delay"))
				queueAddChild(&node, "delay");
			if (ImGui::MenuItem("Teleport To Child"))
				queueAddChild(&node, "teleportToChild");
			if (ImGui::MenuItem("Disable Physics"))
				queueAddChild(&node, "disablePhysics");
			if (ImGui::MenuItem("Enable Physics"))
				queueAddChild(&node, "enablePhysics");
			if (ImGui::MenuItem("Send Event"))
				queueAddChild(&node, "sendEvent");
			if (ImGui::BeginMenu("Lua Task")) {
				renderLuaNodeSubmenu(&node, m_editOps);
				ImGui::EndMenu();
			}
			ImGui::EndPopup();
		}
	}
	if (parent) {
		size_t myIndex = findChildIndex(*parent, &node);
		ImGui::SameLine();
		if (ImGui::SmallButton("-"))
			queueRemoveNode(parent, myIndex);
		if (myIndex > 0) {
			ImGui::SameLine();
			if (ImGui::SmallButton("^"))
				queueMoveNode(parent, myIndex, -1);
		}
		if (myIndex + 1 < parent->children.size()) {
			ImGui::SameLine();
			if (ImGui::SmallButton("v"))
				queueMoveNode(parent, myIndex, +1);
		}
	}
	ImGui::PopID();

	if (!isLeaf && open) {
		for (auto &child : node.children)
			renderTree(child, &node, selectedNode, selectedParent,
				   id);
		ImGui::TreePop();
	}
}

void BehaviorTreeEditor::renderProperties(BehaviorTreeNode *node)
{
	if (!node)
		return;

	ImGui::Separator();
	ImGui::Text("Node Properties");

	const char *types[] = { "sequence",
				"selector",
				"invert",
				"task",
				"check",
				"debugPrint",
				"setAnimationState",
				"isAnimationEnded",
				"setBit",
				"checkBit",
				"setValue",
				"checkValue",
				"blackboardDump",
				"delay",
				"teleportToChild",
				"disablePhysics",
				"enablePhysics",
				"sendEvent",
				"luaTask" };
	int typeIdx = 0;
	for (int i = 0; i < IM_ARRAYSIZE(types); i++) {
		if (node->type == types[i]) {
			typeIdx = i;
			break;
		}
	}
	if (ImGui::Combo("Type", &typeIdx, types, IM_ARRAYSIZE(types)))
		node->type = types[typeIdx];

	if (node->isLeaf()) {
		// For luaTask nodes, show a dropdown of registered Lua nodes
		if (node->type == "luaTask") {
			auto names = editScene::getRegisteredLuaNodeNames();
			if (names.empty()) {
				ImGui::TextDisabled("No Lua nodes registered");
			} else {
				// Find current index
				int current = -1;
				for (size_t i = 0; i < names.size(); i++) {
					if (names[i] == node->name) {
						current = (int)i;
						break;
					}
				}
				// Build a C-string array for ImGui Combo
				std::vector<const char *> cstrs;
				cstrs.reserve(names.size());
				for (const auto &n : names)
					cstrs.push_back(n.c_str());

				if (ImGui::Combo("Lua Node", &current,
						 cstrs.data(),
						 (int)cstrs.size())) {
					if (current >= 0 &&
					    current < (int)names.size())
						node->name = names[current];
				}
				if (ImGui::IsItemHovered())
					ImGui::SetTooltip(
						"Select a registered Lua node handler");
			}
		} else {
			char buf[256];
			snprintf(buf, sizeof(buf), "%s", node->name.c_str());

			const char *label = "Name";
			const char *hint = nullptr;
			if (node->type == "debugPrint") {
				label = "Message";
				hint = "Text to print when node runs";
			} else if (node->type == "setAnimationState") {
				label = "State (SM/State)";
				hint = "Format: stateMachineName/stateName";
			} else if (node->type == "isAnimationEnded") {
				label = "State Machine";
				hint = "Name of state machine to check";
			} else if (node->type == "setBit") {
				label = "Bit";
				hint = "Bit index or name";
			} else if (node->type == "checkBit") {
				label = "Bit";
				hint = "Bit index or name to check";
			} else if (node->type == "setValue") {
				label = "Key";
				hint = "Variable name";
			} else if (node->type == "checkValue") {
				label = "Key";
				hint = "Variable name to check";
			} else if (node->type == "blackboardDump") {
				label = "Label";
				hint = "Optional label prefix";
			} else if (node->type == "teleportToChild") {
				label = "Child Name";
				hint = "Name of child entity under Smart Object to teleport to";
			} else if (node->type == "sendEvent") {
				label = "Event Name";
				hint = "Name of the event to send via EventBus";
			}

			if (ImGui::InputText(label, buf, sizeof(buf)))
				node->name = buf;
			if (hint && ImGui::IsItemHovered())
				ImGui::SetTooltip("%s", hint);
		}

		// Params input for nodes that need it
		if (node->type == "setBit" || node->type == "setValue" ||
		    node->type == "checkValue" || node->type == "delay" ||
		    node->type == "sendEvent" || node->type == "luaTask") {
			char pbuf[256];
			snprintf(pbuf, sizeof(pbuf), "%s",
				 node->params.c_str());
			const char *phint = nullptr;
			if (node->type == "setBit")
				phint = "1 or 0 (default 1)";
			else if (node->type == "setValue")
				phint = "Value: 5, 3.14, or 1,2,3";
			else if (node->type == "checkValue")
				phint = "Comparison: == 5, > 3.14, != 1,2,3";
			else if (node->type == "delay")
				phint = "Duration in seconds (float, default 1.0)";
			else if (node->type == "sendEvent")
				phint = "Event params: key=val key2=val2 ...";
			else if (node->type == "luaTask")
				phint = "Parameters: key=val,key2=val2 passed to Lua function";
			if (ImGui::InputText("Params", pbuf, sizeof(pbuf)))
				node->params = pbuf;
			if (phint && ImGui::IsItemHovered())
				ImGui::SetTooltip("%s", phint);
		}
	}
}

bool BehaviorTreeEditor::renderComponent(flecs::entity entity,
					 BehaviorTreeComponent &bt)
{
	(void)entity;
	ImGui::PushID("BehaviorTree");
	bool modified = false;

	if (ImGui::CollapsingHeader("Behavior Tree",
				    ImGuiTreeNodeFlags_DefaultOpen)) {
		ImGui::Indent();

		char nameBuf[256];
		snprintf(nameBuf, sizeof(nameBuf), "%s", bt.treeName.c_str());
		if (ImGui::InputText("Tree Name", nameBuf, sizeof(nameBuf))) {
			bt.treeName = nameBuf;
			modified = true;
		}

		if (ImGui::Checkbox("Enabled", &bt.enabled))
			modified = true;

		ImGui::Separator();

		static BehaviorTreeNode *selectedNode = nullptr;
		static BehaviorTreeNode *selectedParent = nullptr;
		int id = 0;
		renderTree(bt.root, nullptr, selectedNode, selectedParent, id);

		bool hadEdits = !m_editOps.empty();
		applyEditOps(bt);
		if (hadEdits) {
			selectedNode = nullptr;
			selectedParent = nullptr;
			modified = true;
		}

		renderProperties(selectedNode);

		ImGui::Unindent();
	}

	ImGui::PopID();
	return modified;
}
