/*
 * 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 "ImageLoaderDds.hpp"

#include <dds.hpp>
#include <fmt/format.h>
#include <magic_enum.hpp>

#include <unordered_map>

namespace
{
	std::unordered_map<DXGI_FORMAT, OpenVulkano::DataFormat> DXGI_TO_OV_DATAFORMAT = {
		{ DXGI_FORMAT_R32G32B32A32_FLOAT, OpenVulkano::DataFormat::Format::R32G32B32A32_SFLOAT },
		{ DXGI_FORMAT_R32G32B32A32_UINT, OpenVulkano::DataFormat::Format::R32G32B32A32_UINT },
		{ DXGI_FORMAT_R32G32B32A32_SINT, OpenVulkano::DataFormat::Format::R32G32B32A32_SINT },
		{ DXGI_FORMAT_R32G32B32_FLOAT, OpenVulkano::DataFormat::Format::R32G32B32_SFLOAT },
		{ DXGI_FORMAT_R32G32B32_UINT, OpenVulkano::DataFormat::Format::R32G32B32_UINT },
		{ DXGI_FORMAT_R32G32B32_SINT, OpenVulkano::DataFormat::Format::R32G32B32_SINT },

		{ DXGI_FORMAT_R16G16B16A16_FLOAT, OpenVulkano::DataFormat::Format::R16G16B16A16_SFLOAT },
		{ DXGI_FORMAT_R16G16B16A16_UNORM, OpenVulkano::DataFormat::Format::R16G16B16A16_UNORM },
		{ DXGI_FORMAT_R16G16B16A16_UINT, OpenVulkano::DataFormat::Format::R16G16B16A16_UINT },
		{ DXGI_FORMAT_R16G16B16A16_SNORM, OpenVulkano::DataFormat::Format::R16G16B16A16_SNORM },
		{ DXGI_FORMAT_R16G16B16A16_SINT, OpenVulkano::DataFormat::Format::R16G16B16A16_SINT },

		{ DXGI_FORMAT_R32G32_FLOAT, OpenVulkano::DataFormat::Format::R32G32_SFLOAT },
		{ DXGI_FORMAT_R32G32_UINT, OpenVulkano::DataFormat::Format::R32G32_UINT },
		{ DXGI_FORMAT_R32G32_SINT, OpenVulkano::DataFormat::Format::R32G32_SINT },

		{ DXGI_FORMAT_D32_FLOAT_S8X24_UINT, OpenVulkano::DataFormat::Format::D32_SFLOAT_S8_UINT },

		{ DXGI_FORMAT_R8G8B8A8_UNORM, OpenVulkano::DataFormat::Format::R8G8B8A8_UNORM },
		{ DXGI_FORMAT_R8G8B8A8_UNORM_SRGB, OpenVulkano::DataFormat::Format::R8G8B8A8_SRGB },
		{ DXGI_FORMAT_R8G8B8A8_UINT, OpenVulkano::DataFormat::Format::R8G8B8A8_UINT },
		{ DXGI_FORMAT_R8G8B8A8_SNORM, OpenVulkano::DataFormat::Format::R8G8B8A8_SNORM },
		{ DXGI_FORMAT_R8G8B8A8_SINT, OpenVulkano::DataFormat::Format::R8G8B8A8_SINT },

		{ DXGI_FORMAT_R16G16_FLOAT, OpenVulkano::DataFormat::Format::R16G16_SFLOAT },
		{ DXGI_FORMAT_R16G16_UNORM, OpenVulkano::DataFormat::Format::R16G16_UNORM },
		{ DXGI_FORMAT_R16G16_UINT, OpenVulkano::DataFormat::Format::R16G16_UINT },
		{ DXGI_FORMAT_R16G16_SNORM, OpenVulkano::DataFormat::Format::R16G16_SNORM },
		{ DXGI_FORMAT_R16G16_SINT, OpenVulkano::DataFormat::Format::R16G16_SINT },

		{ DXGI_FORMAT_D32_FLOAT, OpenVulkano::DataFormat::Format::D32_SFLOAT },

		{ DXGI_FORMAT_R32_FLOAT, OpenVulkano::DataFormat::Format::R32_SFLOAT },
		{ DXGI_FORMAT_R32_UINT, OpenVulkano::DataFormat::Format::R32_UINT },
		{ DXGI_FORMAT_R32_SINT, OpenVulkano::DataFormat::Format::R32_SINT },

		{ DXGI_FORMAT_D24_UNORM_S8_UINT, OpenVulkano::DataFormat::Format::D24_UNORM_S8_UINT },

		{ DXGI_FORMAT_R8G8_UNORM, OpenVulkano::DataFormat::Format::R8G8_UNORM },
		{ DXGI_FORMAT_R8G8_UINT, OpenVulkano::DataFormat::Format::R8G8_UINT },
		{ DXGI_FORMAT_R8G8_SNORM, OpenVulkano::DataFormat::Format::R8G8_SNORM },
		{ DXGI_FORMAT_R8G8_SINT, OpenVulkano::DataFormat::Format::R8G8_SINT },

		{ DXGI_FORMAT_D16_UNORM, OpenVulkano::DataFormat::Format::D16_UNORM },
		{ DXGI_FORMAT_R16_FLOAT, OpenVulkano::DataFormat::Format::R16_SFLOAT },
		{ DXGI_FORMAT_R16_UNORM, OpenVulkano::DataFormat::Format::R16_UNORM },
		{ DXGI_FORMAT_R16_UINT, OpenVulkano::DataFormat::Format::R16_UINT },
		{ DXGI_FORMAT_R16_SNORM, OpenVulkano::DataFormat::Format::R16_SNORM },
		{ DXGI_FORMAT_R16_SINT, OpenVulkano::DataFormat::Format::R16_SINT },

		{ DXGI_FORMAT_R8_UNORM, OpenVulkano::DataFormat::Format::R8_UNORM },
		{ DXGI_FORMAT_R8_UINT, OpenVulkano::DataFormat::Format::R8_UINT },
		{ DXGI_FORMAT_R8_SNORM, OpenVulkano::DataFormat::Format::R8_SNORM },
		{ DXGI_FORMAT_R8_SINT, OpenVulkano::DataFormat::Format::R8_SINT },

		{ DXGI_FORMAT_BC1_UNORM, OpenVulkano::DataFormat::Format::BC1_RGBA_UNORM_BLOCK },
		{ DXGI_FORMAT_BC1_UNORM_SRGB, OpenVulkano::DataFormat::Format::BC1_RGBA_SRGB_BLOCK },

		{ DXGI_FORMAT_BC2_UNORM, OpenVulkano::DataFormat::Format::BC2_UNORM_BLOCK },
		{ DXGI_FORMAT_BC2_UNORM_SRGB, OpenVulkano::DataFormat::Format::BC2_SRGB_BLOCK },

		{ DXGI_FORMAT_BC3_UNORM, OpenVulkano::DataFormat::Format::BC3_UNORM_BLOCK },
		{ DXGI_FORMAT_BC3_UNORM_SRGB, OpenVulkano::DataFormat::Format::BC3_SRGB_BLOCK },

		{ DXGI_FORMAT_BC4_UNORM, OpenVulkano::DataFormat::Format::BC4_UNORM_BLOCK },
		{ DXGI_FORMAT_BC4_SNORM, OpenVulkano::DataFormat::Format::BC4_SNORM_BLOCK },

		{ DXGI_FORMAT_BC5_UNORM, OpenVulkano::DataFormat::Format::BC5_UNORM_BLOCK },
		{ DXGI_FORMAT_BC5_SNORM, OpenVulkano::DataFormat::Format::BC5_SNORM_BLOCK },

		{ DXGI_FORMAT_B5G6R5_UNORM, OpenVulkano::DataFormat::Format::B5G6R5_UNORM_PACK16 },
		{ DXGI_FORMAT_B5G5R5A1_UNORM, OpenVulkano::DataFormat::Format::B5G5R5A1_UNORM_PACK16 },

		{ DXGI_FORMAT_B8G8R8A8_UNORM, OpenVulkano::DataFormat::Format::B8G8R8A8_UNORM },
		{ DXGI_FORMAT_B8G8R8X8_UNORM, OpenVulkano::DataFormat::Format::B8G8R8A8_SNORM },

		{ DXGI_FORMAT_B8G8R8A8_UNORM_SRGB, OpenVulkano::DataFormat::Format::B8G8R8A8_SRGB },
		{ DXGI_FORMAT_B8G8R8X8_UNORM_SRGB, OpenVulkano::DataFormat::Format::B8G8R8A8_SRGB },

		{ DXGI_FORMAT_BC6H_UF16, OpenVulkano::DataFormat::Format::BC6H_UFLOAT_BLOCK },
		{ DXGI_FORMAT_BC6H_SF16, OpenVulkano::DataFormat::Format::BC6H_SFLOAT_BLOCK },

		{ DXGI_FORMAT_BC7_UNORM, OpenVulkano::DataFormat::Format::BC7_UNORM_BLOCK },
		{ DXGI_FORMAT_BC7_UNORM_SRGB, OpenVulkano::DataFormat::Format::BC7_SRGB_BLOCK },

		{ DXGI_FORMAT_B4G4R4A4_UNORM, OpenVulkano::DataFormat::Format::B4G4R4A4_UNORM_PACK16 },
	};
}

namespace OpenVulkano::Image
{
	std::unique_ptr<Image> ImageLoaderDds::loadFromFile(const std::string& filePath)
	{
		dds::Image image;
		if (dds::readFile(filePath, &image) != dds::ReadResult::Success)
		{
			throw std::runtime_error("Failed to load DDS texture");
		}
		return ExtractImage(&image);
	}

	std::unique_ptr<Image> ImageLoaderDds::loadFromMemory(const std::vector<uint8_t>& buffer)
	{
		dds::Image image;
		if (dds::readImage(const_cast<uint8_t*>(buffer.data()), buffer.size(), &image) != dds::ReadResult::Success)
		{
			throw std::runtime_error("Failed to load DDS texture");
		}
		return ExtractImage(&image);
	}

	Math::Vector2i ImageLoaderDds::GetImageDimensions(const std::string& filename)
	{
		dds::Image image;
		if (dds::readFile(filename, &image) != dds::ReadResult::Success)
		{
			throw std::runtime_error("Failed to load DDS texture");
		}
		
		Math::Vector2i result;
		result.x = image.width;
		result.y = image.height;
		
		return result;
	}

	std::unique_ptr<Image> ImageLoaderDds::ExtractImage(dds::Image* ddsImage)
	{
		auto image = std::make_unique<Image>();
		
		image->resolution.x = ddsImage->width;
		image->resolution.y = ddsImage->height;
		image->resolution.z = 1;
		
		if (DXGI_TO_OV_DATAFORMAT.contains(ddsImage->format))
		{
			image->dataFormat = DXGI_TO_OV_DATAFORMAT[ddsImage->format];
		}
		else
		{
			throw std::runtime_error(fmt::format("Unhandled DDS texture format: {} ({})",
			                                     magic_enum::enum_name(ddsImage->format),
			                                     static_cast<int>(ddsImage->format)));
		}
		
		if (ddsImage->mipmaps.size() == 0)
		{
			throw std::runtime_error("Failed to read an image from a texture with no mipmaps");
		}

		dds::span texture = ddsImage->mipmaps[0];
		image->data = Array<uint8_t>(texture.size_bytes());
		memcpy(image->data.Data(), texture.data(), texture.size_bytes());

		return image;
	}
}