/*
 * 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 "DataFormat.hpp"
#include "Base/Logger.hpp"
#include <magic_enum.hpp>
#include "Math/Math.hpp"

namespace OpenVulkano
{
	namespace
	{
		std::unordered_map<DataFormat::Format, uint32_t> MakeFormatSizeMap()
		{
			std::unordered_map<DataFormat::Format, uint32_t> map;
			for (auto format: magic_enum::enum_values<DataFormat::Format>())
			{
				uint32_t size = 0;
				if (format == DataFormat::R4G4_UNORM_PACK8 || format == DataFormat::S8_UINT) size = 1;
				else if (format >= DataFormat::R4G4B4A4_UNORM_PACK16 && format <= DataFormat::A1R5G5B5_UNORM_PACK16)
					size = 2;
				else if (format >= DataFormat::R8_UNORM && format <= DataFormat::R8_SRGB) size = 1;
				else if (format >= DataFormat::R8G8_UNORM && format <= DataFormat::R8G8_SRGB) size = 2;
				else if (format >= DataFormat::R8G8B8_UNORM && format <= DataFormat::B8G8R8_SRGB) size = 3;
				else if (format >= DataFormat::R8G8B8A8_UNORM && format <= DataFormat::A2B10G10R10_SINT_PACK32)
					size = 4;
				else if (format >= DataFormat::R16_UNORM && format <= DataFormat::R16_SFLOAT) size = 2;
				else if (format >= DataFormat::R16G16_UNORM && format <= DataFormat::R16G16_SFLOAT) size = 4;
				else if (format >= DataFormat::R16G16B16_UNORM && format <= DataFormat::R16G16B16_SFLOAT) size = 6;
				else if (format >= DataFormat::R16G16B16A16_UNORM && format <= DataFormat::R16G16B16A16_SFLOAT)
					size = 8;
				else if (format >= DataFormat::R32_UINT && format <= DataFormat::R32_SFLOAT) size = 4;
				else if (format >= DataFormat::R32G32_UINT && format <= DataFormat::R32G32_SFLOAT) size = 8;
				else if (format >= DataFormat::R32G32B32_UINT && format <= DataFormat::R32G32B32_SFLOAT) size = 12;
				else if (format >= DataFormat::R32G32B32A32_UINT && format <= DataFormat::R32G32B32A32_SFLOAT)
					size = 16;
				else if (format >= DataFormat::R64_UINT && format <= DataFormat::R64_SFLOAT) size = 8;
				else if (format >= DataFormat::R64G64_UINT && format <= DataFormat::R64G64_SFLOAT) size = 16;
				else if (format >= DataFormat::R64G64B64_UINT && format <= DataFormat::R64G64B64_SFLOAT) size = 24;
				else if (format >= DataFormat::R64G64B64A64_UINT && format <= DataFormat::R64G64B64A64_SFLOAT)
					size = 32;
				else if (format == DataFormat::B10G11R11_UFLOAT_PACK32 || format == DataFormat::E5B9G9R9_UFLOAT_PACK32
				         || format == DataFormat::X8_D24_UNORM_PACK32 || format == DataFormat::D32_SFLOAT
				         || format == DataFormat::D24_UNORM_S8_UINT)
					size = 4;
				else if (format == DataFormat::D16_UNORM) size = 2;
				else if (format == DataFormat::D16_UNORM_S8_UINT) size = 3;
				else if (format == DataFormat::D32_SFLOAT_S8_UINT) size = 5;
				else if (format == DataFormat::BC1_RGB_UNORM_BLOCK || format == DataFormat::BC1_RGB_SRGB_BLOCK
				         || format == DataFormat::BC1_RGBA_UNORM_BLOCK || format == DataFormat::BC1_RGBA_SRGB_BLOCK
				         || format == DataFormat::BC4_SNORM_BLOCK || format == DataFormat::BC4_UNORM_BLOCK)
					size = 8;
				else if (format == DataFormat::BC2_SRGB_BLOCK || format == DataFormat::BC2_UNORM_BLOCK
				         || format == DataFormat::BC3_SRGB_BLOCK || format == DataFormat::BC3_UNORM_BLOCK
				         || format == DataFormat::BC5_SNORM_BLOCK || format == DataFormat::BC5_UNORM_BLOCK
				         || format == DataFormat::BC6H_SFLOAT_BLOCK || format == DataFormat::BC6H_UFLOAT_BLOCK
				         || format == DataFormat::BC7_SRGB_BLOCK || format == DataFormat::BC7_UNORM_BLOCK)
					size = 16;
				else if (format == DataFormat::ETC2_R8G8B8_UNORM_BLOCK || format == DataFormat::ETC2_R8G8B8_SRGB_BLOCK
				         || format == DataFormat::ETC2_R8G8B8A1_UNORM_BLOCK
				         || format == DataFormat::ETC2_R8G8B8A1_SRGB_BLOCK
				         || format == DataFormat::ETC2_R8G8B8A8_UNORM_BLOCK
				         || format == DataFormat::ETC2_R8G8B8A8_SRGB_BLOCK || format == DataFormat::EAC_R11_UNORM_BLOCK
				         || format == DataFormat::EAC_R11_SNORM_BLOCK || format == DataFormat::EAC_R11G11_UNORM_BLOCK
				         || format == DataFormat::EAC_R11G11_SNORM_BLOCK)
					size = 16;
				else if (format == DataFormat::ASTC_4x4_UNORM_BLOCK || format == DataFormat::ASTC_4x4_SFLOAT_BLOCK
				         || format == DataFormat::ASTC_4x4_SRGB_BLOCK || format == DataFormat::ASTC_5x4_UNORM_BLOCK
				         || format == DataFormat::ASTC_5x4_SFLOAT_BLOCK || format == DataFormat::ASTC_5x4_SRGB_BLOCK
				         || format == DataFormat::ASTC_5x5_UNORM_BLOCK || format == DataFormat::ASTC_5x5_SFLOAT_BLOCK
				         || format == DataFormat::ASTC_5x5_SRGB_BLOCK || format == DataFormat::ASTC_6x5_UNORM_BLOCK
				         || format == DataFormat::ASTC_6x5_SFLOAT_BLOCK || format == DataFormat::ASTC_6x5_SRGB_BLOCK
				         || format == DataFormat::ASTC_6x6_UNORM_BLOCK || format == DataFormat::ASTC_6x6_SFLOAT_BLOCK
				         || format == DataFormat::ASTC_6x6_SRGB_BLOCK || format == DataFormat::ASTC_8x5_UNORM_BLOCK
				         || format == DataFormat::ASTC_8x5_SFLOAT_BLOCK || format == DataFormat::ASTC_8x5_SRGB_BLOCK
				         || format == DataFormat::ASTC_8x6_UNORM_BLOCK || format == DataFormat::ASTC_8x6_SFLOAT_BLOCK
				         || format == DataFormat::ASTC_8x6_SRGB_BLOCK || format == DataFormat::ASTC_8x8_UNORM_BLOCK
				         || format == DataFormat::ASTC_8x8_SFLOAT_BLOCK || format == DataFormat::ASTC_8x8_SRGB_BLOCK
				         || format == DataFormat::ASTC_10x5_UNORM_BLOCK || format == DataFormat::ASTC_10x5_SFLOAT_BLOCK
				         || format == DataFormat::ASTC_10x5_SRGB_BLOCK || format == DataFormat::ASTC_10x6_UNORM_BLOCK
				         || format == DataFormat::ASTC_10x6_SFLOAT_BLOCK || format == DataFormat::ASTC_10x6_SRGB_BLOCK
				         || format == DataFormat::ASTC_10x8_UNORM_BLOCK || format == DataFormat::ASTC_10x8_SFLOAT_BLOCK
				         || format == DataFormat::ASTC_10x8_SRGB_BLOCK || format == DataFormat::ASTC_10x10_UNORM_BLOCK
				         || format == DataFormat::ASTC_10x10_SFLOAT_BLOCK || format == DataFormat::ASTC_10x10_SRGB_BLOCK
				         || format == DataFormat::ASTC_12x10_UNORM_BLOCK
				         || format == DataFormat::ASTC_12x10_SFLOAT_BLOCK || format == DataFormat::ASTC_12x10_SRGB_BLOCK
				         || format == DataFormat::ASTC_12x12_UNORM_BLOCK || format == DataFormat::ASTC_12x12_SRGB_BLOCK
				         || format == DataFormat::ASTC_12x12_SFLOAT_BLOCK)
					size = 16;
				else if (format == DataFormat::PVRTC1_2BPP_UNORM_BLOCK_IMG
				         || format == DataFormat::PVRTC1_2BPP_SRGB_BLOCK_IMG
				         || format == DataFormat::PVRTC1_4BPP_SRGB_BLOCK_IMG
				         || format == DataFormat::PVRTC1_4BPP_UNORM_BLOCK_IMG
				         || format == DataFormat::PVRTC2_2BPP_SRGB_BLOCK_IMG
				         || format == DataFormat::PVRTC2_2BPP_UNORM_BLOCK_IMG
				         || format == DataFormat::PVRTC2_4BPP_SRGB_BLOCK_IMG
				         || format == DataFormat::PVRTC2_4BPP_UNORM_BLOCK_IMG)
					size = 8;
				else if (format == DataFormat::B10X6G10X6R10X6G10X6_422_UNORM_4PACK16
				         || format == DataFormat::G10X6B10X6G10X6R10X6_422_UNORM_4PACK16
				         || format == DataFormat::R10X6G10X6B10X6A10X6_UNORM_4PACK16
				         || format == DataFormat::R12X4G12X4B12X4A12X4_UNORM_4PACK16
				         || format == DataFormat::G12X4B12X4G12X4R12X4_422_UNORM_4PACK16
				         || format == DataFormat::B12X4G12X4R12X4G12X4_422_UNORM_4PACK16
				         || format == DataFormat::G16B16G16R16_422_UNORM
				         || format == DataFormat::B16G16R16G16_422_UNORM)
					size = 8;
				else if (format == DataFormat::G10X6_B10X6_R10X6_3PLANE_420_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6R10X6_2PLANE_420_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6_R10X6_3PLANE_422_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6R10X6_2PLANE_422_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6_R10X6_3PLANE_444_UNORM_3PACK16)
					size = 6;

				map[format] = size;
			}
			return map;
		}

		std::unordered_map<DataFormat::Format, Math::Vector2uc> MakeBlockSizeMap()
		{
			std::unordered_map<DataFormat::Format, Math::Vector2uc> map;

			for (auto format: magic_enum::enum_values<DataFormat::Format>())
			{
				Math::Vector2ui size = { 1, 1 };
				if (format == DataFormat::BC1_RGB_UNORM_BLOCK || format == DataFormat::BC1_RGB_SRGB_BLOCK
				    || format == DataFormat::BC1_RGBA_UNORM_BLOCK || format == DataFormat::BC1_RGBA_SRGB_BLOCK
				    || format == DataFormat::BC2_UNORM_BLOCK || format == DataFormat::BC2_SRGB_BLOCK
				    || format == DataFormat::BC3_UNORM_BLOCK || format == DataFormat::BC3_SRGB_BLOCK
				    || format == DataFormat::BC4_UNORM_BLOCK || format == DataFormat::BC4_SNORM_BLOCK
				    || format == DataFormat::BC5_UNORM_BLOCK || format == DataFormat::BC5_SNORM_BLOCK
				    || format == DataFormat::BC6H_UFLOAT_BLOCK || format == DataFormat::BC6H_SFLOAT_BLOCK
				    || format == DataFormat::BC7_UNORM_BLOCK || format == DataFormat::BC7_SRGB_BLOCK
				    || format == DataFormat::ETC2_R8G8B8_UNORM_BLOCK || format == DataFormat::ETC2_R8G8B8_SRGB_BLOCK
				    || format == DataFormat::ETC2_R8G8B8A1_UNORM_BLOCK || format == DataFormat::ETC2_R8G8B8A1_SRGB_BLOCK
				    || format == DataFormat::ETC2_R8G8B8A8_UNORM_BLOCK || format == DataFormat::ETC2_R8G8B8A8_SRGB_BLOCK
				    || format == DataFormat::EAC_R11_UNORM_BLOCK || format == DataFormat::EAC_R11G11_UNORM_BLOCK
				    || format == DataFormat::EAC_R11G11_SNORM_BLOCK || format == DataFormat::ASTC_4x4_UNORM_BLOCK
				    || format == DataFormat::ASTC_4x4_SRGB_BLOCK || format == DataFormat::PVRTC1_4BPP_UNORM_BLOCK_IMG
				    || format == DataFormat::PVRTC2_4BPP_UNORM_BLOCK_IMG
				    || format == DataFormat::PVRTC1_4BPP_SRGB_BLOCK_IMG || format == DataFormat::ASTC_4x4_SFLOAT_BLOCK)
					size = { 4, 4 };
				else if (format == DataFormat::ASTC_5x4_UNORM_BLOCK || format == DataFormat::ASTC_5x4_SRGB_BLOCK)
					size = { 5, 4 };
				else if (format == DataFormat::ASTC_5x5_UNORM_BLOCK || format == DataFormat::ASTC_5x5_SRGB_BLOCK)
					size = { 5, 5 };
				else if (format == DataFormat::ASTC_6x5_UNORM_BLOCK || format == DataFormat::ASTC_6x5_SRGB_BLOCK)
					size = { 6, 5 };
				else if (format == DataFormat::ASTC_6x6_UNORM_BLOCK || format == DataFormat::ASTC_6x6_SRGB_BLOCK)
					size = { 6, 6 };
				else if (format == DataFormat::ASTC_8x5_UNORM_BLOCK || format == DataFormat::ASTC_8x5_SRGB_BLOCK)
					size = { 8, 5 };
				else if (format == DataFormat::ASTC_8x6_UNORM_BLOCK || format == DataFormat::ASTC_8x6_SRGB_BLOCK)
					size = { 8, 6 };
				else if (format == DataFormat::ASTC_8x8_UNORM_BLOCK || format == DataFormat::ASTC_8x8_SRGB_BLOCK)
					size = { 8, 8 };
				else if (format == DataFormat::ASTC_10x5_UNORM_BLOCK || format == DataFormat::ASTC_10x5_SRGB_BLOCK)
					size = { 10, 5 };
				else if (format == DataFormat::ASTC_10x6_UNORM_BLOCK || format == DataFormat::ASTC_10x6_SRGB_BLOCK)
					size = { 10, 6 };
				else if (format == DataFormat::ASTC_10x8_UNORM_BLOCK || format == DataFormat::ASTC_10x8_SRGB_BLOCK)
					size = { 10, 8 };
				else if (format == DataFormat::ASTC_10x10_UNORM_BLOCK || format == DataFormat::ASTC_10x10_SRGB_BLOCK)
					size = { 10, 10 };
				else if (format == DataFormat::ASTC_12x10_UNORM_BLOCK || format == DataFormat::ASTC_12x10_SRGB_BLOCK)
					size = { 12, 10 };
				else if (format == DataFormat::ASTC_12x12_UNORM_BLOCK || format == DataFormat::ASTC_12x12_SRGB_BLOCK)
					size = { 12, 12 };
				else if (format == DataFormat::G8B8G8R8_422_UNORM || format == DataFormat::B8G8R8G8_422_UNORM
				         || format == DataFormat::G8_B8_R8_3PLANE_422_UNORM
				         || format == DataFormat::G8_B8R8_2PLANE_422_UNORM
				         || format == DataFormat::G10X6B10X6G10X6R10X6_422_UNORM_4PACK16
				         || format == DataFormat::B10X6G10X6R10X6G10X6_422_UNORM_4PACK16
				         || format == DataFormat::G10X6_B10X6_R10X6_3PLANE_422_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6R10X6_2PLANE_422_UNORM_3PACK16
				         || format == DataFormat::G12X4B12X4G12X4R12X4_422_UNORM_4PACK16
				         || format == DataFormat::B12X4G12X4R12X4G12X4_422_UNORM_4PACK16
				         || format == DataFormat::G12X4_B12X4_R12X4_3PLANE_422_UNORM_3PACK16
				         || format == DataFormat::G12X4_B12X4R12X4_2PLANE_422_UNORM_3PACK16
				         || format == DataFormat::G16B16G16R16_422_UNORM || format == DataFormat::B16G16R16G16_422_UNORM
				         || format == DataFormat::G16_B16_R16_3PLANE_422_UNORM
				         || format == DataFormat::G16_B16R16_2PLANE_422_UNORM)
					size = { 2, 1 };
				else if (format == DataFormat::G8_B8_R8_3PLANE_420_UNORM
				         || format == DataFormat::G8_B8R8_2PLANE_420_UNORM
				         || format == DataFormat::G10X6_B10X6_R10X6_3PLANE_420_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6R10X6_2PLANE_420_UNORM_3PACK16
				         || format == DataFormat::G12X4_B12X4_R12X4_3PLANE_420_UNORM_3PACK16
				         || format == DataFormat::G12X4_B12X4R12X4_2PLANE_420_UNORM_3PACK16
				         || format == DataFormat::G16_B16_R16_3PLANE_420_UNORM
				         || format == DataFormat::G16_B16R16_2PLANE_420_UNORM)
					size = { 2, 2 };
				else if (format == DataFormat::G8_B8_R8_3PLANE_444_UNORM
				         || format == DataFormat::G8_B8R8_2PLANE_444_UNORM
				         || format == DataFormat::G10X6_B10X6_R10X6_3PLANE_444_UNORM_3PACK16
				         || format == DataFormat::G10X6_B10X6R10X6_2PLANE_444_UNORM_3PACK16
				         || format == DataFormat::G12X4_B12X4_R12X4_3PLANE_444_UNORM_3PACK16
				         || format == DataFormat::G12X4_B12X4R12X4_2PLANE_444_UNORM_3PACK16
				         || format == DataFormat::G16_B16_R16_3PLANE_444_UNORM
				         || format == DataFormat::G16_B16R16_2PLANE_444_UNORM)
					size = { 2, 2 };
				else if (format == DataFormat::PVRTC2_2BPP_UNORM_BLOCK_IMG
				         || format == DataFormat::PVRTC2_2BPP_SRGB_BLOCK_IMG)
					size = { 8, 4 };

				map[format] = size;
			}
			return map;
		}

		const std::unordered_map<DataFormat::Format, uint32_t> FORMAT_SIZE_MAP = MakeFormatSizeMap();
		const std::unordered_map<DataFormat::Format, Math::Vector2uc> BLOCK_SIZE_MAP = MakeBlockSizeMap();
	}

	std::string_view DataFormat::GetName() const { return magic_enum::enum_name(m_format); }

	DataFormat DataFormat::GetFromName(std::string_view name)
	{
		auto result = magic_enum::enum_cast<Format>(name);
		if (result.has_value()) return { result.value() };
		return { UNDEFINED };
	}

	uint32_t DataFormat::GetBytesPerPixel() const
	{
		auto size = FORMAT_SIZE_MAP.find(m_format);
		if (size == FORMAT_SIZE_MAP.end() || size->second == 0) [[unlikely]]
		{
			Logger::RENDER->warn("Failed to get pixel size for format: {}", GetName());
			return 0;
		}
		return size->second;
	}

	bool DataFormat::IsBlockCompressed() const
	{
		for (auto format: magic_enum::enum_values<DataFormat::Format>())
		{
			if (format == m_format)
			{
				return format >= BC1_RGB_UNORM_BLOCK && format <= ASTC_12x12_SRGB_BLOCK
				       || format >= ASTC_4x4_SFLOAT_BLOCK && format <= PVRTC2_4BPP_SRGB_BLOCK_IMG;
			}
		}
		return false;
	}

	size_t DataFormat::CalculatedSize(uint32_t& width, uint32_t& height) const
	{
		if (IsBlockCompressed())
		{
			Math::Vector2uc size = BLOCK_SIZE_MAP.find(m_format)->second;
			Math::Vector2uc blockCount = { std::ceil((float) width / size.x), std::ceil((float) height / size.y) };

			width = blockCount.x * size.x;
			height = blockCount.y * size.y;

			return width * height * GetBytesPerPixel();
		}

		return width * height * GetBytesPerPixel();
	}

#ifndef __APPLE__
	DataFormat DataFormat::GetFromMetalPixelFormat(int formatId) { return UNDEFINED; }
#endif
}