#include "ResourceManager.h"

#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <vector>

#include <assimp/Importer.hpp>
#include <assimp/postprocess.h>
#include <assimp/scene.h>

#include "texture.h"
#include "font.h"
#include "MeshData.h"

ResourceManager::~ResourceManager()
{
	for (auto& [path, texture] : textures) {
		if (!texture) continue;
		texture->Delete();
		delete texture;
	}
	for (Texture* texture : createdTextures) {
		texture->Delete();
		delete texture;
	}
	for (auto& [path, font] : fonts)
		delete font;
	for (auto& [path, mesh] : meshes)
		delete mesh;
}
// Which of a mesh's uv maps its material's base color texture is laid out on, 0 for the first
static unsigned int BaseColorUVChannel(const aiMaterial* material)
{
	int channel = 0;
	if (material->Get(AI_MATKEY_UVWSRC(aiTextureType_BASE_COLOR, 0), channel) != AI_SUCCESS)
		material->Get(AI_MATKEY_UVWSRC(aiTextureType_DIFFUSE, 0), channel);
	return channel > 0 ? (unsigned int)channel : 0u;
}

// One material's base color texture, embedded in the file like a .glb does or in a file next to it.
// nullptr if it has none or it can't be read.
static Texture* LoadMaterialTexture(const aiScene* scene, const aiMaterial* material, const std::string& modelPath)
{
	aiString file;
	if (material->GetTexture(aiTextureType_BASE_COLOR, 0, &file) != AI_SUCCESS &&
		material->GetTexture(aiTextureType_DIFFUSE, 0, &file) != AI_SUCCESS)
		return nullptr;

	try {
		if (const aiTexture* embedded = scene->GetEmbeddedTexture(file.C_Str())) {
			// A height of 0 means it's still encoded, a whole png or jpg mWidth bytes long
			if (embedded->mHeight != 0) {
				std::cout << "RESOURCE MANAGER: " << modelPath << " has an unencoded embedded texture, those aren't supported\n";
				return nullptr;
			}
			return new Texture((const unsigned char*)embedded->pcData, (int)embedded->mWidth, modelPath + "#" + file.C_Str());
		}
		const size_t slash = modelPath.find_last_of("/\\");
		const std::string folder = slash == std::string::npos ? "" : modelPath.substr(0, slash + 1);
		return new Texture((folder + file.C_Str()).c_str());
	}
	catch (const std::exception& error) {
		std::cout << "RESOURCE MANAGER: " << error.what() << "\n";
		return nullptr;
	}
}

// A base color texture for the whole model, embedded in the file like a .glb does or in a file next to it.
// Meshes only take one texture, so a model with several gets the one covering the most of it.
static Texture* LoadModelTexture(const aiScene* scene, const std::string& modelPath)
{
	std::vector<unsigned int> faces(scene->mNumMaterials, 0);
	for (unsigned int m = 0; m < scene->mNumMeshes; m++)
		if (scene->mMeshes[m]->mMaterialIndex < scene->mNumMaterials)
			faces[scene->mMeshes[m]->mMaterialIndex] += scene->mMeshes[m]->mNumFaces;
	std::vector<unsigned int> order(scene->mNumMaterials);
	for (unsigned int m = 0; m < scene->mNumMaterials; m++)
		order[m] = m;
	std::stable_sort(order.begin(), order.end(), [&](unsigned int a, unsigned int b) { return faces[a] > faces[b]; });

	for (unsigned int m : order) {
		aiString file;
		const aiMaterial* material = scene->mMaterials[m];
		if (material->GetTexture(aiTextureType_BASE_COLOR, 0, &file) != AI_SUCCESS &&
			material->GetTexture(aiTextureType_DIFFUSE, 0, &file) != AI_SUCCESS)
			continue;
		return LoadMaterialTexture(scene, material, modelPath);
	}
	return nullptr;
}

// Adds one of assimp's meshes onto the end of data, with the uvs its material's texture uses
static void AppendMesh(const aiScene& scene, const aiMesh& mesh, MeshData& data)
{
	unsigned int uv = mesh.mMaterialIndex < scene.mNumMaterials ? BaseColorUVChannel(scene.mMaterials[mesh.mMaterialIndex]) : 0;
	if (uv >= AI_MAX_NUMBER_OF_TEXTURECOORDS || !mesh.mTextureCoords[uv])
		uv = 0;
	const GLuint firstVertex = (GLuint)data.vertices.size();
	for (unsigned int i = 0; i < mesh.mNumVertices; i++) {
		Vertex vertex;
		vertex.Position = glm::vec3(mesh.mVertices[i].x, mesh.mVertices[i].y, mesh.mVertices[i].z);
		vertex.Normal = mesh.HasNormals()
			? glm::vec3(mesh.mNormals[i].x, mesh.mNormals[i].y, mesh.mNormals[i].z) : glm::vec3(0.0f, 1.0f, 0.0f);
		vertex.TexCoords = mesh.mTextureCoords[uv]
			? glm::vec2(mesh.mTextureCoords[uv][i].x, mesh.mTextureCoords[uv][i].y) : glm::vec2(0.0f);
		data.vertices.push_back(vertex);
	}
	for (unsigned int f = 0; f < mesh.mNumFaces; f++) {
		for (unsigned int j = 0; j < mesh.mFaces[f].mNumIndices; j++)
			data.indices.push_back(firstVertex + mesh.mFaces[f].mIndices[j]);
	}
}

const MeshData* ResourceManager::GetMeshData(const std::string& path, const std::string& name)
{
	const std::string key = path + "#" + name;
	auto existing = meshes.find(key);
	if (existing != meshes.end())
		return existing->second;

	// Not pre-transformed, so the vertices stay around the object's own origin
	Assimp::Importer importer;
	const aiScene* scene = importer.ReadFile(path, aiProcess_Triangulate | aiProcess_GenSmoothNormals |
		aiProcess_JoinIdenticalVertices);

	MeshData* data = nullptr;
	const aiNode* node = scene && scene->mRootNode ? scene->mRootNode->FindNode(name.c_str()) : nullptr;
	if (!node || node->mNumMeshes == 0) {
		std::cout << "RESOURCE MANAGER: no object called " << name << " with a mesh in " << path << "\n";
	}
	else {
		data = new MeshData();
		for (unsigned int m = 0; m < node->mNumMeshes; m++)
			AppendMesh(*scene, *scene->mMeshes[node->mMeshes[m]], *data);
	}

	meshes[key] = data;
	return data;
}

Texture* ResourceManager::GetModelTexture(const std::string& path)
{
	GetMeshData(path);
	auto found = textures.find(path + "#model");
	return found != textures.end() ? found->second : nullptr;
}

const MeshData* ResourceManager::GetMeshData(const std::string& path)
{
	auto existing = meshes.find(path);
	if (existing != meshes.end())
		return existing->second;

	// Pre-transforming bakes each node's placement in, so the merged meshes keep their layout
	Assimp::Importer importer;
	const aiScene* scene = importer.ReadFile(path, aiProcess_Triangulate | aiProcess_GenSmoothNormals |
		aiProcess_JoinIdenticalVertices | aiProcess_PreTransformVertices);

	MeshData* data = nullptr;
	if (!scene || (scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE) || scene->mNumMeshes == 0) {
		std::cout << "RESOURCE MANAGER: couldn't load mesh " << path << ": " << importer.GetErrorString() << "\n";
	}
	else {
		data = new MeshData();
		for (unsigned int m = 0; m < scene->mNumMeshes; m++)
			AppendMesh(*scene, *scene->mMeshes[m], *data);
	}

	// Its texture comes out now too, while the file's open. Kept with the rest, under a name no file has.
	textures[path + "#model"] = scene ? LoadModelTexture(scene, path) : nullptr;

	meshes[path] = data;
	return data;
}

bool ResourceManager::GetModelMarker(const std::string& path, const std::string& name, glm::mat4& transform)
{
	const std::string key = path + "#" + name;
	auto existing = markers.find(key);
	if (existing == markers.end()) {
		// Read without flattening, which would throw away objects with no mesh like empties
		Assimp::Importer importer;
		const aiScene* scene = importer.ReadFile(path, 0);
		const aiNode* node = scene && scene->mRootNode ? scene->mRootNode->FindNode(name.c_str()) : nullptr;

		// Its placement and every parent's, all the way up, the same as flattening gives the meshes
		aiMatrix4x4 global;
		for (const aiNode* at = node; at; at = at->mParent)
			global = at->mTransformation * global;
		glm::mat4 found(1.0f);
		for (int row = 0; row < 4; row++)
			for (int column = 0; column < 4; column++)
				found[column][row] = global[row][column];   // assimp is row major, glm column major

		if (!node)
			std::cout << "RESOURCE MANAGER: no object called " << name << " in " << path << "\n";
		existing = markers.emplace(key, std::make_pair(node != nullptr, found)).first;
	}
	if (existing->second.first)
		transform = existing->second.second;
	return existing->second.first;
}

Texture* ResourceManager::CreateTexture(int width, int height)
{
	createdTextures.push_back(new Texture(width, height));
	return createdTextures.back();
}

Texture* ResourceManager::CopyTexture(const Texture& source)
{
	createdTextures.push_back(new Texture(source));
	return createdTextures.back();
}

// Failed loads are cached as nullptr too, so a missing file only gets reported once
Texture* ResourceManager::GetTexture(const std::string& path)
{
	auto existing = textures.find(path);
	if (existing != textures.end())
		return existing->second;

	Texture* texture = nullptr;
	try {
		texture = new Texture(path.c_str());
	}
	catch (const std::exception& error) {
		std::cout << "RESOURCE MANAGER: " << error.what() << "\n";
	}
	textures[path] = texture;
	return texture;
}

Font* ResourceManager::GetFont(const std::string& path)
{
	auto existing = fonts.find(path);
	if (existing != fonts.end())
		return existing->second;

	Font* font = nullptr;
	try {
		font = new Font(path.c_str());
	}
	catch (const std::exception& error) {
		std::cout << "RESOURCE MANAGER: " << error.what() << ": " << path << "\n";
	}
	fonts[path] = font;
	return font;
}