#include "AnimationTreeSystem.hpp"
#include "CharacterSlotSystem.hpp"
#include "../components/Transform.hpp"
#include "../components/Renderable.hpp"
#include "../components/CharacterSlots.hpp"
#include "../components/Character.hpp"
#include "../components/PlayerController.hpp"
#include <OgreEntity.h>
#include <OgreLogManager.h>
#include <OgreSceneNode.h>
#include <cmath>
#include <iostream>

#ifdef EDITSCENE_VERBOSE_DEBUG
#define EDITSCENE_LOG_ANIM(msg) \
	Ogre::LogManager::getSingleton().logMessage(msg)
#else
#define EDITSCENE_LOG_ANIM(msg)
#endif

static unsigned short findBoneBlendIndex(Ogre::SkeletonInstance *skel,
					 Ogre::Bone *bone)
{
	unsigned short handle = bone->getHandle();
	if (handle < skel->getNumBones() && skel->getBone(handle) == bone)
		return handle;
	/* Fallback: search by pointer in case handles are not sequential */
	for (unsigned short i = 0; i < skel->getNumBones(); ++i) {
		if (skel->getBone(i) == bone)
			return i;
	}
	return (unsigned short)-1;
}

AnimationTreeSystem::AnimationTreeSystem(flecs::world &world,
					 Ogre::SceneManager *sceneMgr)
	: m_world(world)
	, m_sceneMgr(sceneMgr)
{
	m_world.observer<AnimationTreeComponent>("AnimationTreeCleanup")
		.event(flecs::OnRemove)
		.each([this](flecs::entity e, AnimationTreeComponent &) {
			teardownEntity(e);
		});
}

AnimationTreeSystem::~AnimationTreeSystem()
{
	std::vector<flecs::entity_t> toRemove;
	for (auto &pair : m_states)
		toRemove.push_back(pair.first);
	for (auto id : toRemove)
		teardownEntity(m_world.entity(id));
}

void AnimationTreeSystem::initialize()
{
	m_initialized = true;
}

Ogre::Entity *AnimationTreeSystem::findAnimatedEntity(flecs::entity e)
{
	if (!e.is_alive())
		return nullptr;

	CharacterSlotSystem *slotSystem =
		CharacterSlotSystem::getSingletonPtr();
	if (slotSystem) {
		Ogre::Entity *masterEnt = slotSystem->getMasterEntity(e);
		if (masterEnt && masterEnt->hasSkeleton())
			return masterEnt;
	}

	if (e.has<RenderableComponent>()) {
		auto &rc = e.get<RenderableComponent>();
		if (rc.entity && rc.entity->hasSkeleton())
			return rc.entity;
	}

	if (e.has<TransformComponent>()) {
		auto &t = e.get<TransformComponent>();
		if (t.node) {
			for (unsigned int i = 0;
			     i < t.node->numAttachedObjects(); ++i) {
				Ogre::MovableObject *obj =
					t.node->getAttachedObject(i);
				if (obj->getMovableType() == "Entity") {
					Ogre::Entity *ent =
						static_cast<Ogre::Entity *>(
							obj);
					if (ent->hasSkeleton())
						return ent;
				}
			}
		}
	}

	return nullptr;
}

bool AnimationTreeSystem::setupEntity(flecs::entity e,
				      AnimationTreeComponent &at)
{
	Ogre::Entity *ent = findAnimatedEntity(e);
	std::cout << "AnimationTreeSystem::setupEntity: entity=" << e.id()
		  << " ent=" << ent
		  << " hasSkeleton=" << (ent ? ent->hasSkeleton() : false)
		  << std::endl;
	if (!ent || !ent->hasSkeleton()) {
		teardownEntity(e);
		return false;
	}

	EntityAnimTreeState &state = m_states[e.id()];
	/* Preserve root-motion tracking state across rebuilds so
	 * setupEntity() does not produce a one-frame zero-delta
	 * stutter when the mesh is recreated by CharacterSlotSystem.
	 */
	std::unordered_map<Ogre::String, std::pair<Ogre::Vector3, bool> >
		prevRootMotion;
	for (const auto &pair : state.animations) {
		prevRootMotion[pair.first] = { pair.second.prevRootPos,
					       pair.second.hasPrevRootPos };
	}
	state.ogreEntity = ent;
	state.ogreEntityName = ent->getName();
	state.rootBone = nullptr;
	state.animations.clear();

	Ogre::SkeletonInstance *skel = ent->getSkeleton();

	/* Find and freeze root bone */
	const char *rootNames[] = { "Root", "mixamorig:Hips", "Spineroot" };
	for (const char *name : rootNames) {
		try {
			state.rootBone = skel->getBone(name);
			break;
		} catch (...) {
		}
	}
	if (state.rootBone) {
		state.rootBindingPosition =
			state.rootBone->getInitialPosition();
		state.rootBindingOrientation =
			state.rootBone->getInitialOrientation();
		state.rootBindingScale = state.rootBone->getInitialScale();
		state.rootBone->setManuallyControlled(true);
		state.rootBone->setPosition(state.rootBindingPosition);
		state.rootBone->setOrientation(state.rootBindingOrientation);
		state.rootBone->setScale(state.rootBindingScale);
	}

	/* Discover all animations and precompute root tracks */
	Ogre::AnimationStateSet *states = ent->getAllAnimationStates();
	if (states) {
		for (const auto &pair : states->getAnimationStates()) {
			const Ogre::String &animName = pair.first;
			Ogre::AnimationState *as = pair.second;

			/* Use blend mask to exclude root bone from animation
			 * (setManuallyControlled alone does not prevent
			 *  Animation::apply() in OGRE 14) */
			if (state.rootBone) {
				as->destroyBlendMask();
				as->createBlendMask(skel->getNumBones(), 1.0f);
				unsigned short blendIdx = findBoneBlendIndex(
					skel, state.rootBone);
				if (blendIdx != (unsigned short)-1) {
					as->setBlendMaskEntry(blendIdx, 0.0f);
				} else {
					std::cout
						<< "AnimationTreeSystem::setupEntity: WARNING could not find blend mask index for root bone "
						<< state.rootBone->getName()
						<< " handle="
						<< state.rootBone->getHandle()
						<< std::endl;
				}
			}

			AnimationRuntimeInfo info;
			info.ogreAnimState = as;

			try {
				Ogre::Animation *anim =
					skel->getAnimation(animName);
				if (state.rootBone) {
					unsigned short handle =
						state.rootBone->getHandle();
					if (anim->hasNodeTrack(handle)) {
						info.rootTrack =
							anim->getNodeTrack(
								handle);
						if (info.rootTrack
							    ->getNumKeyFrames() >=
						    2) {
							/* Use interpolated keyframes at the exact
							 * time boundaries (0 and length) rather
							 * than the first/last stored keyframes.
							 * Exporters sometimes place keyframes
							 * slightly inside the boundaries, which
							 * would make loopTranslation slightly
							 * wrong and cause a backward lurch on
							 * wrap frames. */
							Ogre::TransformKeyFrame
								tkfBeg(info.rootTrack,
								       0.0f);
							Ogre::TransformKeyFrame
								tkfEnd(info.rootTrack,
								       0.0f);
							info.rootTrack
								->getInterpolatedKeyFrame(
									0.0f,
									&tkfBeg);
							info.rootTrack->getInterpolatedKeyFrame(
								as->getLength(),
								&tkfEnd);
							info.loopTranslation =
								tkfEnd.getTranslate() -
								tkfBeg.getTranslate();
						}
					}
				}
			} catch (...) {
			}

			state.animations[animName] = info;
		}
	}

	/* Restore root-motion tracking state for animations that
	 * existed before the rebuild.
	 */
	for (auto &pair : state.animations) {
		auto it = prevRootMotion.find(pair.first);
		if (it != prevRootMotion.end()) {
			pair.second.prevRootPos = it->second.first;
			pair.second.hasPrevRootPos = it->second.second;
		}
	}

	/* Initialize default state machine states */
	initializeTreeStates(at.root, at);
	std::cout
		<< "  setupEntity done: animations=" << state.animations.size()
		<< " rootBone="
		<< (state.rootBone ? state.rootBone->getName() : "null")
		<< std::endl;
	return true;
}

void AnimationTreeSystem::teardownEntity(flecs::entity e)
{
	auto it = m_states.find(e.id());
	if (it == m_states.end())
		return;

	EntityAnimTreeState &state = it->second;

	/*
	 * The Ogre::Entity may have already been destroyed by
	 * EditorUISystem::deleteEntity before the Flecs OnRemove
	 * observer fires.  In that case the AnimationState pointers
	 * are dangling and we must not touch them.
	 *
	 * We detect this by checking whether the cached Ogre::Entity
	 * name is still known to the scene manager.  We use the
	 * cached name string (ogreEntityName) rather than
	 * ogreEntity->getName() because the pointer may already be
	 * dangling.
	 */
	bool ogreEntityAlive = false;
	if (!state.ogreEntityName.empty()) {
		try {
			ogreEntityAlive =
				m_sceneMgr->hasEntity(state.ogreEntityName);
		} catch (...) {
			ogreEntityAlive = false;
		}
	}

	if (ogreEntityAlive) {
		for (auto &pair : state.animations) {
			if (pair.second.ogreAnimState)
				pair.second.ogreAnimState->destroyBlendMask();
		}

		disableAllAnimations(state);

		if (state.rootBone) {
			state.rootBone->setManuallyControlled(false);
			state.rootBone = nullptr;
		}
	}

	state.animations.clear();
	state.fadeStates.clear();
	state.prevStates.clear();

	m_states.erase(it);
}

void AnimationTreeSystem::disableAllAnimations(EntityAnimTreeState &state)
{
	for (auto &pair : state.animations) {
		if (pair.second.ogreAnimState)
			pair.second.ogreAnimState->setEnabled(false);
	}
}

void AnimationTreeSystem::initializeTreeStates(const AnimationTreeNode &node,
					       AnimationTreeComponent &at)
{
	if (node.type == "stateMachine") {
		if (at.currentStates.find(node.name) ==
		    at.currentStates.end()) {
			for (const auto &child : node.children) {
				if (child.type == "state") {
					at.currentStates[node.name] =
						child.name;
					break;
				}
			}
		}
	}
	for (const auto &child : node.children)
		initializeTreeStates(child, at);
}

void AnimationTreeSystem::resolveTemplate(AnimationTreeComponent &at)
{
	if (at.templateName.empty())
		return;

	AnimationTreeTemplate *templ = nullptr;
	AnimationTreeComponent *templAt = nullptr;

	m_world.query<AnimationTreeTemplate, AnimationTreeComponent>().each(
		[&](flecs::entity, AnimationTreeTemplate &t,
		    AnimationTreeComponent &ta) {
			if (t.name == at.templateName) {
				templ = &t;
				templAt = &ta;
			}
		});

	if (!templ) {
		m_world.query<AnimationTreeTemplate>().each(
			[&](flecs::entity, AnimationTreeTemplate &t) {
				if (t.name == at.templateName)
					templ = &t;
			});
	}

	if (templ && at.templateVersion != templ->version) {
		if (templAt)
			at.root = templAt->root;
		at.templateVersion = templ->version;
		at.dirty = true;
	}
}

void AnimationTreeSystem::update(float deltaTime)
{
	if (!m_initialized)
		return;

	m_world.query<AnimationTreeComponent>().each([this, deltaTime](
							     flecs::entity e,
							     AnimationTreeComponent
								     &at) {
		/* Reset root-motion velocity for all entities.
		 * Root motion velocity will be set below only for
		 * entities with active root motion. */
		if (e.has<CharacterComponent>())
			e.get_mut<CharacterComponent>().linearVelocity =
				Ogre::Vector3::ZERO;
		resolveTemplate(at);

		if (at.dirty) {
			if (setupEntity(e, at))
				at.dirty = false;
		}

		auto it = m_states.find(e.id());
		if (it == m_states.end())
			return;

		/* Validate cached entity pointer -
			 * CharacterSlotSystem rebuilds can destroy and
			 * recreate the Ogre::Entity */
		Ogre::Entity *currentEnt = findAnimatedEntity(e);
		if (currentEnt != it->second.ogreEntity) {
			if (!setupEntity(e, at)) {
				return;
			}
			it = m_states.find(e.id());
			if (it == m_states.end())
				return;
		}

		EntityAnimTreeState &state = it->second;
		if (!state.ogreEntity)
			return;

#ifdef EDITSCENE_VERBOSE_DEBUG
		/* Diagnostic: check if OGRE moved the root bone during rendering */
		if (state.rootBone) {
			Ogre::Vector3 currentPos =
				state.rootBone->getPosition();
			Ogre::Vector3 diff =
				currentPos - state.rootBindingPosition;
			if (diff.squaredLength() > 0.0001f) {
				std::cout
					<< "[ANIM_DBG] ROOT BONE MOVED! entity="
					<< e.id() << " current=" << currentPos
					<< " binding="
					<< state.rootBindingPosition
					<< " diff=" << diff << std::endl;
			}
		}
#endif

		if (!at.enabled) {
			disableAllAnimations(state);
			return;
		}

		/* Sync external state changes (editor) */
		syncStateChanges(at, state);

		/* Evaluate tree */
		EvalContext ctx;
		ctx.deltaTime = deltaTime;
		evaluateNode(at.root, 1.0f, 1.0f, at, state, ctx);

		/* Apply animation weights and advance */
		Ogre::Vector3 totalRootMotion = Ogre::Vector3::ZERO;
		Ogre::SceneNode *sceneNode =
			state.ogreEntity->getParentSceneNode();

		for (auto &pair : ctx.animData) {
			const Ogre::String &animName = pair.first;
			AnimEvalData &data = pair.second;

			auto itAnim = state.animations.find(animName);
			if (itAnim == state.animations.end())
				continue;
			AnimationRuntimeInfo &info = itAnim->second;

			if (!info.ogreAnimState)
				continue;

			bool active = data.weight > 0.001f;
			info.ogreAnimState->setEnabled(active);
			info.ogreAnimState->setWeight(
				Ogre::Math::Clamp(data.weight, 0.0f, 1.0f));

			if (active && data.timeDelta != 0.0f) {
				float lastTime =
					info.ogreAnimState->getTimePosition();
				info.ogreAnimState->addTime(data.timeDelta);
				float thisTime =
					info.ogreAnimState->getTimePosition();

				if (at.useRootMotion && info.rootTrack) {
					Ogre::TransformKeyFrame tkf(nullptr,
								    0.0f);
					info.rootTrack->getInterpolatedKeyFrame(
						thisTime, &tkf);
					Ogre::Vector3 thisPos =
						tkf.getTranslate();

					/*
						 * For looping animations, use the average
						 * cycle velocity (loopTranslation / duration)
						 * instead of exact frame-to-frame displacement.
						 * The raw root bone oscillates back and forth
						 * within a walk cycle; transferring that
						 * oscillation to world position makes the
						 * whole character jitter.  The average velocity
						 * captures only the net forward movement.
						 * Non-looping animations still use exact
						 * displacement so acceleration curves are
						 * preserved.
						 */
					float animLen =
						info.ogreAnimState->getLength();
					if (animLen > 0.0001f &&
					    info.loopTranslation.squaredLength() >
						    0.0001f) {
						Ogre::Vector3 avgVelLocal =
							info.loopTranslation /
							animLen;
						totalRootMotion +=
							avgVelLocal *
							data.timeDelta *
							data.weight;
					} else if (info.hasPrevRootPos) {
						Ogre::Vector3 displacement =
							thisPos -
							info.prevRootPos;
						if (thisTime < lastTime)
							displacement +=
								info.loopTranslation;
						totalRootMotion +=
							displacement *
							data.weight;
					} else {
						/* Fallback for first frame:
							 * instantaneous velocity */
						float sampleDt = 0.001f;
						float nextTime =
							thisTime + sampleDt;
						bool wrapped = false;
						if (nextTime >
						    info.ogreAnimState
							    ->getLength()) {
							nextTime -=
								info.ogreAnimState
									->getLength();
							wrapped = true;
						}
						Ogre::TransformKeyFrame tkfNext(
							nullptr, 0.0f);
						info.rootTrack
							->getInterpolatedKeyFrame(
								nextTime,
								&tkfNext);
						Ogre::Vector3 nextPos =
							tkfNext.getTranslate();
						if (wrapped)
							nextPos +=
								info.loopTranslation;
						Ogre::Vector3 velocityLocal =
							(nextPos - thisPos) /
							sampleDt;
						totalRootMotion +=
							velocityLocal *
							data.timeDelta *
							data.weight;
					}

					info.prevRootPos = thisPos;
					info.hasPrevRootPos = true;
				}
			}
		}

		if (at.useRootMotion && sceneNode) {
			/* Smooth root-motion displacement with EMA to
				 * filter high-frequency hip oscillation within
				 * the walk cycle.  Alpha = 0.15 gives strong
				 * attenuation of ~2 Hz walk-cycle jitter while
				 * still responding to stops within ~6 frames. */
			float alpha = 0.15f;
			if (state.hasSmoothedRootMotion) {
				totalRootMotion = state.smoothedRootMotion *
							  (1.0f - alpha) +
						  totalRootMotion * alpha;
			}
			state.smoothedRootMotion = totalRootMotion;
			state.hasSmoothedRootMotion = true;

			Ogre::Quaternion worldRot =
				sceneNode->_getDerivedOrientation();
			Ogre::Vector3 displacementWorld =
				worldRot * totalRootMotion;

			if (e.has<CharacterComponent>()) {
				auto &cc = e.get_mut<CharacterComponent>();
				cc.useRootMotion = true;
				if (deltaTime > 0.0000001f) {
					cc.linearVelocity =
						displacementWorld / deltaTime;
				}
			}
		}
#ifdef EDITSCENE_VERBOSE_DEBUG
		/* Diagnostic: for player-controlled characters in the idle
		 * locomotion state, log which Ogre animations are still
		 * active and the state-machine fade weights. This helps
		 * catch cases where walk/run keeps playing after idle is
		 * requested. */
		if (e.has<PlayerControlledComponent>()) {
			auto itLoco = at.currentStates.find("locomotion");
			if (itLoco != at.currentStates.end() &&
			    itLoco->second == "idle") {
				Ogre::String animList;
				bool hasNonIdle = false;
				int activeCount = 0;
				for (const auto &pair : ctx.animData) {
					if (pair.second.weight > 0.001f) {
						activeCount++;
						animList += pair.first + "=" +
							    Ogre::StringConverter::toString(
								    pair.second.weight) +
							    " ";
						if (pair.first != "idle")
							hasNonIdle = true;
					}
				}
				auto itFade = state.fadeStates.find("locomotion");
				Ogre::String fadeList;
				bool fading = false;
				if (itFade != state.fadeStates.end()) {
					for (const auto &w : itFade->second.weights) {
						fadeList += w.first + "=" +
							    Ogre::StringConverter::toString(
								    w.second) +
							    " ";
					}
					fadeList += "fadeIn={";
					for (const auto &n : itFade->second.fadeIn) {
						fadeList += n + " ";
						fading = true;
					}
					fadeList += "} fadeOut={";
					for (const auto &n : itFade->second.fadeOut) {
						fadeList += n + " ";
						fading = true;
					}
					fadeList += "}";
				}
				if (activeCount > 1 || hasNonIdle || fading) {
					EDITSCENE_LOG_ANIM("[ANIM_DBG] player idle anims: " +
							   animList + "| fade: " + fadeList);
				}
			}
		}
#endif

		/* Reset root bone to binding pose */
		if (state.rootBone) {
			state.rootBone->setPosition(state.rootBindingPosition);
			state.rootBone->setOrientation(
				state.rootBindingOrientation);
			state.rootBone->setScale(state.rootBindingScale);
		}

		/* Handle end-of-animation transitions */
		checkEndTransitions(e, at, state, ctx);
	});
}

void AnimationTreeSystem::evaluateNode(const AnimationTreeNode &node,
				       float parentWeight, float timeMul,
				       const AnimationTreeComponent &at,
				       EntityAnimTreeState &state,
				       EvalContext &ctx)
{
	if (parentWeight < 0.001f)
		return;

	if (node.type == "output") {
		if (!node.children.empty())
			evaluateNode(node.children[0], parentWeight,
				     timeMul * node.speed, at, state, ctx);
	} else if (node.type == "speed") {
		if (!node.children.empty())
			evaluateNode(node.children[0], parentWeight,
				     timeMul * node.speed, at, state, ctx);
	} else if (node.type == "stateMachine") {
		auto itCurrent = at.currentStates.find(node.name);
		if (itCurrent == at.currentStates.end())
			return;

		auto &fade = state.fadeStates[node.name];
		const Ogre::String &currentName = itCurrent->second;

		/* Ensure weight entries exist */
		for (const auto &child : node.children) {
			if (child.type == "state" &&
			    fade.weights.find(child.name) ==
				    fade.weights.end()) {
				fade.weights[child.name] =
					(child.name == currentName) ? 1.0f :
								      0.0f;
			}
		}

		/* Update fades */
		std::vector<Ogre::String> doneIn;
		for (const auto &inName : fade.fadeIn) {
			fade.weights[inName] += ctx.deltaTime * node.fadeSpeed;
			if (fade.weights[inName] >= 1.0f) {
				fade.weights[inName] = 1.0f;
				doneIn.push_back(inName);
			}
		}
		for (const auto &n : doneIn)
			fade.fadeIn.erase(n);

		std::vector<Ogre::String> doneOut;
		for (const auto &outName : fade.fadeOut) {
			fade.weights[outName] -= ctx.deltaTime * node.fadeSpeed;
			if (fade.weights[outName] <= 0.0f) {
				fade.weights[outName] = 0.0f;
				doneOut.push_back(outName);
			}
		}
		for (const auto &n : doneOut)
			fade.fadeOut.erase(n);

		/* Evaluate active states */
		for (const auto &child : node.children) {
			if (child.type != "state")
				continue;
			float w = fade.weights[child.name];
			if (w > 0.001f)
				evaluateNode(child, parentWeight * w, timeMul,
					     at, state, ctx);
		}

		/* Queue end-transition check */
		if (node.endTransitions.find(currentName) !=
		    node.endTransitions.end())
			ctx.endChecks.push_back({ node.name, currentName });
	} else if (node.type == "state") {
		if (!node.children.empty())
			evaluateNode(node.children[0], parentWeight, timeMul,
				     at, state, ctx);
	} else if (node.type == "animation") {
		auto &data = ctx.animData[node.animationName];
		data.weight += parentWeight;
		data.timeDelta = ctx.deltaTime * timeMul;
	}
}

void AnimationTreeSystem::syncStateChanges(AnimationTreeComponent &at,
					   EntityAnimTreeState &state)
{
	for (auto &pair : at.currentStates) {
		auto itPrev = state.prevStates.find(pair.first);
		if (itPrev == state.prevStates.end() ||
		    itPrev->second != pair.second) {
			/* State changed externally (editor); set up fade */
			auto &fade = state.fadeStates[pair.first];
			if (itPrev != state.prevStates.end())
				fade.fadeOut.insert(itPrev->second);
			fade.fadeIn.insert(pair.second);
			fade.weights[pair.second] = 0.0f;
			state.prevStates[pair.first] = pair.second;
		}
	}
	/* Also track newly added state machines */
	for (auto &pair : at.currentStates) {
		if (state.prevStates.find(pair.first) == state.prevStates.end())
			state.prevStates[pair.first] = pair.second;
	}
}

void AnimationTreeSystem::checkEndTransitions(flecs::entity e,
					      AnimationTreeComponent &at,
					      EntityAnimTreeState &state,
					      EvalContext &ctx)
{
	for (auto &pair : ctx.endChecks) {
		const Ogre::String &smName = pair.first;
		const Ogre::String &stateName = pair.second;

		const AnimationTreeNode *smNode =
			findStateMachineNode(at.root, smName);
		if (!smNode)
			continue;

		auto itTrans = smNode->endTransitions.find(stateName);
		if (itTrans == smNode->endTransitions.end())
			continue;

		const AnimationTreeNode *stNode =
			findStateNode(*smNode, stateName);
		if (!stNode)
			continue;
		const AnimationTreeNode *animNode = findAnimationNode(*stNode);
		if (!animNode)
			continue;

		auto itAnim = state.animations.find(animNode->animationName);
		if (itAnim == state.animations.end())
			continue;
		Ogre::AnimationState *as = itAnim->second.ogreAnimState;
		if (!as)
			continue;

		float t = as->getTimePosition();
		float len = as->getLength();
		if (len > 0.0f && t >= len * 0.99f)
			setStateInternal(e, at, state, smName, itTrans->second,
					 true);
	}
}

void AnimationTreeSystem::setState(flecs::entity e,
				   const Ogre::String &stateMachineName,
				   const Ogre::String &stateName, bool reset)
{
	auto it = m_states.find(e.id());
	if (it == m_states.end())
		return;
	if (!e.has<AnimationTreeComponent>())
		return;
	auto &at = e.get_mut<AnimationTreeComponent>();
	setStateInternal(e, at, it->second, stateMachineName, stateName, reset);
}

void AnimationTreeSystem::setStateInternal(flecs::entity e,
					   AnimationTreeComponent &at,
					   EntityAnimTreeState &state,
					   const Ogre::String &stateMachineName,
					   const Ogre::String &stateName,
					   bool reset)
{
	(void)e;
	auto itCurrent = at.currentStates.find(stateMachineName);
	if (itCurrent == at.currentStates.end())
		return;
	if (itCurrent->second == stateName)
		return;

	auto &fade = state.fadeStates[stateMachineName];
	const Ogre::String &oldState = itCurrent->second;

	/* If the new state is currently fading out, stop fading it out; if the
	 * old state is still fading in, stop fading it in. Otherwise a state can
	 * end up in both fadeIn and fadeOut simultaneously and its weight gets
	 * stuck. */
	fade.fadeOut.erase(stateName);
	fade.fadeIn.erase(oldState);

	fade.fadeOut.insert(oldState);
	fade.fadeIn.insert(stateName);
	if (fade.weights.find(stateName) == fade.weights.end())
		fade.weights[stateName] = 0.0f;

	itCurrent->second = stateName;
	state.prevStates[stateMachineName] = stateName;

	if (reset) {
		/* Reset the new state's animation */
		const AnimationTreeNode *smNode =
			findStateMachineNode(at.root, stateMachineName);
		if (smNode) {
			const AnimationTreeNode *stNode =
				findStateNode(*smNode, stateName);
			if (stNode) {
				const AnimationTreeNode *animNode =
					findAnimationNode(*stNode);
				if (animNode) {
					auto itAnim = state.animations.find(
						animNode->animationName);
					if (itAnim != state.animations.end() &&
					    itAnim->second.ogreAnimState) {
						itAnim->second.ogreAnimState
							->setTimePosition(0.0f);
						/* Reset root motion tracking
						 * so the first frame of the
						 * new animation doesn't
						 * produce a large delta from
						 * the previous animation's
						 * root position. */
						itAnim->second.hasPrevRootPos =
							false;
					}
				}
			}
		}
	}
}

Ogre::String
AnimationTreeSystem::getCurrentState(flecs::entity e,
				     const Ogre::String &stateMachineName)
{
	if (!e.has<AnimationTreeComponent>())
		return "";
	auto &at = e.get<AnimationTreeComponent>();
	auto it = at.currentStates.find(stateMachineName);
	if (it != at.currentStates.end())
		return it->second;
	return "";
}

const AnimationTreeNode *
AnimationTreeSystem::findStateMachineNode(const AnimationTreeNode &root,
					  const Ogre::String &name) const
{
	if (root.type == "stateMachine" && root.name == name)
		return &root;
	for (const auto &child : root.children) {
		const AnimationTreeNode *found =
			findStateMachineNode(child, name);
		if (found)
			return found;
	}
	return nullptr;
}

const AnimationTreeNode *
AnimationTreeSystem::findStateNode(const AnimationTreeNode &smNode,
				   const Ogre::String &stateName) const
{
	for (const auto &child : smNode.children) {
		if (child.type == "state" && child.name == stateName)
			return &child;
	}
	return nullptr;
}

const AnimationTreeNode *
AnimationTreeSystem::findAnimationNode(const AnimationTreeNode &stateNode) const
{
	if (stateNode.type == "animation")
		return &stateNode;
	for (const auto &child : stateNode.children) {
		const AnimationTreeNode *found = findAnimationNode(child);
		if (found)
			return found;
	}
	return nullptr;
}
