/*
 * 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 "ImageLoader.hpp"
#include "Base/Logger.hpp"
#include <memory>

#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>

namespace OpenVulkano::Image
{
	std::unique_ptr<Image> IImageLoader::loadData(const uint8_t* data, int size, int desiredChannels)
	{
		Image result;
		int rows, cols, channels;
		stbi_set_flip_vertically_on_load(true);
		uint8_t* pixelData = stbi_load_from_memory(data, static_cast<int>(size), &rows, &cols, &channels, desiredChannels);
		if (desiredChannels != 0 && channels < desiredChannels)
		{
			Logger::INPUT->warn(
			"Stbi load image channels mismatch. Desired channels = {}, actual amount of channels in image = {}",
			desiredChannels, channels);
		}
		switch (channels)
		{
			case 1:
				result.dataFormat = OpenVulkano::DataFormat::R8_UNORM;
				break;
			case 2:
				result.dataFormat = OpenVulkano::DataFormat::R8G8_UNORM; 
				break;
			case 3:
			case 4:
				result.dataFormat = OpenVulkano::DataFormat::R8G8B8A8_UNORM;
				break;
		}
		result.resolution.x = cols;
		result.resolution.y = rows;
		result.resolution.z = 1;
		if (channels == 3)
		{ 
			result.data = OpenVulkano::Array<uint8_t>(cols * rows * 4);
			for (size_t srcPos = 0, dstPos = 0; srcPos < cols * rows * 3; srcPos += 3, dstPos += 4)
			{
				result.data[dstPos] = pixelData[srcPos];
				result.data[dstPos + 1] = pixelData[srcPos + 1];
				result.data[dstPos + 2] = pixelData[srcPos + 2];
				result.data[dstPos + 3] = 255;
			}
		}
		else
		{ 
			result.data = OpenVulkano::Array<uint8_t>(cols * rows * channels);
			std::memcpy(result.data.Data(), pixelData, result.data.Size());
		}
		stbi_image_free(pixelData);
		return std::make_unique<Image>(std::move(result));
	}

	Math::Vector2i IImageLoader::GetDimensionsInternal(const std::string& filename)
	{
		Math::Vector2i res = {};
		int channels;
		if (!stbi_info(filename.c_str(), &res.x, &res.y, &channels))
		{
			Logger::FILESYS->error("Failed to read image header of file {}", filename);
			return Math::Vector2i{ -1, -1 };
		}
		return res;
	}
}