#include "VulkanTexture.hpp"

#include "Logger.h"
#include "VulkanBuffers.h"
#include "VulkanSwapChain.h"

#include <stb_image.h>

VulkanTexture::VulkanTexture(
	VkDevice		 Device,
	VkPhysicalDevice PhysicalDevice,
	VkQueue			 GraphicsQueue,
	VulkanBuffers&	 Buffers,
	VulkanSwapChain& SwapChain
)
	: Device(Device)
	, PhysicalDevice(PhysicalDevice)
	, GraphicsQueue(GraphicsQueue)
	, Buffers(Buffers)
	, SwapChain(SwapChain)
{
}

VulkanTexture::~VulkanTexture() {}

void VulkanTexture::Cleanup()
{
	CleanupTextures();
	CleanupDepth();
}

void VulkanTexture::CleanupTextures()
{
	vkDestroySampler(Device, Sampler, nullptr);
	vkDestroyImageView(Device, ImageView, nullptr);
	vkDestroyImage(Device, Image, nullptr);
	vkFreeMemory(Device, ImageMemory, nullptr);
}

void VulkanTexture::CleanupDepth()
{
	vkDestroyImageView(Device, DepthImageView, nullptr);
	vkDestroyImage(Device, DepthImage, nullptr);
	vkFreeMemory(Device, DepthImageMemory, nullptr);
}

void VulkanTexture::LoadFromFile(const char* Filename)
{
	int		 TextureWidth, TextureHeight, TextureChannels;
	stbi_uc* Pixels =
		stbi_load(Filename, &TextureWidth, &TextureHeight, &TextureChannels, STBI_rgb_alpha);
	VkDeviceSize ImageSize = TextureWidth * TextureHeight * 4;

	if (!Pixels)
	{
		Log::Error("Failed to load texture image!");
	}
	else
	{
		Log::Info("Successfully loaded texture image.");
	}

	VkBuffer	   StagingBuffer;
	VkDeviceMemory StagingBufferMemory;

	Buffers.CreateBuffer(
		ImageSize,
		VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
		VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
		StagingBuffer,
		StagingBufferMemory
	);

	void* Data;
	vkMapMemory(Device, StagingBufferMemory, 0, ImageSize, 0, &Data);
	memcpy(Data, Pixels, static_cast<size_t>(ImageSize));
	vkUnmapMemory(Device, StagingBufferMemory);

	stbi_image_free(Pixels);

	CreateImage(
		TextureWidth,
		TextureHeight,
		VK_FORMAT_R8G8B8A8_SRGB,
		VK_IMAGE_TILING_OPTIMAL,
		VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
		VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		Image,
		ImageMemory
	);

	Buffers.TransitionImageLayout(
		Image,
		VK_FORMAT_R8G8B8A8_SRGB,
		VK_IMAGE_LAYOUT_UNDEFINED,
		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		HasStencilComponent(VK_FORMAT_R8G8B8A8_SRGB)
	);

	Buffers.CopyBufferToImage(
		StagingBuffer,
		Image,
		static_cast<uint32_t>(TextureWidth),
		static_cast<uint32_t>(TextureHeight)
	);

	Buffers.TransitionImageLayout(
		Image,
		VK_FORMAT_R8G8B8A8_SRGB,
		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
		HasStencilComponent(VK_FORMAT_R8G8B8A8_SRGB)
	);

	vkDestroyBuffer(Device, StagingBuffer, nullptr);
	vkFreeMemory(Device, StagingBufferMemory, nullptr);
}

void VulkanTexture::CreateImageView()
{
	ImageView =
		SwapChain.CreateImageView(Image, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_ASPECT_COLOR_BIT);
}

void VulkanTexture::CreateSampler()
{
	VkSamplerCreateInfo SamplerInfo{};
	SamplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
	SamplerInfo.magFilter = VK_FILTER_LINEAR;
	SamplerInfo.minFilter = VK_FILTER_LINEAR;
	SamplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
	SamplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
	SamplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
	SamplerInfo.anisotropyEnable = VK_TRUE;

	VkPhysicalDeviceProperties Properties{};
	vkGetPhysicalDeviceProperties(PhysicalDevice, &Properties);

	SamplerInfo.maxAnisotropy = Properties.limits.maxSamplerAnisotropy;
	SamplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
	SamplerInfo.unnormalizedCoordinates = VK_FALSE;
	SamplerInfo.compareEnable = VK_FALSE;
	SamplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;
	SamplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
	SamplerInfo.mipLodBias = 0.0f;
	SamplerInfo.minLod = 0.0f;
	SamplerInfo.maxLod = 0.0f;

	if (vkCreateSampler(Device, &SamplerInfo, nullptr, &Sampler) != VK_SUCCESS)
	{
		Log::Error("Failed to create texture sampler!");
	}
	else
	{
		Log::Info("Successfully created texture sampler.");
	}
}

void VulkanTexture::CreateDepthResources()
{
	VkFormat   DepthFormat = FindDepthFormat();
	VkExtent2D Extents = SwapChain.GetSwapChainExtent();

	CreateImage(
		Extents.width,
		Extents.height,
		DepthFormat,
		VK_IMAGE_TILING_OPTIMAL,
		VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
		VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
		DepthImage,
		DepthImageMemory
	);

	DepthImageView = SwapChain.CreateImageView(DepthImage, DepthFormat, VK_IMAGE_ASPECT_DEPTH_BIT);

	Buffers.TransitionImageLayout(
		DepthImage,
		DepthFormat,
		VK_IMAGE_LAYOUT_UNDEFINED,
		VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL,
		HasStencilComponent(DepthFormat)
	);
}

VkFormat VulkanTexture::FindSupportedFormat(
	const std::vector<VkFormat>& Candidates,
	VkImageTiling				 Tiling,
	VkFormatFeatureFlags		 Features
)
{
	for (VkFormat Format : Candidates)
	{
		VkFormatProperties Properties;
		vkGetPhysicalDeviceFormatProperties(PhysicalDevice, Format, &Properties);

		if (Tiling == VK_IMAGE_TILING_LINEAR
			&& (Properties.linearTilingFeatures & Features) == Features)
		{
			return Format;
		}
		else if (Tiling == VK_IMAGE_TILING_OPTIMAL
				 && (Properties.optimalTilingFeatures & Features) == Features)
		{
			return Format;
		}
	}

	Log::Error("Failed to find supported format!");
}

VkFormat VulkanTexture::FindDepthFormat()
{
	return FindSupportedFormat(
		{ VK_FORMAT_D32_SFLOAT, VK_FORMAT_D32_SFLOAT_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT },
		VK_IMAGE_TILING_OPTIMAL,
		VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT
	);
}

bool VulkanTexture::HasStencilComponent(VkFormat Format)
{
	return Format == VK_FORMAT_D32_SFLOAT_S8_UINT || Format == VK_FORMAT_D24_UNORM_S8_UINT;
}

void VulkanTexture::CreateImage(
	uint32_t			  Width,
	uint32_t			  Height,
	VkFormat			  Format,
	VkImageTiling		  Tiling,
	VkImageUsageFlags	  Usage,
	VkMemoryPropertyFlags Properties,
	VkImage&			  Image,
	VkDeviceMemory&		  ImageMemory
)
{

	VkImageCreateInfo ImageInfo{};
	ImageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
	ImageInfo.imageType = VK_IMAGE_TYPE_2D;
	ImageInfo.extent.width = Width;
	ImageInfo.extent.height = Height;
	ImageInfo.extent.depth = 1;
	ImageInfo.mipLevels = 1;
	ImageInfo.arrayLayers = 1;
	ImageInfo.format = Format;
	ImageInfo.tiling = Tiling;
	ImageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
	ImageInfo.usage = Usage;
	ImageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
	ImageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	if (vkCreateImage(Device, &ImageInfo, nullptr, &Image) != VK_SUCCESS)
	{
		Log::Error("Failed to create image!");
	}
	else
	{
		Log::Info("Successfully created image.");
	}

	VkMemoryRequirements MemoryRequirements;
	vkGetImageMemoryRequirements(Device, Image, &MemoryRequirements);

	VkMemoryAllocateInfo AllocateInfo{};
	AllocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	AllocateInfo.allocationSize = MemoryRequirements.size;
	AllocateInfo.memoryTypeIndex =
		Buffers.FindMemoryType(MemoryRequirements.memoryTypeBits, Properties);

	if (vkAllocateMemory(Device, &AllocateInfo, nullptr, &ImageMemory) != VK_SUCCESS)
	{
		Log::Error("Failed to allocate image memory!");
	}
	else
	{
		Log::Info("Successfully allocated imaged memory.");
	}

	vkBindImageMemory(Device, Image, ImageMemory, 0);
}
