/*
 * 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 "ImageLoaderJpeg.hpp"
#include "Base/Utils.hpp"
#include "Base/Logger.hpp"
#include "Data/Containers/Array.hpp"
#include <fstream>
#include <cstring>

#if __has_include("turbojpeg.h")
	#include <turbojpeg.h>
#else
	#define STB_IMAGE_IMPLEMENTATION
	#include <stb_image.h>
#endif

namespace OpenVulkano::Image
{
	std::unique_ptr<Image> ImageLoaderJpeg::loadFromFile(const std::string& filePath)
	{
		const auto file = Utils::ReadFile(filePath);
		return loadJpeg(reinterpret_cast<const uint8_t*>(file.Data()), file.Size());
	}

	std::unique_ptr<Image> ImageLoaderJpeg::loadFromMemory(const std::vector<uint8_t>& buffer)
	{
		return loadJpeg(buffer.data(), buffer.size());
	}

	Math::Vector2i ImageLoaderJpeg::GetImageDimensions(const std::string& filename)
	{
#if __has_include("turbojpeg.h")
		auto image = Utils::ReadFile(filename);
		tjhandle jpegDecompressor = tjInitDecompress();
		if (!jpegDecompressor)
		{
			Logger::FILESYS->error("Failed to read jpeg header. Error: {}", tjGetErrorStr());
			return Math::Vector2i { -1, -1 };
		}
		Math::Vector2i res = {};
		int size = 0, jpegSubsamp = 0;
		int status = tjDecompressHeader2(jpegDecompressor, reinterpret_cast<unsigned char*>(image.Data()), image.Size(),
		                                 &res.x, &res.y, &jpegSubsamp);
		if (status != 0)
		{
			Logger::FILESYS->error("Failed to read jpeg header. Error: {}", tjGetErrorStr());
			tjDestroy(jpegDecompressor);
			return Math::Vector2i { -1, -1 };
		}
		tjDestroy(jpegDecompressor);
		return res;
#else
		return IImageLoader::GetDimensionsInternal(filename);
#endif
	}

	std::unique_ptr<Image> ImageLoaderJpeg::loadJpeg(const uint8_t* data, size_t size)
	{
#if __has_include("turbojpeg.h")
		{
			Image result;

			int rows = 0, cols = 0;
			unsigned char* compressedImage = const_cast<uint8_t*>(data);

			int jpegSubsamp = 0;
			tjhandle jpegDecompressor = tjInitDecompress();
			if (!jpegDecompressor)
			{
				Logger::FILESYS->error("Failed to read jpeg header. Error: {}", tjGetErrorStr());
				return nullptr;
			}
			int status = tjDecompressHeader2(jpegDecompressor, compressedImage, size, &cols, &rows, &jpegSubsamp);
			if (status != 0)
			{
				Logger::FILESYS->error("Failed to read jpeg header. Error: {}", tjGetErrorStr());
				return nullptr;
			}

			const int channels = 4;
			result.data = OpenVulkano::Array<uint8_t>(cols * rows * channels);
			result.dataFormat = OpenVulkano::DataFormat::B8G8R8A8_UINT;

			tjDecompress2(jpegDecompressor, compressedImage, size, result.data.Data(), cols, 0 /*pitch*/, rows,
			              TJPF_BGRA, TJFLAG_FASTDCT);
			tjDestroy(jpegDecompressor);
			result.resolution.x = cols;
			result.resolution.y = rows;
			result.resolution.z = 1;

			return std::make_unique<Image>(std::move(result));
		}
#else
		{
			return loadData(data, static_cast<int>(size));
		}
#endif
	}
}