#include "AnimationTreeEditor.hpp"
#include "../systems/AnimationTreeSystem.hpp"
#include "../components/AnimationTreeTemplate.hpp"
#include <imgui.h>

AnimationTreeEditor::AnimationTreeEditor(Ogre::SceneManager *sceneMgr)
	: m_sceneMgr(sceneMgr)
{
}

std::vector<Ogre::String> AnimationTreeEditor::getAvailableAnimations(
	flecs::entity entity)
{
	std::vector<Ogre::String> result;
	Ogre::Entity *ent =
		AnimationTreeSystem::findAnimatedEntity(entity);
	if (ent && ent->hasSkeleton()) {
		Ogre::AnimationStateSet *states =
			ent->getAllAnimationStates();
		if (states) {
			for (const auto &pair : states->getAnimationStates())
				result.push_back(pair.first);
		}
	}
	return result;
}

static ImU32 nodeTypeColorU32(const Ogre::String &type)
{
	if (type == "output")
		return IM_COL32(0xFF, 0x88, 0x88, 0xFF);
	if (type == "stateMachine")
		return IM_COL32(0x88, 0xBB, 0xFF, 0xFF);
	if (type == "state")
		return IM_COL32(0x88, 0xFF, 0x88, 0xFF);
	if (type == "speed")
		return IM_COL32(0xFF, 0xEE, 0x88, 0xFF);
	return IM_COL32(0xFF, 0xFF, 0xFF, 0xFF);
}

static void makeLabel(const AnimationTreeNode &node, char *buf,
		      size_t bufSize)
{
	if (node.type == "animation") {
		snprintf(buf, bufSize, "[%s] %s", node.type.c_str(),
			 node.animationName.empty() ?
				 "(none)" :
				 node.animationName.c_str());
	} else if (node.type == "speed") {
		snprintf(buf, bufSize, "[%s] %.2fx %s", node.type.c_str(),
			 node.speed,
			 node.name.empty() ? "" : node.name.c_str());
	} else if (node.type == "stateMachine") {
		snprintf(buf, bufSize, "[%s] %s (fade %.1f)",
			 node.type.c_str(), node.name.c_str(),
			 node.fadeSpeed);
	} else {
		snprintf(buf, bufSize, "[%s] %s", node.type.c_str(),
			 node.name.empty() ? "" : node.name.c_str());
	}
}

static int countChildrenOfType(const AnimationTreeNode &node,
			       const Ogre::String &type)
{
	int n = 0;
	for (const auto &c : node.children)
		if (c.type == type)
			n++;
	return n;
}

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

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

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

void AnimationTreeEditor::queueMoveNode(AnimationTreeNode *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 AnimationTreeEditor::applyEditOps(AnimationTreeComponent &at)
{
	bool changed = false;
	for (auto &op : m_editOps) {
		if (!op.targetNode)
			continue;
		auto &vec = op.targetNode->children;
		switch (op.type) {
		case TreeEditOp::AddChild: {
			AnimationTreeNode child;
			child.type = op.childType;
			if (op.childType == "stateMachine") {
				child.name = "sm" +
					     std::to_string(countChildrenOfType(
						     *op.targetNode,
						     "stateMachine"));
				child.fadeSpeed = 7.5f;
			} else if (op.childType == "state") {
				child.name = "state" +
					     std::to_string(countChildrenOfType(
						     *op.targetNode, "state"));
			} else if (op.childType == "animation") {
				child.animationName = "";
			} else if (op.childType == "speed") {
				child.speed = 1.0f;
			}
			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)
		at.dirty = true;
}

static bool canHaveChildren(const Ogre::String &type)
{
	return type == "output" || type == "stateMachine" ||
	       type == "state" || type == "speed";
}

void AnimationTreeEditor::renderTree(AnimationTreeNode &node,
				     AnimationTreeNode *parent,
				     AnimationTreeNode *&selectedNode,
				     AnimationTreeNode *&selectedParent,
				     int &id)
{
	int myId = id++;
	bool isLeaf = node.children.empty();
	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));

	/* Color the whole tree node label */
	ImGui::PushStyleColor(ImGuiCol_Text,
			      ImGui::ColorConvertU32ToFloat4(
				      nodeTypeColorU32(node.type)));
	bool open = ImGui::TreeNodeEx(
		("##animnode" + std::to_string(myId)).c_str(), flags,
		"%s", label);
	ImGui::PopStyleColor();

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

	/* Right-click context menu — must be attached to TreeNodeEx
	 * (the last item), before we add buttons on the same line */
	if (ImGui::BeginPopupContextItem()) {
		if (canHaveChildren(node.type)) {
			if (node.type == "output" ||
			    node.type == "stateMachine" ||
			    node.type == "state") {
				if (ImGui::MenuItem("Add State Machine"))
					queueAddChild(&node, "stateMachine");
			}
			if (node.type == "stateMachine" ||
			    node.type == "output" ||
			    node.type == "state") {
				if (ImGui::MenuItem("Add State"))
					queueAddChild(&node, "state");
			}
			if (node.type == "output" ||
			    node.type == "state" ||
			    node.type == "speed") {
				if (ImGui::MenuItem("Add Animation"))
					queueAddChild(&node, "animation");
			}
			if (node.type == "output" ||
			    node.type == "state") {
				if (ImGui::MenuItem("Add Speed"))
					queueAddChild(&node, "speed");
			}
		}
		if (parent) {
			size_t idx = findChildIndex(*parent, &node);
			if (ImGui::MenuItem("Remove")) {
				queueRemoveNode(parent, idx);
				if (selectedNode == &node) {
					selectedNode = nullptr;
					selectedParent = nullptr;
				}
			}
			if (ImGui::MenuItem("Move Up"))
				queueMoveNode(parent, idx, -1);
			if (ImGui::MenuItem("Move Down"))
				queueMoveNode(parent, idx, +1);
		}
		ImGui::EndPopup();
	}

	/* Inline buttons on the same line as the tree node label */
	ImGui::PushID(myId);
	if (canHaveChildren(node.type)) {
		ImGui::SameLine();
		if (ImGui::SmallButton("+")) {
			ImGui::OpenPopup("AddChild");
		}
		if (ImGui::BeginPopup("AddChild")) {
			if (node.type == "output" ||
			    node.type == "stateMachine" ||
			    node.type == "state") {
				if (ImGui::MenuItem("State Machine"))
					queueAddChild(&node, "stateMachine");
			}
			if (node.type == "stateMachine" ||
			    node.type == "output" ||
			    node.type == "state") {
				if (ImGui::MenuItem("State"))
					queueAddChild(&node, "state");
			}
			if (node.type == "output" ||
			    node.type == "state" ||
			    node.type == "speed") {
				if (ImGui::MenuItem("Animation"))
					queueAddChild(&node, "animation");
			}
			if (node.type == "output" ||
			    node.type == "state") {
				if (ImGui::MenuItem("Speed"))
					queueAddChild(&node, "speed");
			}
			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();

	/* Render children and TreePop (only when node is open) */
	if (!isLeaf && open) {
		for (auto &child : node.children)
			renderTree(child, &node, selectedNode,
				   selectedParent, id);
		ImGui::TreePop();
	}
}

void AnimationTreeEditor::renderProperties(
	AnimationTreeNode *node, AnimationTreeComponent &at,
	flecs::entity entity)
{
	if (!node)
		return;

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

	bool modified = false;

	/* Type selector */
	const char *types[] = {"output",   "stateMachine", "state",
			       "speed",    "animation"};
	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];
		modified = true;
	}

	/* Name */
	if (node->type == "stateMachine" || node->type == "state") {
		char buf[256];
		strncpy(buf, node->name.c_str(), sizeof(buf) - 1);
		buf[sizeof(buf) - 1] = '\0';
		if (ImGui::InputText("Name", buf, sizeof(buf))) {
			node->name = buf;
			modified = true;
		}
	}

	/* Animation name */
	if (node->type == "animation") {
		std::vector<Ogre::String> anims =
			getAvailableAnimations(entity);
		Ogre::String current = node->animationName;
		Ogre::String preview = current.empty() ? "(none)" :
						       current;
		if (ImGui::BeginCombo("Animation", preview.c_str())) {
			bool noneSel = current.empty();
			if (ImGui::Selectable("(none)", noneSel)) {
				node->animationName = "";
				modified = true;
			}
			for (const auto &name : anims) {
				bool sel = (current == name);
				if (ImGui::Selectable(name.c_str(), sel)) {
					node->animationName = name;
					modified = true;
				}
			}
			ImGui::EndCombo();
		}
	}

	/* Speed */
	if (node->type == "output" || node->type == "speed") {
		if (ImGui::SliderFloat("Speed", &node->speed, -2.0f,
				       10.0f, "%.2f"))
			modified = true;
	}

	/* Fade speed */
	if (node->type == "stateMachine") {
		if (ImGui::SliderFloat("Fade Speed", &node->fadeSpeed,
				       0.1f, 30.0f, "%.1f"))
			modified = true;
	}

	/* End transitions (for state machines) */
	if (node->type == "stateMachine" && !node->children.empty()) {
		ImGui::Separator();
		ImGui::Text("End Transitions");
		ImGui::TextDisabled(
			"Auto-switch state when animation ends");

		for (const auto &child : node->children) {
			if (child.type != "state")
				continue;
			char buf[256];
			auto it = node->endTransitions.find(child.name);
			if (it != node->endTransitions.end())
				strncpy(buf, it->second.c_str(),
					sizeof(buf) - 1);
			else
				buf[0] = '\0';
			buf[sizeof(buf) - 1] = '\0';

			Ogre::String lbl = "On '" + child.name +
					   "' end →";
			if (ImGui::InputText(lbl.c_str(), buf,
					     sizeof(buf))) {
				if (buf[0] == '\0')
					node->endTransitions.erase(
						child.name);
				else
					node->endTransitions[child.name] =
						buf;
				modified = true;
			}
		}
	}

	if (modified)
		at.dirty = true;
}

void AnimationTreeEditor::renderStatePreview(
	AnimationTreeComponent &at)
{
	std::vector<AnimationTreeNode *> stateMachines;
	at.root.collectStateMachines(stateMachines);
	if (stateMachines.empty())
		return;

	ImGui::Separator();
	ImGui::Text("State Preview");

	for (auto *sm : stateMachines) {
		if (!sm)
			continue;

		auto itCurrent = at.currentStates.find(sm->name);
		Ogre::String currentState =
			(itCurrent != at.currentStates.end()) ?
				itCurrent->second :
				"";

		Ogre::String label = sm->name + "##" + sm->name;
		Ogre::String preview = currentState.empty() ?
					       "(none)" :
					       currentState;

		if (ImGui::BeginCombo(label.c_str(), preview.c_str())) {
			for (const auto &child : sm->children) {
				if (child.type != "state")
					continue;
				bool sel = (currentState == child.name);
				if (ImGui::Selectable(child.name.c_str(),
						      sel)) {
					at.currentStates[sm->name] =
						child.name;
					at.dirty = true;
				}
			}
			ImGui::EndCombo();
		}
	}
}

bool AnimationTreeEditor::renderComponent(
	flecs::entity entity, AnimationTreeComponent &at)
{
	ImGui::PushID("AnimTree");
	bool modified = false;

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

		/* Template reference */
		char templBuf[256];
		snprintf(templBuf, sizeof(templBuf), "%s",
			 at.templateName.c_str());
		if (ImGui::InputText("Template Name", templBuf,
				     sizeof(templBuf))) {
			at.templateName = templBuf;
			modified = true;
		}
		if (ImGui::IsItemHovered()) {
			ImGui::SetTooltip(
				"Leave empty for local tree. Set to template name to copy from a template entity.");
		}

		bool isTemplateEntity = entity.has<AnimationTreeTemplate>();
		if (isTemplateEntity) {
			auto &templ = entity.get_mut<AnimationTreeTemplate>();
			char nameBuf[256];
			snprintf(nameBuf, sizeof(nameBuf), "%s",
				 templ.name.c_str());
			if (ImGui::InputText("Template Publish Name", nameBuf,
					     sizeof(nameBuf))) {
				templ.name = nameBuf;
				modified = true;
			}
			ImGui::TextDisabled("Template version: %llu",
					    (unsigned long long)templ.version);
		}

		ImGui::Separator();

		/* Global toggles */
		if (ImGui::Checkbox("Enabled", &at.enabled))
			modified = true;
		if (ImGui::Checkbox("Use Root Motion",
				    &at.useRootMotion)) {
			modified = true;
			at.dirty = true;
		}

		ImGui::Separator();

		/* Tree view */
		static AnimationTreeNode *selectedNode = nullptr;
		static AnimationTreeNode *selectedParent = nullptr;
		int id = 0;
		renderTree(at.root, nullptr, selectedNode,
			   selectedParent, id);

		/* Apply all deferred edits. Vector reallocations may
		 * invalidate pointers, so clear selection afterwards. */
		bool hadEdits = !m_editOps.empty();
		applyEditOps(at);
		if (hadEdits) {
			selectedNode = nullptr;
			selectedParent = nullptr;
		}

		/* Node properties */
		renderProperties(selectedNode, at, entity);

		/* State preview */
		renderStatePreview(at);

		ImGui::Unindent();
	}

	/* Auto-publish template changes */
	if (modified && entity.has<AnimationTreeTemplate>()) {
		auto &templ = entity.get_mut<AnimationTreeTemplate>();
		templ.version++;
	}

	ImGui::PopID();
	return modified;
}
