/*
 * 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 <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)
	{
		std::ifstream file(filePath, std::ios::binary);
		if (!file) { throw std::runtime_error("Could not open file: " + filePath); }
		std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
		return loadJpeg(buffer.data(), buffer.size());
	}

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

	std::unique_ptr<Image> ImageLoaderJpeg::loadJpeg(const uint8_t* data, size_t size)
	{
		Image result;

		int rows, cols;
#if __has_include("turbojpeg.h")
		{
			unsigned char* compressedImage = const_cast<uint8_t*>(data);

			int jpegSubsamp;
			tjhandle jpegDecompressor = tjInitDecompress();
			tjDecompressHeader2(jpegDecompressor, compressedImage, size, &cols, &rows, &jpegSubsamp);

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

			tjDecompress2(jpegDecompressor, compressedImage, size, result.data.Data(), cols, 0 /*pitch*/, rows,
			              TJPF_RGBA, TJFLAG_FASTDCT);
			tjDestroy(jpegDecompressor);
		}
#else
		{
			int channels;
			uint8_t* pixelData = stbi_load_from_memory(data, size, &cols, &rows, &channels, 3);
			result.data = OpenVulkano::Array<uint8_t>(cols * rows * channels);
			result.dataFormat = channels == 3 ? OpenVulkano::DataFormat::R8G8B8_UINT :
			                                    OpenVulkano::DataFormat::R8G8B8A8_UINT;
			std::memcpy(result.data.Data(), pixelData, result.data.Size());
			stbi_image_free(pixelData);
		}
#endif
		result.resolution.x = cols;
		result.resolution.y = rows;
		result.resolution.z = 1;

		return std::make_unique<Image>(std::move(result));
	}
}