/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#include "MeshLoader.hpp"
#include "Scene/Geometry.hpp"
#include "Base/Logger.hpp"
#if __has_include("assimp/Importer.hpp")
#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/mesh.h>
#include <assimp/postprocess.h>
#define ASSIMP_AVAILABLE
#endif

#include "tinyusdz.hh"
#include "tydra/attribute-eval.hh"
#include "tydra/render-data.hh"
#include "tydra/scene-access.hh"
#include "tydra/shader-network.hh"
#include "usdShade.hh"
#include "pprinter.hh"
#include "prim-pprint.hh"
#include "value-pprint.hh"
#include "value-types.hh"

namespace
{
	static bool ends_with(std::string_view str, std::string_view suffix)
	{
		return str.size() >= suffix.size() && str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0;
	}
}

namespace OpenVulkano::Scene
{
	void MeshLoader::ParseAssimpFile(Geometry *geometry, const std::string& file)
	{
#ifdef ASSIMP_AVAILABLE
		Assimp::Importer importer;

		const uint32_t flags = aiProcess_CalcTangentSpace | aiProcess_Triangulate | aiProcess_JoinIdenticalVertices | aiProcess_GenNormals |
			aiProcess_ImproveCacheLocality | aiProcess_RemoveRedundantMaterials | aiProcess_GenUVCoords | aiProcess_TransformUVCoords |
			aiProcess_ConvertToLeftHanded | aiProcess_PreTransformVertices | aiProcess_OptimizeGraph;

		const aiScene* scene = importer.ReadFile(file, flags);
		if (!scene) throw std::runtime_error("Failed to load file \"" + file + "\" Error: " + importer.GetErrorString());
		if (!scene->HasMeshes()) throw std::runtime_error("File \"" + file + "\" does not have any meshes");
		if (scene->mNumMeshes > 1) Logger::DATA->warn("File {0} contains more than one mesh. Only first one will be loaded", file);
		aiMesh *mesh = scene->mMeshes[0];

		geometry->aabb.Reset();
		geometry->Init(mesh->mNumVertices, mesh->mNumFaces * 3); // Reserve the space for the data
		for (unsigned int i = 0; i < mesh->mNumVertices; i++)
		{
			geometry->vertices[i].Set(mesh->mVertices[i]);
			if (mesh->HasNormals()) geometry->vertices[i].SetNormal(mesh->mNormals[i]);
			if (mesh->HasTangentsAndBitangents())
			{
				geometry->vertices[i].SetTangentAndBiTangent(mesh->mTangents[i], mesh->mBitangents[i]);
			}
			if (mesh->HasTextureCoords(0)) geometry->vertices[i].SetTextureCoordinates(mesh->mTextureCoords[0][i]);
			if (mesh->HasVertexColors(0)) geometry->vertices[i].SetColor(mesh->mColors[0][i]);
			geometry->aabb.Grow(geometry->vertices[i].position);
		}

		for (unsigned int i = 0; i < mesh->mNumFaces; i++)
		{
			const aiFace face = mesh->mFaces[i];
			if (face.mNumIndices != 3) throw std::runtime_error("Mesh is not a triangle mesh!");
			for (unsigned int j = 0; j < face.mNumIndices; j++)
			{
				if (geometry->indexType == VertexIndexType::UINT16)
				{
					static_cast<uint16_t*>(geometry->indices)[i * face.mNumIndices + j] = static_cast<uint16_t>(face.mIndices[j]);
				}
				else
				{
					static_cast<uint32_t*>(geometry->indices)[i * face.mNumIndices + j] = face.mIndices[j];
				}
			}
		}

		//TODO load bones
		//TODO load materials
		importer.FreeScene();
#else
		throw std::runtime_error("ASSIMP not available!");
#endif
	}

	void MeshLoader::ParseUSDFile(Geometry *geometry, const std::string& file)
	{
		tinyusdz::Stage stage;
		std::string warning, err;

		auto result = tinyusdz::LoadUSDFromFile(file, &stage, &warning, &err);
		if (!result)
		{
			throw std::runtime_error("Failed to load USD file: " + file);
		}

		for (auto &prim : stage.root_prims())
		{
			if (prim.type_name() == "Xform")
			{
				for (auto& child : prim.children())
				{
					auto mesh = child.as<tinyusdz::GeomMesh>();
					if (child.prim_type_name() == "Mesh" && mesh)
					{
						const auto& points = mesh->get_points();
						const auto& indices = mesh->get_faceVertexIndices();
					
						geometry->Init(static_cast<uint32_t>(points.size()), static_cast<uint32_t>(indices.size()));

						for (uint32_t i = 0; i < geometry->vertexCount; ++i)
						{
							auto point = points[i];
							OpenVulkano::Math::Vector3f p;
							p.x = point.x;
							p.y = point.y;
							p.z = point.z;
							geometry->vertices[i].position = p;
							geometry->aabb.Grow(p);
						}

						if (geometry->indexType == VertexIndexType::UINT16)
						{
							uint16_t* indices16 = static_cast<uint16_t*>(geometry->indices);
							for (uint32_t i = 0; i < geometry->indexCount; ++i)
							{
								indices16[i] = static_cast<uint16_t>(indices[i]);
							}
						}
						else
						{
							uint32_t* indices32 = static_cast<uint32_t*>(geometry->indices);
							for (uint32_t i = 0; i < geometry->indexCount; ++i)
							{
								indices32[i] = indices[i];
							}
						}
						return;
					}
				}
			}
		}
		throw std::runtime_error("No mesh found inside a xform in USD file: " + file);
	}

	Geometry* MeshLoader::LoadFromFile(const std::string& file)
	{
		Geometry* geometry = new Geometry();

		if (ends_with(file, ".usd") || ends_with(file, ".usda") || ends_with(file, ".usdc") || ends_with(file, ".usdz"))
		{
			ParseUSDFile(geometry, file);
		}
		else
		{
			ParseAssimpFile(geometry, file);
		}

		return geometry;
	}
}