/*
 * 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 "Scene/Geometry.hpp"
#include "Scene/Vertex.hpp"
#include "IO/Archive/ArchiveWriter.hpp"
#include <fstream>
#include <fmt/core.h>
#include <tinyusdz.hh>
#include <pprinter.hh>
#include <prim-pprint.hh>
#include <value-pprint.hh>
#include <tiny-format.hh>
#include <tiny-variant.hh>
#include <usda-reader.hh>
#include <usda-writer.hh>
#include <usdc-reader.hh>
#include <usdc-writer.hh>

namespace
{
	uint32_t GetIndexFromGeometry(OpenVulkano::Scene::Geometry* geometry, int index)
	{
		uint32_t result = 0;

		if (geometry->indexType == OpenVulkano::Scene::VertexIndexType::UINT16)
		{
			uint16_t *indices = static_cast<uint16_t *>(geometry->indices);
			result = indices[index];
		}
		else if (geometry->indexType == OpenVulkano::Scene::VertexIndexType::UINT32)
		{
			uint32_t *indices = static_cast<uint32_t *>(geometry->indices);
			result = indices[index];
		}
		else
		{
			throw std::runtime_error("Invalid geometry index type");
		}

		return result;
	}

	std::string ReadEntireFile(const std::string& fname)
	{
		FILE *file = fopen(fname.c_str(), "rb");
		std::string buffer;
		if(file)
		{
			fseek(file, 0, SEEK_END);
			size_t size = ftell(file);
			fseek(file, 0, SEEK_SET);
			buffer = std::string(size, ' ');
			fread(&buffer[0], size, 1, file);
			fclose(file);
		}
		return buffer;
	}
	
	void WriteGeometryWithoutTexturesUsingTinyusdz(OpenVulkano::Scene::Geometry* geometry, const std::string& filePath)
	{
		tinyusdz::Stage stage;
		tinyusdz::Xform xform;
		tinyusdz::GeomMesh mesh;
		mesh.name = "TheMesh";

		std::vector<tinyusdz::value::point3f> pts(geometry->vertexCount);
		std::vector<int> indices(geometry->indexCount);
		tinyusdz::Attribute uvAttr;
		std::vector<tinyusdz::value::texcoord2f> uvs(geometry->vertexCount);

		for (uint32_t i = 0; i < geometry->vertexCount; ++i)
		{
			const OpenVulkano::Vertex& v = geometry->vertices[i];
			pts[i].x = v.position.x;
			pts[i].y = v.position.y;
			pts[i].z = v.position.z;
			uvs[i] = { v.textureCoordinates.x, v.textureCoordinates.y };
		}

		mesh.points.set_value(pts);
		uvAttr.set_value(uvs);

		std::vector<int> counts(geometry->indexCount / 3, 3); // NOTE(vb): The value 3 is kind of arbitrary, but this array must be in the mesh!
		mesh.faceVertexCounts.set_value(counts);

		for (uint32_t i = 0; i < geometry->indexCount; ++i)
		{
			uint32_t index = GetIndexFromGeometry(geometry, i);
			indices[i] = index;
		}
		mesh.faceVertexIndices.set_value(indices);

		uvAttr.metas().interpolation = tinyusdz::Interpolation::FaceVarying;
		tinyusdz::Property uvProp(uvAttr);
		mesh.props.emplace("primvars:UVMap", uvProp);

		tinyusdz::Prim xformPrim(xform);
		tinyusdz::Prim meshPrim(mesh);
		std::string err;
		if (!xformPrim.add_child(std::move(meshPrim), true, &err))
		{
			throw std::runtime_error("Failed to construct scene: " + err);
		}

		if (!stage.add_root_prim(std::move(xformPrim)))
		{
			throw std::runtime_error("Failed to add prim to stage root: " + stage.get_error());
		}

		stage.metas().defaultPrim = tinyusdz::value::token(xformPrim.element_name());
		stage.metas().comment = "Generated by OpenVulkanoCpp";
		if (!stage.commit())
		{
			throw std::runtime_error("Failed to commit stage: " + stage.get_error());
		}

		std::ofstream file(filePath);
		if (!file.is_open())
			throw std::runtime_error("Failed to open file '" + filePath + "' for writing!");
		std::string scene = to_string(stage);
		file << scene << "\n";
		file.close();
	}

	std::vector<uint8_t> ConvertToUSDC(const std::string& contents)
	{
		std::vector<uint8_t> result;
		tinyusdz::Stage stage;
		std::string warn, err;
		assert(tinyusdz::LoadUSDAFromMemory((const uint8_t *)contents.data(), (const size_t)contents.size(), "", &stage, &warn, &err));
		bool ret = tinyusdz::usdc::SaveAsUSDCToMemory(stage, &result, &warn, &err); // As for now, this reports that it's not implemented...
		return result;
	}

	std::string GetUsdContents(OpenVulkano::Scene::Geometry* geometry, const std::string texturePath = "")
	{
		std::ostringstream file;
		std::ostringstream points, normals, indices, texCoords, faceCounts;
		points << std::fixed << std::setprecision(6);
		normals << std::fixed << std::setprecision(6);

		for (size_t i = 0; i < geometry->vertexCount; ++i)
		{
			const auto& v = geometry->vertices[i];
			points << "(" << v.position.x << ", " << v.position.y << ", " << v.position.z << ")";
			normals << "(" << v.normal.x << ", " << v.normal.y << ", " << v.normal.z << ")";
			if (i < geometry->vertexCount - 1)
			{
				points << ", ";
				normals << ", ";
			}
		}

		for (size_t i = 0; i < geometry->indexCount; ++i)
		{
			indices << GetIndexFromGeometry(geometry, i);
			if (i < geometry->indexCount - 1)
			{
				indices << ", ";
			}
			if ((i + 1) % 3 == 0)
			{
				faceCounts << "3";
				if (i < geometry->indexCount - 1)
				{
					faceCounts << ", ";
				}
			}
		}

		texCoords << std::fixed << std::setprecision(6);

		for (size_t i = 0; i < geometry->indexCount; ++i)
		{
			const size_t vertexIndex = GetIndexFromGeometry(geometry, i);
			const auto& v = geometry->vertices[vertexIndex];
			texCoords << "(" << v.textureCoordinates.x << ", " << v.textureCoordinates.y << ")";
			if (i < geometry->indexCount - 1)
			{
				texCoords << ", ";
			}
		}

		file << R"(#usda 1.0
(
    defaultPrim = "root"
    doc = "Exported from OpenVulkano"
    metersPerUnit = 1
    upAxis = "Z"
)

def Xform "root" (
    customData = {
        dictionary Blender = {
            bool generated = 1
        }
    }
)
{
    def Xform "model"
    {
        custom string userProperties:blender:object_name = "model"
        float3 xformOp:rotateXYZ = (89.99999, -0, 0)
        float3 xformOp:scale = (1, 1, 1)
        double3 xformOp:translate = (0, 0, 0)
        uniform token[] xformOpOrder = ["xformOp:translate", "xformOp:rotateXYZ", "xformOp:scale"]

        def Mesh "model" (
            active = true
            prepend apiSchemas = ["MaterialBindingAPI"]
        )
        {
            uniform bool doubleSided = 1
            float3[] extent = [(-0.5, -0.5, 0), (0.5, 0.5, 0)]
            int[] faceVertexCounts = [)"
         << faceCounts.str() << R"(]
            int[] faceVertexIndices = [)"
         << indices.str() << R"(]
            rel material:binding = </root/_materials/Material0>
            normal3f[] normals = [)"
         << normals.str() << R"(] (
                interpolation = "faceVarying"
            )
            point3f[] points = [)"
         << points.str() << R"(]
            texCoord2f[] primvars:st = [)"
         << texCoords.str() << R"(] (
                interpolation = "faceVarying"
            )
            uniform token subdivisionScheme = "none"
            custom string userProperties:blender:data_name = "model"
        }
    }

    def Scope "_materials"
    {
        def Material "Material0"
        {
            token outputs:surface.connect = </root/_materials/Material0/Principled_BSDF.outputs:surface>
            custom string userProperties:blender:data_name = "Material0"

            def Shader "Principled_BSDF"
            {
                uniform token info:id = "UsdPreviewSurface"
                float inputs:clearcoat = 0
                float inputs:clearcoatRoughness = 0.03
                color3f inputs:diffuseColor.connect = </root/_materials/Material0/Image_Texture.outputs:rgb>
                float inputs:ior = 1.5
                float inputs:metallic = 0
                float inputs:opacity = 1
                float inputs:roughness = 1
                float inputs:specular = 0
                token outputs:surface
            }

            def Shader "Image_Texture"
            {
                uniform token info:id = "UsdUVTexture"
                asset inputs:file = @./texture.png@
                token inputs:sourceColorSpace = "sRGB"
                float2 inputs:st.connect = </root/_materials/Material0/uvmap.outputs:result>
                token inputs:wrapS = "repeat"
                token inputs:wrapT = "repeat"
                float3 outputs:rgb
            }

            def Shader "uvmap"
            {
                uniform token info:id = "UsdPrimvarReader_float2"
                string inputs:varname = "st"
                float2 outputs:result
            }
        }
    }
}
)";
		return file.str();
	}
}

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

		if (!file.is_open())
			throw std::runtime_error("Failed to open file '" + filePath + "' for writing!");
		
		// Vertices
		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const OpenVulkano::Vertex& v = geometry->vertices[i];
			std::string line;
			line = fmt::format("v {} {} {}\n", v.position.x, v.position.y, v.position.z);
			file << line;
		}

		// Normals
		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const OpenVulkano::Vertex& v = geometry->vertices[i];
			std::string line;
			line = fmt::format("vn {} {} {}\n", v.normal.x, v.normal.y, v.normal.z);
			file << line;
		}

		// TexCoords
		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const OpenVulkano::Vertex& v = geometry->vertices[i];
			std::string line;
			line = fmt::format("vt {} {}\n", v.textureCoordinates.x, v.textureCoordinates.y);
			file << line;
		}

		// Indices
		for (int i = 0; i < geometry->indexCount; i += 3)
		{
			uint32_t i0 = GetIndexFromGeometry(geometry, i + 0) + 1;
			uint32_t i1 = GetIndexFromGeometry(geometry, i + 1) + 1;
			uint32_t i2 = GetIndexFromGeometry(geometry, i + 2) + 1;
			std::string line = fmt::format("f {0}/{0}/{0} {1}/{1}/{1} {2}/{2}/{2}\n", i0, i1, i2);
			file << line;
		}

		file.close();
	}

	void MeshWriter::WriteAsUSD(Geometry* geometry, const std::string& filePath)
	{
		std::ofstream file(filePath);
		if (!file.is_open())
			throw std::runtime_error("Failed to open file '" + filePath + "' for writing!");
		std::string scene = GetUsdContents(geometry);
		file << scene << "\n";
		file.close();
	}

	void MeshWriter::WriteObjAsZip(Geometry* geometry, const std::string& zipPath, const std::string& texturePath)
	{
		std::stringstream mtlContent;
		std::string materialName = "Material0";
		mtlContent << "newmtl " << materialName << "\n";
		mtlContent << "Ka 1.000 1.000 1.000\n"; // Ambient
		mtlContent << "Kd 1.000 1.000 1.000\n"; // Diffuse
		mtlContent << "Ks 0.000 0.000 0.000\n"; // Specular
		if (!texturePath.empty())
		{
			mtlContent << "map_Ka texture.png\n"; // Ambient map
			mtlContent << "map_Kd texture.png\n"; // Texture map
		}

		std::stringstream objContent;
		objContent << "# OBJ file generated by OpenVulkanoCpp\n";
		objContent << "mtllib material.mtl\n";
		objContent << "usemtl " << materialName << "\n";
		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const auto& v = geometry->vertices[i];
			objContent << fmt::format("v {} {} {}\n", v.position.x, v.position.y, v.position.z);
		}
		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const auto& v = geometry->vertices[i];
			objContent << fmt::format("vn {} {} {}\n", v.normal.x, v.normal.y, v.normal.z);
		}
		for (int i = 0; i < geometry->vertexCount; ++i)
		{
			const auto& v = geometry->vertices[i];
			objContent << fmt::format("vt {} {}\n", v.textureCoordinates.x, v.textureCoordinates.y);
		}
		for (int i = 0; i < geometry->indexCount; i += 3)
		{
			uint32_t i0 = GetIndexFromGeometry(geometry, i) + 1;
			uint32_t i1 = GetIndexFromGeometry(geometry, i + 1) + 1;
			uint32_t i2 = GetIndexFromGeometry(geometry, i + 2) + 1;
			objContent << fmt::format("f {0}/{0}/{0} {1}/{1}/{1} {2}/{2}/{2}\n", i0, i1, i2);
		}

		OpenVulkano::ArchiveWriter zipWriter(zipPath.c_str());

		auto objDesc = OpenVulkano::FileDescription::MakeDescriptionForFile("model.obj", objContent.str().size());
		zipWriter.AddFile(objDesc, objContent.str().data());

		auto mtlDesc = OpenVulkano::FileDescription::MakeDescriptionForFile("material.mtl", mtlContent.str().size());
		zipWriter.AddFile(mtlDesc, mtlContent.str().data());

		if (!texturePath.empty() && std::filesystem::exists(texturePath))
		{
			auto textureFileSize = std::filesystem::file_size(texturePath);
			auto texDesc = OpenVulkano::FileDescription::MakeDescriptionForFile("texture.png", textureFileSize);
			zipWriter.AddFile(texDesc, ReadEntireFile(texturePath).c_str());
		}
	}

	void MeshWriter::WriteAsUSDZ(Geometry* geometry, const std::string& usdzPath, const std::string& texturePath)
	{
		OpenVulkano::ArchiveConfiguration config;
		config.type = OpenVulkano::ArchiveType::ZIP;
		OpenVulkano::ArchiveWriter zipWriter(usdzPath.c_str(), config);
		
		std::string usd = GetUsdContents(geometry, texturePath);
		auto usdDesc = OpenVulkano::FileDescription::MakeDescriptionForFile("geometry.usda", usd.size());
		zipWriter.AddFile(usdDesc, usd.data());

		if (!texturePath.empty() && std::filesystem::exists(texturePath))
		{
			auto textureFileSize = std::filesystem::file_size(texturePath);
			auto texDesc = OpenVulkano::FileDescription::MakeDescriptionForFile("texture.png", textureFileSize);
			zipWriter.AddFile(texDesc, ReadEntireFile(texturePath).c_str());
		}
	}
}