/*
 * 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/UsdEncoder.hpp"
#include "Scene/ObjEncoder.hpp"
#include "IO/Archive/ArchiveWriter.hpp"
#include "IO/Archive/ZipWriter.hpp"

#include <fstream>
#include <iostream>
#include <fmt/core.h>
#include <assimp/scene.h>
#include <assimp/Exporter.hpp>
#include <assimp/mesh.h>
#include <assimp/material.h>
#include <assimp/postprocess.h>

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::MakeDescriptionForFile("model.obj", objContentsStr.size());
			zipWriter.AddFile(objDesc, objContentsStr.data());
		}
		{
			FileDescription mtlDesc = FileDescription::MakeDescriptionForFile("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::MakeDescriptionForFile("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::MakeDescriptionForFile("geometry.usda", usdFileStr.size());
			zipWriter.AddFile(usdDesc, usdFileStr.data());
		}

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

	void MeshWriter::WriteAsFBX(Geometry* geometry, const std::string& texturePath, const std::string& fbxPath)
	{
		aiScene scene;
		scene.mRootNode = new aiNode();

		aiMesh* mesh = new aiMesh();
		mesh->mNumVertices = geometry->vertexCount;
		mesh->mVertices = new aiVector3D[geometry->vertexCount];
		mesh->mNormals = new aiVector3D[geometry->vertexCount];
		mesh->mTextureCoords[0] = new aiVector3D[geometry->vertexCount];
		mesh->mMaterialIndex = 0;
		mesh->mPrimitiveTypes = aiPrimitiveType_TRIANGLE;
		mesh->mNumUVComponents[0] = 2;

		float scaling = 100; // fbx units are centimeters...
		for (uint32_t i = 0; i < geometry->vertexCount; ++i)
		{
			const 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);
			mesh->mTextureCoords[0][i] = aiVector3D(vertex.textureCoordinates.x, vertex.textureCoordinates.y, 0.0f);
		}

		mesh->mNumFaces = geometry->indexCount / 3;
		mesh->mFaces = new aiFace[mesh->mNumFaces];

		for (uint32_t i = 0; i < mesh->mNumFaces; ++i)
		{
			aiFace& face = mesh->mFaces[i];
			face.mNumIndices = 3;
			face.mIndices = new unsigned int[3];

			face.mIndices[0] = geometry->GetIndex(i * 3 + 0);
			face.mIndices[1] = geometry->GetIndex(i * 3 + 1);
			face.mIndices[2] = geometry->GetIndex(i * 3 + 2);
		}

		scene.mMeshes = new aiMesh*[1];
		scene.mMeshes[0] = mesh;
		scene.mNumMeshes = 1;
		scene.mRootNode->mMeshes = new unsigned int[1];
		scene.mRootNode->mMeshes[0] = 0;
		scene.mRootNode->mNumMeshes = 1;

		aiMaterial* material = new aiMaterial();
		scene.mMaterials = new aiMaterial*[1];
		scene.mMaterials[0] = material;
		scene.mNumMaterials = 1;

		aiString externalPath(texturePath);
		scene.mMaterials[0]->AddProperty(&externalPath, AI_MATKEY_TEXTURE(aiTextureType_DIFFUSE, 0));

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

#define SAFE_DELETE_ARRAY(arr) \
	do                         \
	{                          \
		if (arr)               \
			delete[] arr;      \
		arr = nullptr;         \
	} while (0)

		SAFE_DELETE_ARRAY(mesh->mVertices);
		SAFE_DELETE_ARRAY(mesh->mNormals);
		SAFE_DELETE_ARRAY(mesh->mTextureCoords[0]);

		for (uint32_t i = 0; i < mesh->mNumFaces; ++i)
		{
			SAFE_DELETE_ARRAY(mesh->mFaces[i].mIndices);
		}

		SAFE_DELETE_ARRAY(mesh->mFaces);

		delete mesh;
		mesh = nullptr;

		delete material;
		material = nullptr;

		SAFE_DELETE_ARRAY(scene.mMeshes);
		SAFE_DELETE_ARRAY(scene.mRootNode->mMeshes);
		SAFE_DELETE_ARRAY(scene.mMaterials);
		delete scene.mRootNode;
		scene.mRootNode = nullptr;
#undef SAFE_DELETE_ARRAY

		if (result != aiReturn_SUCCESS)
		{
			throw std::runtime_error("Unable to write a fbx file to " + fbxPath + ": " + exporter.GetErrorString());
		}
	}
}