#include "CharacterSlotSystem.hpp"
#include "../components/Transform.hpp"
#include "../components/AnimationTree.hpp"
#include <OgreAnimation.h>
#include <OgreAnimationState.h>
#include <OgreAnimationTrack.h>
#include <OgreDataStream.h>
#include <OgreEntity.h>
#include <OgreKeyFrame.h>
#include <OgreLogManager.h>
#include <OgreMesh.h>
#include <OgreResourceGroupManager.h>
#include <OgreSceneNode.h>
#include <algorithm>

bool CharacterSlotSystem::s_catalogLoaded = false;
nlohmann::json CharacterSlotSystem::s_bodyParts = nlohmann::json::object();
std::set<Ogre::String> CharacterSlotSystem::s_meshNames;

CharacterSlotSystem::CharacterSlotSystem(flecs::world &world,
						 Ogre::SceneManager *sceneMgr)
	: m_world(world)
	, m_sceneMgr(sceneMgr)
{
	m_world.observer<CharacterSlotsComponent>("CharacterSlotsCleanup")
		.event(flecs::OnRemove)
		.each([this](flecs::entity e, CharacterSlotsComponent &) {
			destroyCharacterParts(e);
		});
}

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

void CharacterSlotSystem::initialize()
{
	if (m_initialized)
		return;
	loadCatalog();
	m_initialized = true;
}

void CharacterSlotSystem::loadCatalog()
{
	if (s_catalogLoaded)
		return;

	s_bodyParts = nlohmann::json::object();
	s_meshNames.clear();

	Ogre::ResourceGroupManager &rgm =
		Ogre::ResourceGroupManager::getSingleton();
	Ogre::StringVector groups = rgm.getResourceGroups();
	std::vector<Ogre::String> partNames;

	for (const auto &group : groups) {
		Ogre::StringVectorPtr names =
			rgm.findResourceNames(group, "body_part_*.json");
		if (names)
			partNames.insert(partNames.end(), names->begin(),
					 names->end());
	}


	for (const auto &name : partNames) {
		Ogre::String group = rgm.findGroupContainingResource(name);
		Ogre::DataStreamPtr stream = rgm.openResource(name, group);
		Ogre::String jsonStr = stream->getAsString();
		try {
			nlohmann::json jdata = nlohmann::json::parse(jsonStr);
			if (!jdata.contains("age") || !jdata.contains("sex") ||
			    !jdata.contains("slot") || !jdata.contains("mesh"))
				continue;

			Ogre::String age = jdata["age"].get<Ogre::String>();
			Ogre::String sex = jdata["sex"].get<Ogre::String>();
			Ogre::String slot = jdata["slot"].get<Ogre::String>();
			Ogre::String mesh = jdata["mesh"].get<Ogre::String>();

			if (!s_bodyParts.contains(age))
				s_bodyParts[age] = nlohmann::json::object();
			if (!s_bodyParts[age].contains(sex))
				s_bodyParts[age][sex] = nlohmann::json::object();
			if (!s_bodyParts[age][sex].contains(slot))
				s_bodyParts[age][sex][slot] =
					nlohmann::json::array();

			s_bodyParts[age][sex][slot].push_back(jdata);
			s_meshNames.insert(mesh);

			/* Preload mesh into Characters group */
			try {
				Ogre::MeshManager::getSingleton().load(mesh,
								       "Characters");
			} catch (...) {
			}
		} catch (...) {
			continue;
		}
	}


	s_catalogLoaded = true;
}

bool CharacterSlotSystem::isCatalogLoaded()
{
	return s_catalogLoaded;
}

std::vector<Ogre::String> CharacterSlotSystem::getAges()
{
	std::vector<Ogre::String> ages;
	if (!s_catalogLoaded)
		return ages;
	for (auto &el : s_bodyParts.items())
		ages.push_back(el.key());
	return ages;
}

std::vector<Ogre::String> CharacterSlotSystem::getSexes(
	const Ogre::String &age)
{
	std::vector<Ogre::String> sexes;
	if (!s_catalogLoaded || !s_bodyParts.contains(age))
		return sexes;
	for (auto &el : s_bodyParts[age].items())
		sexes.push_back(el.key());
	return sexes;
}

std::vector<Ogre::String> CharacterSlotSystem::getSlots(
	const Ogre::String &age, const Ogre::String &sex)
{
	std::vector<Ogre::String> slots;
	if (!s_catalogLoaded || !s_bodyParts.contains(age) ||
	    !s_bodyParts[age].contains(sex))
		return slots;
	for (auto &el : s_bodyParts[age][sex].items())
		slots.push_back(el.key());
	return slots;
}

std::vector<Ogre::String> CharacterSlotSystem::getMeshesForLayer(
	const Ogre::String &age, const Ogre::String &sex,
	const Ogre::String &slot, int layer)
{
	std::vector<Ogre::String> meshes;
	if (!s_catalogLoaded || !s_bodyParts.contains(age) ||
	    !s_bodyParts[age].contains(sex) ||
	    !s_bodyParts[age][sex].contains(slot))
		return meshes;

	for (const auto &entry : s_bodyParts[age][sex][slot]) {
		int entryLayer = entry.value("layer", 0);
		if (entryLayer == layer)
			meshes.push_back(entry.value("mesh", ""));
	}
	return meshes;
}

Ogre::String CharacterSlotSystem::getMeshLabel(
	const Ogre::String &age, const Ogre::String &sex,
	const Ogre::String &slot, const Ogre::String &mesh)
{
	if (!s_catalogLoaded || !s_bodyParts.contains(age) ||
	    !s_bodyParts[age].contains(sex) ||
	    !s_bodyParts[age][sex].contains(slot))
		return mesh;

	for (const auto &entry : s_bodyParts[age][sex][slot]) {
		if (entry.value("mesh", "") == mesh) {
			const auto &garments =
				entry.value("garments", nlohmann::json::array());
			if (garments.empty())
				return "nude";
			Ogre::String label;
			for (size_t i = 0; i < garments.size(); ++i) {
				if (i > 0)
					label += " + ";
				label += garments[i].get<std::string>();
			}
			return label;
		}
	}
	return mesh;
}

std::vector<Ogre::String> CharacterSlotSystem::getMeshes(
	const Ogre::String &age, const Ogre::String &sex,
	const Ogre::String &slot)
{
	std::vector<Ogre::String> meshes;
	if (!s_catalogLoaded || !s_bodyParts.contains(age) ||
	    !s_bodyParts[age].contains(sex) ||
	    !s_bodyParts[age][sex].contains(slot))
		return meshes;
	for (const auto &entry : s_bodyParts[age][sex][slot])
		meshes.push_back(entry["mesh"].get<Ogre::String>());
	return meshes;
}

std::vector<Ogre::String> CharacterSlotSystem::getShapeKeyNames(
	const Ogre::String &age, const Ogre::String &sex)
{
	std::set<Ogre::String> keySet;
	if (!s_catalogLoaded || !s_bodyParts.contains(age) ||
	    !s_bodyParts[age].contains(sex))
		return {};

	for (auto &slotEl : s_bodyParts[age][sex].items()) {
		for (const auto &entry : slotEl.value()) {
			const auto &keys =
				entry.value("shape_keys", nlohmann::json::array());
			for (const auto &k : keys)
				keySet.insert(k.get<Ogre::String>());
		}
	}

	return std::vector<Ogre::String>(keySet.begin(), keySet.end());
}

Ogre::String CharacterSlotSystem::resolveMesh(
	const Ogre::String &age, const Ogre::String &sex,
	const Ogre::String &slot, const SlotSelection &sel,
	int outfitLevel)
{
	if (!sel.explicitMesh.empty())
		return sel.explicitMesh;

	if (!s_catalogLoaded || !s_bodyParts.contains(age) ||
	    !s_bodyParts[age].contains(sex) ||
	    !s_bodyParts[age][sex].contains(slot))
		return "";

	const auto &slotEntries = s_bodyParts[age][sex][slot];

	/* outfitLevel: 0=nude, 1=lingerie, 2=clothed */
	if (outfitLevel >= 2 && sel.layer2Mesh != "none" &&
	    !sel.layer2Mesh.empty()) {
		for (const auto &entry : slotEntries) {
			if (entry.value("mesh", "") == sel.layer2Mesh)
				return sel.layer2Mesh;
		}
	}
	if (outfitLevel >= 1 && sel.layer1Mesh != "none" &&
	    !sel.layer1Mesh.empty()) {
		for (const auto &entry : slotEntries) {
			if (entry.value("mesh", "") == sel.layer1Mesh)
				return sel.layer1Mesh;
		}
	}

	/* Fallback to layer 0 (nude base) — prefer canonical base mesh */
	Ogre::String canonicalBase = sex + "_" + slot + ".mesh";
	Ogre::String firstLayer0;
	for (const auto &entry : slotEntries) {
		if (entry.value("layer", 0) == 0) {
			Ogre::String mesh = entry["mesh"].get<Ogre::String>();
			if (mesh == canonicalBase)
				return mesh;
			if (firstLayer0.empty())
				firstLayer0 = mesh;
		}
	}
	if (!firstLayer0.empty())
		return firstLayer0;

	/* Last resort: first available entry */
	if (!slotEntries.empty())
		return slotEntries[0]["mesh"].get<Ogre::String>();

	return "";
}

void CharacterSlotSystem::update()
{
	if (!m_initialized)
		return;

	int total = 0;
	int dirtyCount = 0;
	m_world.query<CharacterSlotsComponent>().each(
		[&](flecs::entity e, CharacterSlotsComponent &cs) {
			total++;
			if (cs.dirty) {
				dirtyCount++;
				buildCharacter(e, cs);
				cs.dirty = false;
			}
		});
}

static const nlohmann::json *findCatalogEntry(
	const Ogre::String &age, const Ogre::String &sex,
	const Ogre::String &slot, const Ogre::String &mesh)
{
	if (!CharacterSlotSystem::isCatalogLoaded())
		return nullptr;
	const nlohmann::json &cat = CharacterSlotSystem::getCatalog();
	if (!cat.contains(age) || !cat[age].contains(sex) ||
	    !cat[age][sex].contains(slot))
		return nullptr;
	for (const auto &entry : cat[age][sex][slot]) {
		if (entry.value("mesh", "") == mesh)
			return &entry;
	}
	return nullptr;
}

static void ensureMeshPoseAnimation(const Ogre::String &meshName)
{
	Ogre::MeshPtr mesh;
	try {
		mesh = Ogre::MeshManager::getSingleton().load(meshName,
								      "Characters");
	} catch (...) {
		return;
	}
	if (!mesh || mesh->getPoseCount() == 0)
		return;
	try {
		mesh->getAnimation("ShapeKeys");
	} catch (...) {
		Ogre::Animation *anim = mesh->createAnimation("ShapeKeys", 1.0f);
		Ogre::VertexAnimationTrack *track = anim->createVertexTrack(
			0, Ogre::VAT_POSE);
		Ogre::VertexPoseKeyFrame *kf =
			track->createVertexPoseKeyFrame(0.0f);
		for (size_t i = 0; i < mesh->getPoseCount(); ++i)
			kf->addPoseReference(static_cast<ushort>(i), 0.0f);
	}
}

void CharacterSlotSystem::buildCharacter(flecs::entity e,
						 CharacterSlotsComponent &cs)
{
	destroyCharacterParts(e);

	/* Migrate old slots map to slotSelections if needed */
	if (cs.slotSelections.empty() && !cs.slots.empty()) {
		for (const auto &pair : cs.slots) {
			SlotSelection sel;
			sel.explicitMesh = pair.second;
			cs.slotSelections[pair.first] = sel;
		}
	}

	/* Populate default slots from catalog if still empty */
	if (cs.slotSelections.empty()) {
		auto slots = getSlots(cs.age, cs.sex);
		for (const auto &slot : slots) {
			SlotSelection sel;
			cs.slotSelections[slot] = sel;
		}
	}

	if (!e.has<TransformComponent>())
		return;

	auto &transform = e.get_mut<TransformComponent>();
	if (!transform.node)
		return;

	/* Determine master slot (face preferred, else first non-empty) */
	Ogre::String masterSlot;
	if (cs.slotSelections.find("face") != cs.slotSelections.end()) {
		Ogre::String mesh = resolveMesh(cs.age, cs.sex, "face",
						cs.slotSelections["face"],
						cs.outfitLevel);
		if (!mesh.empty())
			masterSlot = "face";
	}
	if (masterSlot.empty()) {
		for (const auto &pair : cs.slotSelections) {
			Ogre::String mesh = resolveMesh(cs.age, cs.sex, pair.first,
							pair.second,
							cs.outfitLevel);
			if (!mesh.empty()) {
				masterSlot = pair.first;
				break;
			}
		}
	}

	if (masterSlot.empty()) {
		return;
	}

	Ogre::String masterMesh = resolveMesh(cs.age, cs.sex, masterSlot,
					      cs.slotSelections[masterSlot],
					      cs.outfitLevel);

	if (masterMesh.empty())
		return;

	/* Pre-create pose animation on mesh so entity knows about it */
	ensureMeshPoseAnimation(masterMesh);

	Ogre::Entity *masterEnt = nullptr;
	try {
		Ogre::MeshPtr meshPtr = Ogre::MeshManager::getSingleton().load(
			masterMesh, "Characters");
		masterEnt = m_sceneMgr->createEntity(meshPtr);
		transform.node->attachObject(masterEnt);
		m_entities[e.id()].parts[masterSlot] = masterEnt;
		cs.masterEntity = masterEnt;

		/* Setup pose animation for shape keys */
		const nlohmann::json *entry = findCatalogEntry(
			cs.age, cs.sex, masterSlot, masterMesh);
		applyShapeKeys(e, masterEnt, entry);

		/* Notify AnimationTreeSystem that entity changed */
		if (e.has<AnimationTreeComponent>())
			e.get_mut<AnimationTreeComponent>().dirty = true;
	} catch (const Ogre::Exception &ex) {
		Ogre::LogManager::getSingleton().logMessage(
			"CharacterSlotSystem: Failed to load master mesh '" +
			masterMesh + "': " + ex.getDescription());
		return;
	}

	for (const auto &pair : cs.slotSelections) {
		const Ogre::String &slot = pair.first;
		const SlotSelection &sel = pair.second;

		if (slot == masterSlot)
			continue;

		Ogre::String mesh = resolveMesh(cs.age, cs.sex, slot, sel,
						cs.outfitLevel);
		if (mesh.empty())
			continue;

		try {
			ensureMeshPoseAnimation(mesh);
			Ogre::MeshPtr partMesh =
				Ogre::MeshManager::getSingleton().load(
					mesh, "Characters");
			Ogre::Entity *partEnt = m_sceneMgr->createEntity(partMesh);
			partEnt->shareSkeletonInstanceWith(masterEnt);
			transform.node->attachObject(partEnt);
			m_entities[e.id()].parts[slot] = partEnt;
			const nlohmann::json *entry = findCatalogEntry(
				cs.age, cs.sex, slot, mesh);
			applyShapeKeys(e, partEnt, entry);
		} catch (const Ogre::Exception &ex) {
			Ogre::LogManager::getSingleton().logMessage(
				"CharacterSlotSystem: Failed to load part '" +
				slot + "' mesh '" + mesh +
				"': " + ex.getDescription());
		}
	}
}

void CharacterSlotSystem::applyShapeKeys(flecs::entity e,
						 Ogre::Entity *ent,
						 const nlohmann::json *entry)
{
	if (!ent || !entry)
		return;

	Ogre::MeshPtr mesh = ent->getMesh();
	if (!mesh || mesh->getPoseCount() == 0)
		return;

	/* Create a pose animation track if one doesn't exist */
	Ogre::Animation *anim = nullptr;
	try {
		anim = mesh->getAnimation("ShapeKeys");
	} catch (...) {
		anim = mesh->createAnimation("ShapeKeys", 1.0f);
		Ogre::VertexAnimationTrack *track = anim->createVertexTrack(
			0, Ogre::VAT_POSE);
		Ogre::VertexPoseKeyFrame *kf =
			track->createVertexPoseKeyFrame(0.0f);
		for (size_t i = 0; i < mesh->getPoseCount(); ++i)
			kf->addPoseReference(static_cast<ushort>(i), 0.0f);
	}

	if (!ent->hasAnimationState("ShapeKeys"))
		return;
	Ogre::AnimationState *as = ent->getAnimationState("ShapeKeys");
	as->setEnabled(true);
	as->setLoop(false);

	/* Build name -> pose index map from catalog */
	const auto &shapeKeys =
		entry->value("shape_keys", nlohmann::json::array());
	std::unordered_map<Ogre::String, size_t> nameToIndex;
	for (size_t i = 0; i < shapeKeys.size(); ++i)
		nameToIndex[shapeKeys[i].get<Ogre::String>()] = i;

	/* Apply weights from CharacterShapeKeysComponent */
	if (e.has<CharacterShapeKeysComponent>()) {
		auto &skc = e.get_mut<CharacterShapeKeysComponent>();
		for (const auto &pair : skc.weights) {
			auto it = nameToIndex.find(pair.first);
			if (it == nameToIndex.end())
				continue;
			if (it->second >= mesh->getPoseCount())
				continue;
			/* Update the keyframe's pose reference influence */
			Ogre::VertexAnimationTrack *track =
				anim->getVertexTrack(0);
			if (track) {
				Ogre::VertexPoseKeyFrame *kf =
					track->getVertexPoseKeyFrame(0);
				if (kf)
					kf->updatePoseReference(
						static_cast<ushort>(it->second),
						pair.second);
			}
		}
	}

	/* Force OGRE to update the entity with new pose weights */
	as->setTimePosition(0.0f);
}

void CharacterSlotSystem::destroyCharacterParts(flecs::entity e)
{
	auto it = m_entities.find(e.id());
	if (it == m_entities.end())
		return;

	for (auto &pair : it->second.parts) {
		Ogre::Entity *ent = pair.second;
		if (ent) {
			try {
				ent->detachFromParent();
				m_sceneMgr->destroyEntity(ent);
			} catch (...) {
			}
		}
	}

	it->second.parts.clear();
	m_entities.erase(it);
}

Ogre::Entity *CharacterSlotSystem::getSlotEntity(flecs::entity e,
						  const Ogre::String &slot)
{
	auto it = m_entities.find(e.id());
	if (it == m_entities.end())
		return nullptr;
	auto pit = it->second.parts.find(slot);
	if (pit == it->second.parts.end())
		return nullptr;
	return pit->second;
}

void CharacterSlotSystem::setSlotVisible(flecs::entity e,
						 const Ogre::String &slot,
						 bool visible)
{
	Ogre::Entity *ent = getSlotEntity(e, slot);
	if (ent) {
		ent->setVisible(visible);
	}
}
