/*
 *  Copyright (c) 2017-2019 Collabora Ltd.
 *  Copyright (c) 2024 mittorn
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sub license, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice (including the
 * next paragraph) shall be included in all copies or substantial portions
 * of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT.
 * IN NO EVENT SHALL PRECISION INSIGHT AND/OR ITS SUPPLIERS BE LIABLE FOR
 * ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */
#ifndef VULKAN_TEXTURE_UTL_H
#define VULKAN_TEXTURE_UTL_H
#include "vulkan_utl.h"
struct VulkanTexture
{
	// keep device for destroy
	VkDevice owning_device = NULL;
	VkImage image = NULL;
	VkImageLayout image_layout = VK_IMAGE_LAYOUT_UNDEFINED;
	VkDeviceMemory device_memory = NULL;
	VkSampler sampler = NULL;
	VkImageView view = NULL;

	uint32_t width = 0, height = 0;
	uint32_t mip_levels = 0;
	uint32_t layer_count = 0;

	bool created_from_image = false;
	static VkAccessFlags GetAccessFlags(VkImageLayout layout)
	{
		switch (layout) {
			case VK_IMAGE_LAYOUT_UNDEFINED:
			case VK_IMAGE_LAYOUT_PRESENT_SRC_KHR:
			case VK_IMAGE_LAYOUT_GENERAL: return 0;
			case VK_IMAGE_LAYOUT_PREINITIALIZED: return VK_ACCESS_HOST_WRITE_BIT;
			case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
				return VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
			case VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL:
			return VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
			case VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL: return VK_ACCESS_TRANSFER_READ_BIT;
			case VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL:
				return VK_ACCESS_TRANSFER_WRITE_BIT;
			case VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL:
				return VK_ACCESS_SHADER_READ_BIT;
			default: printf("Unhandled access mask case for layout %d.\n", layout);
		}
		return 0;
	}

	static void SetImageLayout(VkCommandBuffer cmd_buffer, VkImage image, VkImageLayout src_layout, VkImageLayout dst_layout, VkImageSubresourceRange subresource_range)
	{
		VkImageMemoryBarrier imageMemoryBarrier = {
			.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
			.srcAccessMask = GetAccessFlags(src_layout),
			.dstAccessMask = GetAccessFlags(dst_layout),
			.oldLayout = src_layout,
			.newLayout = dst_layout,
			.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
			.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
			.image = image,
			.subresourceRange = subresource_range,
		};

		vkCmdPipelineBarrier(cmd_buffer, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
							 VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, NULL, 0, NULL,
							 1, &imageMemoryBarrier);
	}

	VkDescriptorImageInfo GetDescriptor()
	{
		VkDescriptorImageInfo descriptor = {
		  .sampler = sampler,
		  .imageView = view,
		  .imageLayout = image_layout,
		};
		return descriptor;
	}

	void Destroy()
	{
		if (view)
		  vkDestroyImageView(owning_device, view, NULL);
		view = NULL;
		if (!created_from_image && image)
		  vkDestroyImage(owning_device, image, NULL);
		image = NULL;
		if (sampler)
		  vkDestroySampler(owning_device, sampler, NULL);
		sampler = NULL;
		if (device_memory)
		  vkFreeMemory(owning_device, device_memory, NULL);
		device_memory = 0;
	}

	void CreateImage(VkDevice device, VkFormat format)
	{
		VkImageCreateInfo info = {
			.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
			.imageType = VK_IMAGE_TYPE_2D,
			.format = format,
			.extent = { .width = width, .height = height, .depth = 1 },
			.mipLevels = mip_levels,
			.arrayLayers = 1,
			.samples = VK_SAMPLE_COUNT_1_BIT,
			.tiling = VK_IMAGE_TILING_OPTIMAL,
			.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT,
			.sharingMode = VK_SHARING_MODE_EXCLUSIVE,
			.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED
		};
		owning_device = device;

		VK_CHECK_RESULT(vkCreateImage(device, &info, NULL, &image));

	}

	static void AllocateImageMemory(VkImage image,
							 VulkanDevice dev,
							 VkMemoryRequirements *mem_reqs,
							 VkMemoryAllocateInfo *mem_info,
							 VkDeviceMemory *out_memory)
	{
		vkGetImageMemoryRequirements(dev.device, image, mem_reqs);

		mem_info->allocationSize = mem_reqs->size;

		dev.GetMemoryType(mem_reqs->memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, &mem_info->memoryTypeIndex);
		VK_CHECK_RESULT(vkAllocateMemory(dev.device, mem_info, NULL, out_memory));
		VK_CHECK_RESULT(vkBindImageMemory(dev.device, image, *out_memory, 0));
	}

	void TransferImage(VulkanDevice &dev, VkQueue copy_queue, VkBuffer staging_buffer, VkImageLayout dest_layout, VkBufferImageCopy *buffer_image_copies)
	{
		VkImageSubresourceRange subresource_range = {
			.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
			.baseMipLevel = 0,
			.levelCount = mip_levels,
			.layerCount = 1,
		};

		VkCommandBuffer copy_cmd = dev.CreateCommandBuffer();
		SetImageLayout(copy_cmd, image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, subresource_range);

		vkCmdCopyBufferToImage(copy_cmd, staging_buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
							   mip_levels, buffer_image_copies);

		image_layout = dest_layout;
		SetImageLayout(copy_cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, dest_layout, subresource_range);

		dev.FlushCommandBuffer(copy_cmd, copy_queue);
	}

	uint8_t *BeginUpload(VulkanDevice &dev, VkMemoryRequirements &mem_reqs, VkMemoryAllocateInfo &mem_info, VkBuffer &staging_buffer, VkDeviceMemory &staging_memory, VkDeviceSize size)
	{
		VkBufferCreateInfo bufferCreateInfo = {
			.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
			.size = size,
			.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
			.sharingMode = VK_SHARING_MODE_EXCLUSIVE
		};
		VK_CHECK_RESULT(vkCreateBuffer(dev.device, &bufferCreateInfo, NULL, &staging_buffer));

		vkGetBufferMemoryRequirements(dev.device, staging_buffer, &mem_reqs);

		uint32_t type_index;
		dev.GetMemoryType(mem_reqs.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
						  VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, &type_index);

		mem_info = VkMemoryAllocateInfo{
			.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
			.allocationSize = mem_reqs.size,
			.memoryTypeIndex = type_index,
		};

		VK_CHECK_RESULT(vkAllocateMemory(dev.device, &mem_info, NULL, &staging_memory));
		VK_CHECK_RESULT(vkBindBufferMemory(dev.device, staging_buffer, staging_memory, 0));

		uint8_t *data;
		VK_CHECK_RESULT(vkMapMemory(dev.device, staging_memory, 0, mem_reqs.size, 0, (void **)&data));
		return data;
	}

	void EndUpload(VulkanDevice &dev, VkQueue copy_queue, bool alloc_mem, VkImageLayout dest_layout, VkMemoryRequirements &mem_reqs, VkMemoryAllocateInfo &mem_info, VkBuffer &staging_buffer, VkDeviceMemory &staging_memory, unsigned int depth, unsigned int *mip_offsets)
	{

		VkBufferImageCopy buffer_image_copies[mip_levels];
		vkUnmapMemory(dev.device, staging_memory);
		unsigned int offset = 0;
		if(!mip_offsets)
			mip_offsets = &offset;
		for (uint32_t i = 0; i < mip_levels; i++) {
			buffer_image_copies[i] = VkBufferImageCopy {
				.bufferOffset = mip_offsets[i],
				.imageSubresource =
				{
					.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
					.mipLevel = i,
					.baseArrayLayer = 0,
					.layerCount = 1,
				},
				.imageExtent = {
					.width = width >> i,
					.height = height >> i,
					.depth = depth  >> i,
				},
			};
		}

		if (alloc_mem)
			AllocateImageMemory(image, dev, &mem_reqs, &mem_info, &device_memory);
		TransferImage(dev, copy_queue, staging_buffer, dest_layout, buffer_image_copies);
		vkFreeMemory(dev.device, staging_memory, NULL);
		vkDestroyBuffer(dev.device, staging_buffer, NULL);
	}

	void FromBuffer(VulkanDevice &dev, int w, int h, int mips, VkFormat format, VkQueue copy_queue, VkImageLayout dest_layout, uint8_t *buffer, VkDeviceSize size)
	{
		VkMemoryRequirements mem_reqs;
		VkMemoryAllocateInfo mem_info;
		VkBuffer staging_buffer;
		VkDeviceMemory staging_memory;
		width = w;
		height = h;
		mip_levels = mips;
		CreateImage(dev.device, format);
		uint8_t *data = BeginUpload(dev, mem_reqs, mem_info,staging_buffer, staging_memory, size);
		memcpy(data, buffer, size);
		EndUpload(dev, copy_queue, true, dest_layout, mem_reqs, mem_info, staging_buffer, staging_memory, 1, NULL);
		CreateSampler(dev.device);
		CreateImageView(dev.device, format);
		created_from_image = false;
	}

	void CreateSampler(VkDevice device)
	{
		VkSamplerCreateInfo info = {
			.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
			.magFilter = VK_FILTER_LINEAR,
			.minFilter = VK_FILTER_LINEAR,
			.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR,
			.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT,
			.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT,
			.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT,
			.mipLodBias = 0.0f,
			.anisotropyEnable = VK_TRUE,
			.maxAnisotropy = 8,
			.compareOp = VK_COMPARE_OP_NEVER,
			.minLod = 0.0f,
			.maxLod = (float)mip_levels,
			.borderColor = VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE,
		};
		owning_device = device;

		VK_CHECK_RESULT(vkCreateSampler(device, &info, NULL, &sampler));

	}
	void CreateImageView(VkDevice device, VkFormat format)
	{
		VkImageViewCreateInfo info = {
			.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
			.image = image,
			.viewType = VK_IMAGE_VIEW_TYPE_2D,
			.format = format,
			.components = {
				.r = VK_COMPONENT_SWIZZLE_R,
				.g = VK_COMPONENT_SWIZZLE_G,
				.b = VK_COMPONENT_SWIZZLE_B,
				.a = VK_COMPONENT_SWIZZLE_A,
			},
			.subresourceRange = {
				.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
				.baseMipLevel = 0,
				.levelCount = mip_levels,
				.baseArrayLayer = 0,
				.layerCount = 1,
			},
		};
		owning_device = device;
		VK_CHECK_RESULT(vkCreateImageView(device, &info, NULL, &view));
	}
};

#endif // VULKAN_TEXTURE_UTL_H