#ifndef EDITSCENE_ANIMATIONTREESYSTEM_HPP
#define EDITSCENE_ANIMATIONTREESYSTEM_HPP
#pragma once

#include <flecs.h>
#include <Ogre.h>
#include <memory>
#include <unordered_map>
#include <set>

#include "../components/AnimationTree.hpp"
#include "../components/AnimationTreeTemplate.hpp"


/**
 * System that evaluates an AnimationTreeComponent each frame.
 *
 * Walks the tree recursively to:
 *   - Determine active animation states and their blend weights
 *   - Advance Ogre::AnimationState::addTime() for active animations
 *   - Compute and apply root motion (optional)
 *   - Handle end-of-animation auto-transitions
 */
class AnimationTreeSystem {
public:
	AnimationTreeSystem(flecs::world &world, Ogre::SceneManager *sceneMgr);
	~AnimationTreeSystem();

	void initialize();
	void update(float deltaTime);

	/**
	 * Manually switch a state machine to a new state.
	 * Sets up cross-fade automatically.
	 */
	void setState(flecs::entity e, const Ogre::String &stateMachineName,
		      const Ogre::String &stateName, bool reset = false);

	/**
	 * Find the Ogre entity with a skeleton for a given flecs entity.
	 */
	static Ogre::Entity *findAnimatedEntity(flecs::entity e);

	/**
	 * Get current state name of a state machine (for editor display).
	 */
	Ogre::String getCurrentState(flecs::entity e,
				     const Ogre::String &stateMachineName);

private:
	struct AnimationRuntimeInfo {
		Ogre::AnimationState *ogreAnimState = nullptr;
		Ogre::NodeAnimationTrack *rootTrack = nullptr;
		Ogre::Vector3 loopTranslation = Ogre::Vector3::ZERO;
		/* Previous root bone position for root motion delta
		 * computation. Used to detect animation wrapping and
		 * compute smooth deltas across loop boundaries. */
		Ogre::Vector3 prevRootPos = Ogre::Vector3::ZERO;
		bool hasPrevRootPos = false;
	};

	struct FadeInfo {
		std::set<Ogre::String> fadeIn;
		std::set<Ogre::String> fadeOut;
		std::unordered_map<Ogre::String, float> weights;
	};

	struct EntityAnimTreeState {
		Ogre::Entity *ogreEntity = nullptr;
		Ogre::String ogreEntityName; /* cached for safe teardown */
		Ogre::Bone *rootBone = nullptr;
		Ogre::Vector3 rootBindingPosition;
		Ogre::Quaternion rootBindingOrientation;
		Ogre::Vector3 rootBindingScale;

		/* Root motion unapply/reapply state */
		Ogre::Vector3 appliedRootTranslation = Ogre::Vector3::ZERO;
		Ogre::Quaternion appliedRootRotation =
			Ogre::Quaternion::IDENTITY;

		/* Previous root-motion velocity for smoothing */
		Ogre::Vector3 prevRootMotionVel = Ogre::Vector3::ZERO;
		/* Smoothed root-motion displacement (EMA) */
		Ogre::Vector3 smoothedRootMotion = Ogre::Vector3::ZERO;
		bool hasSmoothedRootMotion = false;

		std::unordered_map<Ogre::String, AnimationRuntimeInfo>
			animations;
		std::unordered_map<Ogre::String, FadeInfo> fadeStates;
		/* Track previous currentStates to detect external changes */
		std::unordered_map<Ogre::String, Ogre::String> prevStates;
	};

	struct AnimEvalData {
		float weight = 0.0f;
		float timeDelta = 0.0f;
	};

	struct EvalContext {
		float deltaTime = 0.0f;
		std::unordered_map<Ogre::String, AnimEvalData> animData;
		std::vector<std::pair<Ogre::String, Ogre::String> > endChecks;
	};

	bool setupEntity(flecs::entity e, AnimationTreeComponent &at);
	void teardownEntity(flecs::entity e);
	void disableAllAnimations(EntityAnimTreeState &state);
	void initializeTreeStates(const AnimationTreeNode &node,
				  AnimationTreeComponent &at);
	void evaluateNode(const AnimationTreeNode &node, float parentWeight,
			  float timeMul, const AnimationTreeComponent &at,
			  EntityAnimTreeState &state, EvalContext &ctx);
	void checkEndTransitions(flecs::entity e, AnimationTreeComponent &at,
				 EntityAnimTreeState &state, EvalContext &ctx);
	void syncStateChanges(AnimationTreeComponent &at,
			      EntityAnimTreeState &state);
	void setStateInternal(flecs::entity e, AnimationTreeComponent &at,
			      EntityAnimTreeState &state,
			      const Ogre::String &stateMachineName,
			      const Ogre::String &stateName, bool reset);

	/* Resolve template reference and copy tree if template changed */
	void resolveTemplate(AnimationTreeComponent &at);

	const AnimationTreeNode *
	findStateMachineNode(const AnimationTreeNode &root,
			     const Ogre::String &name) const;
	const AnimationTreeNode *
	findStateNode(const AnimationTreeNode &smNode,
		      const Ogre::String &stateName) const;
	const AnimationTreeNode *
	findAnimationNode(const AnimationTreeNode &stateNode) const;

	flecs::world &m_world;
	Ogre::SceneManager *m_sceneMgr;
	bool m_initialized = false;

	std::unordered_map<flecs::entity_t, EntityAnimTreeState> m_states;
};

#endif // EDITSCENE_ANIMATIONTREESYSTEM_HPP
