/*
 * 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 "FontAtlas.hpp"
#include "Base/Logger.hpp"
#include "Base/Wrapper.hpp"
#include "Image/ImageLoader.hpp"
#include "Extensions/STBZlibCompressor.hpp"
#include <fstream>
#include <zstd.h>
#define STBI_MSC_SECURE_CRT
#define STB_IMAGE_WRITE_IMPLEMENTATION
#define STBIW_ZLIB_COMPRESS Extensions::STBZlibCompressor
#include <stb_image_write.h>

namespace OpenVulkano::Scene
{
	namespace
	{
		struct FontAtlasFileFlags
		{
			uint8_t packed = 1;
			uint8_t version = 1;
			uint16_t reserved = 0;
		};
	}

	FontAtlas::FontAtlas(const std::filesystem::path& path)
	{
		auto content = Utils::ReadFile(path);
		Load(content);
	}

	void FontAtlas::Save(const std::filesystem::path& path) const
	{
		if (!*this) { Logger::DATA->warn("Can't save empty font atlas!"); return; };

		std::fstream fs(path, std::ios_base::out | std::ios_base::binary);

		// Save image
		{
			stbi_flip_vertically_on_write(1);
			int pngSize;
			UniqueMalloc<uint8_t> png(stbi_write_png_to_mem(static_cast<const uint8_t*>(m_texture.textureBuffer), m_texture.format.GetBytesPerPixel() * m_texture.resolution.x, m_texture.resolution.x, m_texture.resolution.y, m_texture.format.GetBytesPerPixel(), &pngSize));
			fs.write(reinterpret_cast<const std::ostream::char_type*>(png.get()), pngSize);
		}

		// Save metadata
		{
			const Array<char> metadata = SerializeMetadata();
			Array<char> compressed(ZSTD_compressBound(metadata.Size()));
			size_t compressedSize = ZSTD_compress(compressed.Data(), compressed.Size(), metadata.Data(), metadata.Size(), 5);
			if (ZSTD_isError(compressedSize)) throw std::runtime_error(std::string("Compression failed: ") + ZSTD_getErrorName(compressedSize));
			fs.write(compressed.Data(), compressedSize);
			fs.write(reinterpret_cast<const char*>(&compressedSize), sizeof(uint64_t));
		}

		FontAtlasFileFlags flags;
		fs.write(reinterpret_cast<const char*>(&flags), sizeof(FontAtlasFileFlags));
		fs.close();
	}

	Array<char> FontAtlas::SerializeMetadata() const
	{
		Array<char> data(sizeof(Metadata) + m_glyphs.size() * (sizeof(GlyphInfo) + sizeof(uint32_t)));
		char* ptr = data.Data();
		memcpy(ptr, &m_metadata, sizeof(Metadata));
		ptr += sizeof(Metadata);
		for (const auto& [key, val] : m_glyphs)
		{
			memcpy(ptr, &key, sizeof(uint32_t));
			ptr += sizeof(uint32_t);
			memcpy(ptr, &val, sizeof(GlyphInfo));
			ptr += sizeof(GlyphInfo);
		}
		return data;
	}

	void FontAtlas::DeserializeMetadata(const std::span<char>& data)
	{
		const char* d = data.data();
		const char* end = data.data() + data.size();
		uint32_t unicode = 0;

		std::memcpy(&m_metadata, d, sizeof(Metadata));
		d += sizeof(Metadata);
		while (d < end)
		{
			std::memcpy(&unicode, d, sizeof(uint32_t));
			d += sizeof(uint32_t);
			GlyphInfo& info = m_glyphs[unicode];
			std::memcpy(&info, d, sizeof(GlyphInfo));
			d += sizeof(GlyphInfo);
		}
	}


	void FontAtlas::Load(const std::span<char> data)
	{
		if (data.size() < 16) { Logger::DATA->warn("Font atlas file is invalid!"); return; };
		FontAtlasFileFlags flags;
		std::memcpy(&flags, data.data() + data.size() - sizeof(flags), sizeof(flags));
		size_t headerSize = sizeof(FontAtlasFileFlags) + sizeof(uint64_t);
		uint64_t metadataSize = *reinterpret_cast<uint64_t*>(data.data() + data.size() - headerSize);
		char* metadata = data.data() + data.size() - headerSize - metadataSize;

		if (flags.packed == 0) throw std::runtime_error("No longer support loading of unpacked font atlas!");
		LoadImage({ data.data(), data.size() - headerSize - metadataSize });

		std::span metadataSpan(metadata, metadataSize);
		if (flags.version == 0) LoadLegacy(metadataSpan);
		else LoadNew(metadataSpan);
		if (GetAtlasType() >= FontAtlasType::BITMAP) m_texture.m_samplerConfig = &SamplerConfig::NEAREST;
	}

	void FontAtlas::LoadLegacy(const std::span<char> data)
	{
		const char* d = data.data();
		const char* end = data.data() + data.size();
		uint32_t unicode = 0;

		std::memcpy(&m_metadata, d, sizeof(Metadata));
		d += sizeof(Metadata);
		while (d < end)
		{
			std::memcpy(&unicode, d, sizeof(uint32_t));
			d += sizeof(uint32_t);
			GlyphInfo& info = m_glyphs[unicode];
			for (int i = 0; i < 4; i++)
			{
				std::memcpy(&info.pos[i], d, sizeof(GlyphInfo::pos) / 4);
				d += sizeof(GlyphInfo::pos) / 2;
			}
			std::memcpy(&info.uv, d, sizeof(GlyphInfo::uv));
			d += sizeof(GlyphInfo::uv);
			std::memcpy(&info.advance, d, sizeof(GlyphInfo::advance));
			d += sizeof(GlyphInfo::advance) * 2;
		}
	}

	void FontAtlas::LoadNew(const std::span<char> data)
	{
		size_t originalSize = ZSTD_getFrameContentSize(data.data(), data.size());
		if (originalSize == ZSTD_CONTENTSIZE_ERROR) throw std::runtime_error("Not compressed by zstd");
		if (originalSize == ZSTD_CONTENTSIZE_UNKNOWN) throw std::runtime_error("Original size unknown");

		Array<char> metadata(originalSize);
		size_t decompressedBytes = ZSTD_decompress(metadata.Data(), metadata.Size(), data.data(), data.size());
		if (ZSTD_isError(decompressedBytes))  throw std::runtime_error(std::string("Decompression failed: ") + ZSTD_getErrorName(decompressedBytes));
		if (decompressedBytes != originalSize) throw std::runtime_error("Decompressed size does not match expected size");
		DeserializeMetadata(metadata);
	}

	void FontAtlas::LoadImage(const std::span<char> data)
	{
		auto img = Image::IImageLoader::loadData(reinterpret_cast<const uint8_t*>(data.data()), data.size());
		m_imgData = std::move(img->data);
		m_texture.format = img->dataFormat;
		m_texture.resolution = img->resolution;
		m_texture.size = m_imgData.Size();
		m_texture.textureBuffer = m_imgData.Data();
	}

	void FontAtlas::Init(const Math::Vector2ui textureResolution, const double lineHeight, const FontAtlasType atlasType)
	{
		m_metadata = { lineHeight, atlasType };
		m_texture.format = atlasType.GetChannelCount() == 1 ? DataFormat::R8_UNORM : DataFormat::R8G8B8A8_UNORM;
		m_texture.resolution = { textureResolution, 1 };
		m_imgData = Array<uint8_t>(m_texture.format.CalculatedSize(m_texture.resolution.x, m_texture.resolution.y));
		m_texture.textureBuffer = m_imgData.Data();
		m_texture.size = m_imgData.Size();
		if (atlasType >= FontAtlasType::BITMAP) m_texture.m_samplerConfig = &SamplerConfig::NEAREST;
	}
}