/*
 * 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 "MeshWriter.hpp"
#include "Base/Utils.hpp"
#include "IO/MemMappedFile.hpp"
#include "Scene/Geometry.hpp"
#include "Scene/Vertex.hpp"
#include "Scene/Export/UsdEncoder.hpp"
#include "ObjEncoder.hpp"
#include "IO/Archive/ArchiveWriter.hpp"
#include "IO/Archive/ZipWriter.hpp"

#include <fstream>
#include <fmt/core.h>

#if __has_include("assimp/Exporter.hpp")
#include <assimp/scene.h>
#include <assimp/Exporter.hpp>
#include <assimp/mesh.h>
#include <assimp/material.h>
#endif

namespace
{
#if __has_include("assimp/Exporter.hpp")
	void SetupAssimpScene(OpenVulkano::Scene::Geometry* geometry, aiScene& scene, aiNode& rootNode, aiMesh& mesh,
	                      std::unique_ptr<unsigned int[]>& indices, bool withTexCoords, float scaling)
	{
		mesh.mVertices = nullptr;
		mesh.mNormals = nullptr;
		mesh.mTangents = nullptr;
		mesh.mBitangents = nullptr;
		mesh.mFaces = nullptr;
		for(unsigned int i = 0; i < AI_MAX_NUMBER_OF_TEXTURECOORDS; ++i) {
			mesh.mTextureCoords[i] = nullptr;
			mesh.mNumUVComponents[i] = 0;
		}
		for(unsigned int i = 0; i < AI_MAX_NUMBER_OF_COLOR_SETS; ++i) {
			mesh.mColors[i] = nullptr;
		}

		rootNode.mName = aiString("RootNode");
		rootNode.mChildren = nullptr;
		rootNode.mNumChildren = 0;
		scene.mRootNode = &rootNode;

		mesh.mName = aiString("Mesh");
		mesh.mNumVertices = geometry->vertexCount;
		mesh.mMaterialIndex = 0;
		mesh.mPrimitiveTypes = aiPrimitiveType_TRIANGLE;

		mesh.mVertices = new aiVector3D[geometry->vertexCount];
		mesh.mNormals = new aiVector3D[geometry->vertexCount];

		if (withTexCoords)
		{
			mesh.mNumUVComponents[0] = 2;
			mesh.mTextureCoords[0] = new aiVector3D[geometry->vertexCount];
		}

		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const OpenVulkano::Vertex& vertex = geometry->vertices[i];
			mesh.mVertices[i] = aiVector3D(vertex.position.x, vertex.position.y, vertex.position.z) * scaling;
			mesh.mNormals[i] = aiVector3D(vertex.normal.x, vertex.normal.y, vertex.normal.z);
			if (withTexCoords)
			{
				mesh.mTextureCoords[0][i] = aiVector3D(vertex.textureCoordinates.x, vertex.textureCoordinates.y, 0.0f);
			}
		}

		mesh.mNumFaces = geometry->indexCount / 3;
		mesh.mFaces = new aiFace[mesh.mNumFaces];
		indices = std::make_unique<unsigned int[]>(geometry->indexCount);

		for (unsigned int i = 0; i < geometry->indexCount; ++i)
		{
			indices[i] = geometry->GetIndex(i);
		}

		for (unsigned int i = 0; i < mesh.mNumFaces; ++i)
		{
			mesh.mFaces[i].mNumIndices = 3;
			mesh.mFaces[i].mIndices = &indices[i * 3];
		}

		scene.mNumMeshes = 1;
		scene.mMeshes = new aiMesh*[1] { &mesh };

		scene.mNumMaterials = 1;
		scene.mMaterials = new aiMaterial*[1] { new aiMaterial() };

		rootNode.mNumMeshes = 1;
		rootNode.mMeshes = new unsigned int[1] { 0 };
	}
#endif
}

namespace OpenVulkano::Scene
{
	void MeshWriter::WriteAsOBJ(Geometry* geometry, const std::string& filePath)
	{
		std::ofstream file(filePath);

		if (!file.is_open()) [[unlikely]]
			throw std::runtime_error("Failed to open file '" + filePath + "' for writing!");
		
		WriteObjContents(geometry, "", file);
		file.close();
	}

	void MeshWriter::WriteAsUSD(Geometry* geometry, const std::string& filePath)
	{
		std::ofstream file(filePath);
		if (!file.is_open()) [[unlikely]]
			throw std::runtime_error("Failed to open file '" + filePath + "' for writing!");
		WriteUsdContents(file, geometry);
		file.close();
	}

	void MeshWriter::WriteObjAsZip(Geometry* geometry, const std::string& texturePath, const std::string& zipPath)
	{
		OpenVulkano::ArchiveWriter zipWriter(zipPath.c_str());

		{
			std::stringstream objContents;
			WriteObjContents(geometry, DEFAULT_OBJ_MATERIAL_NAME, objContents);
			std::string objContentsStr = objContents.str();
			FileDescription objDesc = FileDescription::MkFile("model.obj", objContentsStr.size());
			zipWriter.AddFile(objDesc, objContentsStr.data());
		}
		{
			FileDescription mtlDesc = FileDescription::MkFile("material.mtl", DEFAULT_OBJ_MATERIAL_CONTENTS.size());
			zipWriter.AddFile(mtlDesc, DEFAULT_OBJ_MATERIAL_CONTENTS.data());
		}

		if (!texturePath.empty() && std::filesystem::exists(texturePath))
		{
			MemMappedFile textureFile(texturePath);
			FileDescription texDesc = FileDescription::MkFile("texture.png", textureFile.Size());
			zipWriter.AddFile(texDesc, textureFile.Data());
		}
	}

	void MeshWriter::WriteAsUSDZ(Geometry* geometry, const std::string& texturePath, const std::string& usdzPath)
	{
		OpenVulkano::ZipWriter zipWriter(usdzPath, true);
		
		{
			std::stringstream usdFile;
			WriteUsdContents(usdFile, geometry);
			std::string usdFileStr = usdFile.str();
			FileDescription usdDesc = FileDescription::MkFile("geometry.usda", usdFileStr.size());
			zipWriter.AddFile(usdDesc, usdFileStr.data());
		}

		if (!texturePath.empty() && std::filesystem::exists(texturePath))
		{
			MemMappedFile textureFile(texturePath);
			FileDescription texDesc = FileDescription::MkFile("texture.png", textureFile.Size());
			zipWriter.AddFile(texDesc, textureFile.Data());
		}
	}

	void MeshWriter::WriteAsSTL(Geometry* geometry, const std::string& filePath, bool binary)
	{
#if __has_include("assimp/Exporter.hpp")
		std::unique_ptr<aiScene> scene(new aiScene());
		std::unique_ptr<aiNode> rootNode(new aiNode());
		std::unique_ptr<aiMesh> mesh(new aiMesh());
		std::unique_ptr<unsigned int[]> indices;

		SetupAssimpScene(geometry, *scene, *rootNode, *mesh, indices, false, 1.0f);

		Assimp::Exporter exporter;
		const char* formatId = binary ? "stlb" : "stl";
		aiReturn result = exporter.Export(scene.get(), formatId, filePath);

		scene->mRootNode = nullptr;
		scene->mMeshes = nullptr;
		for (uint32_t i = 0; i < mesh->mNumFaces; ++i)
		{
			aiFace& face = mesh->mFaces[i];
			face.mIndices = nullptr;
		}

		if (result != aiReturn_SUCCESS)
		{
			throw std::runtime_error("Unable to write STL file to " + filePath + ": " + exporter.GetErrorString());
		}
#else
		throw std::runtime_error("Unable to export to STL: Assimp is not available!");
#endif
	}

	void MeshWriter::WriteAsFBX(Geometry* geometry, const std::string& texturePath, const std::string& fbxPath)
	{
#if __has_include("assimp/Exporter.hpp")
		std::unique_ptr<aiScene> scene(new aiScene());
		std::unique_ptr<aiNode> rootNode(new aiNode());
		std::unique_ptr<aiMesh> mesh(new aiMesh());
		std::unique_ptr<unsigned int[]> indices;

		SetupAssimpScene(geometry, *scene, *rootNode, *mesh, indices, true, 100.0f);

		if (!texturePath.empty())
		{
			aiString externalPath(texturePath);
			scene->mMaterials[0]->AddProperty(&externalPath, AI_MATKEY_TEXTURE(aiTextureType_DIFFUSE, 0));
		}

		Assimp::Exporter exporter;
		aiReturn result = exporter.Export(scene.get(), "fbx", fbxPath);

		scene->mRootNode = nullptr;
		scene->mMeshes = nullptr;
		for (uint32_t i = 0; i < mesh->mNumFaces; ++i)
		{
			aiFace& face = mesh->mFaces[i];
			face.mIndices = nullptr;
		}

		if (result != aiReturn_SUCCESS)
		{
			throw std::runtime_error("Unable to write a fbx file to " + fbxPath + ": " + exporter.GetErrorString());
		}
#else
		throw std::runtime_error("Unable to convert the scene to FBX: Assimp is not available!");
#endif
	}
}