#include "HairPhysicsSystem.hpp"
#include "CharacterSlotSystem.hpp"
#include "CharacterRegistry.hpp"
#include "../components/CharacterIdentity.hpp"
#include "../components/CharacterSlots.hpp"
#include "../components/Character.hpp"
#include <OgreEntity.h>
#include <OgreLogManager.h>
#include <OgreSceneNode.h>
#include <OgreTagPoint.h>

#include <queue>

namespace {
	/**
	 * Find the master bone the hair is attached to.
	 * Prefer the TagPoint's parent bone; fall back to the named head bone.
	 */
	Ogre::Bone *findAttachBone(Ogre::Entity *hairEnt,
				   Ogre::Entity *masterEnt)
	{
		Ogre::TagPoint *tagPoint = dynamic_cast<Ogre::TagPoint *>(
			hairEnt->getParentNode());
		if (tagPoint) {
			Ogre::Bone *parentBone = dynamic_cast<Ogre::Bone *>(
				tagPoint->getParent());
			if (parentBone)
				return parentBone;
		}

		if (masterEnt && masterEnt->hasSkeleton()) {
			Ogre::SkeletonInstance *skel = masterEnt->getSkeleton();
			if (skel) {
				try {
					return skel->getBone("mixamorig:Head");
				} catch (...) {
				}
			}
		}
		return nullptr;
	}

	/**
	 * World transform of the attachment bone, including the master entity's
	 * scene-node transform.
	 */
	Ogre::Matrix4 getAttachBoneWorldMatrix(Ogre::Entity *masterEnt,
					       Ogre::Bone *attachBone)
	{
		if (!attachBone)
			return Ogre::Matrix4::IDENTITY;

		Ogre::Matrix4 boneWorld = attachBone->_getFullTransform();
		if (masterEnt) {
			Ogre::SceneNode *node = masterEnt->getParentSceneNode();
			if (node)
				boneWorld = node->_getFullTransform() * boneWorld;
			}
		return boneWorld;
	}
} // namespace

#include <Jolt/Physics/Ragdoll/Ragdoll.h>
#include <Jolt/Physics/Collision/Shape/CapsuleShape.h>
#include <Jolt/Physics/Constraints/SwingTwistConstraint.h>
#include <Jolt/Skeleton/SkeletonPose.h>

uint32_t HairPhysicsSystem::s_nextCollisionGroup = 1;

HairPhysicsSystem::HairPhysicsSystem(flecs::world &world,
				     Ogre::SceneManager *sceneMgr,
				     JoltPhysicsWrapper *physics,
				     CharacterSlotSystem *slotSystem)
	: m_world(world)
	, m_sceneMgr(sceneMgr)
	, m_physics(physics)
	, m_slotSystem(slotSystem)
{
}

HairPhysicsSystem::~HairPhysicsSystem()
{
	for (auto &entityPair : m_states) {
		for (auto &slotPair : entityPair.second) {
			destroyRagdoll(slotPair.second);
		}
	}
	m_states.clear();
}

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

void HairPhysicsSystem::prePhysicsUpdate()
{
	if (!m_initialized || !m_physics)
		return;

	m_world.query<CharacterIdentityComponent>().each(
		[this](flecs::entity e, CharacterIdentityComponent &ci) {
			Ogre::Entity *masterEnt =
				m_slotSystem->getMasterEntity(e);
			if (!masterEnt)
				return;

			const CharacterRegistry::CharacterRecord *rec =
				CharacterRegistry::getSingleton().findCharacter(
					ci.registryId);
			if (!rec)
				return;

			auto slots = CharacterSlotSystem::getSlots(
				rec->age, rec->inlineSex);
			for (const auto &slot : slots) {
				Ogre::Entity *partEnt =
					m_slotSystem->getSlotEntity(e, slot);
				if (!partEnt || !partEnt->hasSkeleton())
					continue;

				auto &slotMap = m_states[e.id()];
				auto it = slotMap.find(slot);
				if (it == slotMap.end() ||
				    it->second.hairEntity != partEnt) {
					/* Need to create or recreate ragdoll */
					if (it != slotMap.end())
						destroyRagdoll(it->second);
					createRagdoll(e, slot, partEnt,
						      masterEnt);
					it = slotMap.find(slot);
				}

				if (it != slotMap.end())
					syncRootToHead(it->second);
			}

			/* Clean up ragdolls for slots that no longer have
			 * skeleton parts. */
			auto &slotMap = m_states[e.id()];
			std::vector<Ogre::String> toRemove;
			for (auto &pair : slotMap) {
				Ogre::Entity *partEnt =
					m_slotSystem->getSlotEntity(e,
								    pair.first);
				if (!partEnt || !partEnt->hasSkeleton()) {
					toRemove.push_back(pair.first);
					continue;
				}
			}
			for (const auto &slot : toRemove) {
				destroyRagdoll(slotMap[slot]);
				slotMap.erase(slot);
			}
		});
}

void HairPhysicsSystem::postPhysicsUpdate()
{
	if (!m_initialized || !m_physics)
		return;

	for (auto &entityPair : m_states) {
		for (auto &slotPair : entityPair.second) {
			syncPhysicsToSkeleton(slotPair.second);
		}
	}
}

void HairPhysicsSystem::createRagdoll(flecs::entity e,
				      const Ogre::String &slot,
				      Ogre::Entity *hairEnt,
				      Ogre::Entity *masterEnt)
{
	Ogre::SkeletonInstance *skel = hairEnt->getSkeleton();
	if (!skel)
		return;

	/* Build ordered bone list via BFS so parents always precede children. */
	std::vector<Ogre::Bone *> orderedBones;
	std::unordered_map<Ogre::Bone *, size_t> boneToIndex;
	std::queue<Ogre::Bone *> bfsQueue;

	for (unsigned short i = 0; i < skel->getNumBones(); ++i) {
		Ogre::Bone *b = skel->getBone(i);
		if (!b->getParent()) {
			bfsQueue.push(b);
		}
	}

	while (!bfsQueue.empty()) {
		Ogre::Bone *b = bfsQueue.front();
		bfsQueue.pop();
		boneToIndex[b] = orderedBones.size();
		orderedBones.push_back(b);
		for (unsigned short i = 0; i < b->numChildren(); ++i) {
			Ogre::Bone *child = static_cast<Ogre::Bone *>(b->getChild(i));
			if (child)
				bfsQueue.push(child);
		}
	}

	if (orderedBones.empty())
		return;

	/* Sanity check: if the skeleton has too many bones, it's probably
	 * a shared body skeleton rather than a hair-specific skeleton.
	 * Skip ragdoll creation to avoid physics cache overflow. */
	constexpr size_t maxHairBones = 32;
	if (orderedBones.size() > maxHairBones)
		return;

	/* Build JPH skeleton and bind-world matrices. */
	JPH::Ref<JPH::Skeleton> skeleton = new JPH::Skeleton;
	std::vector<JPH::Mat44> bindWorldMatrices;
	bindWorldMatrices.resize(orderedBones.size());

	for (size_t i = 0; i < orderedBones.size(); ++i) {
		Ogre::Bone *bone = orderedBones[i];
		Ogre::Bone *parent = static_cast<Ogre::Bone *>(bone->getParent());
		int parentIndex = -1;
		if (parent) {
			auto it = boneToIndex.find(parent);
			if (it != boneToIndex.end())
				parentIndex = (int)it->second;
		}
		skeleton->AddJoint(bone->getName().c_str(), parentIndex);

		/* Build bind-world matrix from initial pose. */
		Ogre::Matrix4 localMat = Ogre::Matrix4::IDENTITY;
		localMat.makeTransform(
			bone->getInitialPosition(),
			bone->getInitialScale(),
			bone->getInitialOrientation());
		if (parent) {
			bindWorldMatrices[i] =
				bindWorldMatrices[boneToIndex[parent]] *
				ogreMatrixToJolt(localMat);
		} else {
			bindWorldMatrices[i] = ogreMatrixToJolt(localMat);
		}
	}

	/* World transform of the attachment bone (head) on the master entity. */
	Ogre::Bone *attachBone = findAttachBone(hairEnt, masterEnt);
	if (!attachBone) {
		Ogre::LogManager::getSingleton().logMessage(
			"HairPhysicsSystem: cannot find attachment bone for " +
			slot + ", skipping ragdoll");
		return;
	}
	Ogre::Matrix4 attachBoneMat = getAttachBoneWorldMatrix(masterEnt,
						       attachBone);
	JPH::Mat44 attachBoneJolt = ogreMatrixToJolt(attachBoneMat);

	/* Build ragdoll settings. */
	JPH::Ref<JPH::RagdollSettings> settings = new JPH::RagdollSettings;
	settings->mSkeleton = skeleton;
	settings->mParts.resize(orderedBones.size());

	std::vector<int> parentIndices;
	parentIndices.resize(orderedBones.size(), -1);

	for (size_t i = 0; i < orderedBones.size(); ++i) {
		Ogre::Bone *bone = orderedBones[i];
		JPH::RagdollSettings::Part &part = settings->mParts[i];

		/* Capsule size heuristic from average child distance. */
		float halfHeight = 0.05f;
		float radius = 0.02f;
		if (bone->numChildren() > 0) {
			float avgLen = 0.0f;
			int count = 0;
			for (unsigned short ci = 0; ci < bone->numChildren();
			     ++ci) {
				Ogre::Bone *child = static_cast<Ogre::Bone *>(
					bone->getChild(ci));
				if (!child)
					continue;
				Ogre::Vector3 diff =
					child->getInitialPosition();
				avgLen += diff.length();
				++count;
			}
			if (count > 0) {
				avgLen /= count;
				halfHeight = avgLen * 0.5f;
				radius = avgLen * 0.15f;
			}
		}
		/* Clamp to reasonable bounds. */
		if (halfHeight < 0.01f)
			halfHeight = 0.01f;
		if (radius < 0.005f)
			radius = 0.005f;
		if (radius > halfHeight * 0.9f)
			radius = halfHeight * 0.9f;

		part.SetShape(new JPH::CapsuleShape(halfHeight, radius));

		/* World-space bind position. */
		JPH::Mat44 worldMat = attachBoneJolt * bindWorldMatrices[i];
		part.mPosition = JPH::RVec3(worldMat.GetTranslation());
		part.mRotation = worldMat.GetRotationSafe().GetQuaternion().Normalized();
		part.mMotionType = (i == 0) ?
					   JPH::EMotionType::Kinematic :
					   JPH::EMotionType::Dynamic;
		part.mLinearDamping = 0.5f;
		part.mAngularDamping = 0.5f;
		part.mGravityFactor = 0.5f;
		part.mMaxLinearVelocity = 5.0f;
		part.mMaxAngularVelocity = JPH::DegreesToRadians(720.0f);

		Ogre::Bone *parent = static_cast<Ogre::Bone *>(bone->getParent());
		if (parent) {
			auto it = boneToIndex.find(parent);
			if (it != boneToIndex.end())
				parentIndices[i] = (int)it->second;
		}

		if (i > 0 && parentIndices[i] >= 0) {
			JPH::SwingTwistConstraintSettings *constraint =
				new JPH::SwingTwistConstraintSettings;
			constraint->mSpace = JPH::EConstraintSpace::LocalToBodyCOM;
			/* Position in parent local space (bone's initial pos). */
			constraint->mPosition1 = JPH::RVec3(
				JoltPhysics::convert(
					bone->getInitialPosition()));
			constraint->mPosition2 = JPH::RVec3::sZero();
			constraint->mTwistAxis2 = JPH::Vec3::sAxisY();
			JPH::Quat childLocalRot = JoltPhysics::convert(
				bone->getInitialOrientation());
			constraint->mTwistAxis1 = childLocalRot *
						  constraint->mTwistAxis2;
			JPH::Vec3 planeAxis1 = childLocalRot *
					       JPH::Vec3::sAxisX();
			constraint->mPlaneAxis1 = planeAxis1;
			constraint->mPlaneAxis2 = JPH::Vec3::sAxisX();
			constraint->mTwistMinAngle = -JPH::DegreesToRadians(
				45.0f);
			constraint->mTwistMaxAngle = JPH::DegreesToRadians(
				45.0f);
			constraint->mNormalHalfConeAngle =
				JPH::DegreesToRadians(30.0f);
			constraint->mPlaneHalfConeAngle =
				JPH::DegreesToRadians(30.0f);
			part.mToParent = constraint;
		}
	}

	/* Collision group setup.
	 * Subgroup 0 = character capsule, 1 = head sphere, 2 = chest sphere,
	 * 3+ = hair joints. */
	uint32_t collisionGroupId;
	constexpr uint32_t subgroupHairStart = 3;
	JPH::GroupFilterTable *groupFilter = nullptr;
	if (e.has<CharacterComponent>()) {
		auto &cc = e.get_mut<CharacterComponent>();
		if (cc.collisionGroupId == 0)
			cc.collisionGroupId =
				static_cast<uint32_t>(e.id()) |
				0x80000000;
		collisionGroupId = cc.collisionGroupId;
		groupFilter = m_physics->getOrCreateGroupFilter(
			collisionGroupId);
		/* Disable body (subgroup 0), head (1) and chest (2) vs every
		 * hair joint. These kinematic spheres sit at the attachment
		 * points; overlap with the hair chain creates an explosive
		 * feedback loop that throws the joints across the map. */
		for (size_t i = 0; i < orderedBones.size(); ++i) {
			uint32_t hairSub = subgroupHairStart + (uint32_t)i;
			groupFilter->DisableCollision(0, hairSub);
			groupFilter->DisableCollision(1, hairSub);
			groupFilter->DisableCollision(2, hairSub);
		}
		/* Disable all hair-joint vs hair-joint collisions. Siblings
		 * and parents/children start close together and can otherwise
		 * generate explosive contact impulses. */
		for (size_t i = 0; i < orderedBones.size(); ++i) {
			for (size_t j = i + 1; j < orderedBones.size(); ++j) {
				groupFilter->DisableCollision(
					subgroupHairStart + (uint32_t)i,
					subgroupHairStart + (uint32_t)j);
			}
		}
	} else {
		collisionGroupId = s_nextCollisionGroup++;
	}

	for (size_t i = 0; i < orderedBones.size(); ++i) {
		settings->mParts[i].mObjectLayer = Layers::HAIR;
		if (groupFilter) {
			settings->mParts[i].mCollisionGroup.SetGroupFilter(
				groupFilter);
			settings->mParts[i].mCollisionGroup.SetGroupID(
				collisionGroupId);
			settings->mParts[i].mCollisionGroup.SetSubGroupID(
				subgroupHairStart + (uint32_t)i);
		}
	}
	if (!e.has<CharacterComponent>())
		settings->DisableParentChildCollisions();

	settings->Stabilize();
	settings->CalculateConstraintPriorities();
	settings->CalculateBodyIndexToConstraintIndex();

	JPH::PhysicsSystem *physSystem = m_physics->getPhysicsSystem();
	JPH::Ref<JPH::Ragdoll> ragdoll = settings->CreateRagdoll(
		collisionGroupId, 0, physSystem);

	/* Store skeleton bones for later pose sync. */
	HairRagdollState state;
	state.ragdoll = ragdoll;
	state.entityId = e.id();
	state.hairEntity = hairEnt;
	state.masterEntity = masterEnt;
	state.slotName = slot;
	{
		Ogre::Bone *rootBone = orderedBones[0];
		Ogre::Matrix4 rootLocal = Ogre::Matrix4::IDENTITY;
		rootLocal.makeTransform(rootBone->getInitialPosition(),
					rootBone->getInitialScale(),
					rootBone->getInitialOrientation());
		state.rootBindTransform = rootLocal;
	}
	state.boneNames.reserve(orderedBones.size());
	for (Ogre::Bone *bone : orderedBones)
		state.boneNames.push_back(bone->getName());
	state.parentIndices = parentIndices;

	/* Ensure all bones are manually controlled. */
	for (Ogre::Bone *bone : orderedBones) {
		bone->setManuallyControlled(true);
	}

	ragdoll->AddToPhysicsSystem(JPH::EActivation::Activate);

	Ogre::LogManager::getSingleton().logMessage(
		"HairPhysicsSystem: created ragdoll for " + slot +
		" ('" + hairEnt->getMesh()->getName() + "', " +
		std::to_string(orderedBones.size()) + " bones, " +
		std::to_string(ragdoll->GetConstraintCount()) +
		" constraints)");

	m_states[e.id()][slot] = std::move(state);
}

void HairPhysicsSystem::destroyRagdoll(HairRagdollState &state)
{
	if (state.ragdoll) {
		state.ragdoll->RemoveFromPhysicsSystem();
		state.ragdoll = nullptr;
	}
	state.boneNames.clear();
	state.parentIndices.clear();
}

bool HairPhysicsSystem::isStateValid(const HairRagdollState &state) const
{
	if (!state.hairEntity || state.entityId == 0 || !m_slotSystem)
		return false;

	flecs::entity e = m_world.entity(state.entityId);
	if (!e.is_alive() || !e.has<CharacterIdentityComponent>())
		return false;

	Ogre::Entity *current = m_slotSystem->getSlotEntity(e, state.slotName);
	return current == state.hairEntity && current->hasSkeleton();
}

void HairPhysicsSystem::syncRootToHead(HairRagdollState &state)
{
	if (!state.ragdoll || !isStateValid(state))
		return;

	flecs::entity e = m_world.entity(state.entityId);
	Ogre::Entity *masterEnt = m_slotSystem->getMasterEntity(e);
	if (!masterEnt)
		return;

	Ogre::Bone *attachBone = findAttachBone(state.hairEntity, masterEnt);
	if (!attachBone)
		return;

	/* The root body was created at attachBone * rootBindTransform, so we
	 * must keep it there each frame. */
	Ogre::Matrix4 attachBoneMat = getAttachBoneWorldMatrix(masterEnt,
							       attachBone);
	Ogre::Matrix4 rootWorldMat = attachBoneMat * state.rootBindTransform;
	Ogre::Vector3 pos = rootWorldMat.getTrans();
	Ogre::Quaternion rot(rootWorldMat.linear());
	rot.normalise();

	JPH::BodyID rootBody = state.ragdoll->GetBodyID(0);
	m_physics->setPositionAndRotation(rootBody, pos, rot, true);
}

void HairPhysicsSystem::syncPhysicsToSkeleton(HairRagdollState &state)
{
	if (!state.ragdoll || state.boneNames.empty() ||
	    !isStateValid(state))
		return;

	Ogre::SkeletonInstance *skel = state.hairEntity->getSkeleton();
	if (!skel)
		return;

	JPH::SkeletonPose pose;
	pose.SetSkeleton(
		state.ragdoll->GetRagdollSettings()->GetSkeleton());
	state.ragdoll->GetPose(pose);

	const JPH::SkeletonPose::Mat44Vector &jointMatrices =
		pose.GetJointMatrices();
	JPH::RVec3 rootOffset = pose.GetRootOffset();
	Ogre::Matrix4 rootOffsetMat = Ogre::Matrix4::IDENTITY;
	rootOffsetMat.makeTransform(JoltPhysics::convert(rootOffset),
				    Ogre::Vector3::UNIT_SCALE,
				    Ogre::Quaternion::IDENTITY);

	/* Get the attachment transform so we can compute hair-entity-local
	 * transforms for the root bone. */
	flecs::entity e = m_world.entity(state.entityId);
	Ogre::Entity *masterEnt = m_slotSystem->getMasterEntity(e);

	Ogre::Bone *attachBone = findAttachBone(state.hairEntity, masterEnt);
	Ogre::Matrix4 attachBoneInv = Ogre::Matrix4::IDENTITY;
	if (attachBone) {
		Ogre::Matrix4 attachBoneMat = getAttachBoneWorldMatrix(
			masterEnt, attachBone);
		attachBoneInv = attachBoneMat.inverse();
	}

	std::vector<Ogre::Matrix4> worldMats(state.boneNames.size());
	for (size_t i = 0; i < state.boneNames.size(); ++i)
		worldMats[i] = rootOffsetMat * joltMatrixToOgre(jointMatrices[i]);

	for (size_t i = 0; i < state.boneNames.size(); ++i) {
		Ogre::Bone *bone = skel->getBone(state.boneNames[i]);
		if (!bone)
			continue;

		Ogre::Matrix4 localMat;
		if (state.parentIndices[i] >= 0) {
			/* Joint transforms from Jolt are in skeleton world space;
			 * Ogre bones expect parent-local space. */
			localMat = worldMats[state.parentIndices[i]].inverse() *
				   worldMats[i];
		} else {
			/* Root bone is in hair-entity local space. */
			localMat = attachBoneInv * worldMats[i];
		}

		Ogre::Vector3 pos, scl(Ogre::Vector3::UNIT_SCALE);
		Ogre::Quaternion rot;
		pos = localMat.getTrans();
		rot = Ogre::Quaternion(localMat.linear());
		rot.normalise();
		/* Extract scale from basis vectors. */
		scl.x = Ogre::Vector3(localMat[0][0], localMat[0][1],
				      localMat[0][2]).length();
		scl.y = Ogre::Vector3(localMat[1][0], localMat[1][1],
				      localMat[1][2]).length();
		scl.z = Ogre::Vector3(localMat[2][0], localMat[2][1],
				      localMat[2][2]).length();

		bone->setManuallyControlled(true);
		bone->setPosition(pos);
		bone->setOrientation(rot);
		bone->setScale(scl);
	}
}

JPH::Mat44 HairPhysicsSystem::ogreMatrixToJolt(
	const Ogre::Matrix4 &m) const
{
	JPH::Mat44 jm;
	for (int row = 0; row < 4; ++row) {
		for (int col = 0; col < 4; ++col) {
			jm(row, col) = m[row][col];
		}
	}
	return jm;
}

Ogre::Matrix4 HairPhysicsSystem::joltMatrixToOgre(
	const JPH::Mat44 &m) const
{
	Ogre::Matrix4 om;
	for (int row = 0; row < 4; ++row) {
		for (int col = 0; col < 4; ++col) {
			om[row][col] = m(row, col);
		}
	}
	return om;
}
