#pragma once
#include <vulkan/vulkan.h>
#include <GLFW/glfw3.h>
#include <GLM/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <chrono>
#include <vector>
#include <array>
#include <iostream>
#include <stdexcept>
#include <cstdlib>
#include <optional> 
#include <set>
#include <cstdint> // Necessary for uint32_t
#include <limits> // Necessary for std::numeric_limits
#include <algorithm> // Necessary for std::clamp
#include <fstream>

struct TextureData
{
    VkImage textureImage;
    VkDeviceMemory textureImageMemory;
    VkImageView textureImageView;
    VkSampler textureSampler;
    VkDevice device;

    TextureData(VkDevice& vk_device) : device(vk_device)
    {
    }

    void destroyVulkanAssets()
    {
        vkDestroySampler(device, textureSampler, nullptr);
        vkDestroyImageView(device, textureImageView, nullptr);
        vkDestroyImage(device, textureImage, nullptr);
        vkFreeMemory(device, textureImageMemory, nullptr);
    }
};

class TextureLoader
{
public:
    VkPhysicalDevice m_physicalDevice; 
    VkDevice m_device;
    VkQueue m_graphicsQueue;
    VkCommandPool m_commandPool;


    TextureLoader(VkPhysicalDevice physicalDevice, VkDevice device, VkQueue graphicsQueue, VkCommandPool commandPool) :
        m_physicalDevice(physicalDevice),
        m_device(device),
        m_graphicsQueue(graphicsQueue),
        m_commandPool(commandPool)
    {}

    void load(std::string filePath, TextureData& data)
    {
        createTextureImage(filePath, data.textureImage, data.textureImageMemory);
        createTextureImageView(data.textureImageView, data.textureImage);
        createTextureSampler(data.textureSampler);
    }



    void createTextureImage(const std::string& filePath, VkImage& in_textureImage,  VkDeviceMemory& in_textureImageMemory)
    {
        // Load image with library in pixels object
        /*
        The STBI_rgb_alpha value forces the image to be loaded with an alpha channel, even if it doesn't have one
        The pointer that is returned is the first element in an array of pixel values.
        The pixels are laid out row by row with 4 bytes per pixel in the case of STBI_rgb_alpha for a total of texWidth* texHeight * 4 values.
        */
        int texWidth, texHeight, texChannels;
        stbi_uc* pixels = stbi_load(filePath.c_str(), &texWidth, &texHeight, &texChannels, STBI_rgb_alpha);
        VkDeviceSize imageSize = texWidth * texHeight * 4;
        if (!pixels) throw std::runtime_error("failed to load texture image!");


        // create a staging buffer in host visible memory so that we can use vkMapMemory and copy the pixels to it.
        VkBuffer stagingBuffer;
        VkDeviceMemory stagingBufferMemory;
        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(m_device, stagingBufferMemory, 0, imageSize, 0, &data);
        memcpy(data, pixels, static_cast<size_t>(imageSize));
        vkUnmapMemory(m_device, stagingBufferMemory);

        // Clean up STBI storage
        stbi_image_free(pixels);

        // Allocates and binds memory accessed by m_textureImage
        createImage(texWidth, texHeight, 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, in_textureImage, in_textureImageMemory);

        // Configure Image Layout
        transitionImageLayout(in_textureImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

        //copy data from staging buffer to image
        copyBufferToImage(stagingBuffer, in_textureImage, static_cast<uint32_t>(texWidth), static_cast<uint32_t>(texHeight));

        /*To be able to start sampling from the texture image in the shader, we need one last transition to prepare it for shader access */
        transitionImageLayout(in_textureImage, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        vkDestroyBuffer(m_device, stagingBuffer, nullptr);
        vkFreeMemory(m_device, stagingBufferMemory, nullptr);
    }

        // Helper for createTextureImage() to create and bind buffers    
        void createBuffer(VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags properties, VkBuffer& buffer, VkDeviceMemory& bufferMemory) {
            VkBufferCreateInfo bufferInfo{};
            bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
            bufferInfo.size = size;
            bufferInfo.usage = usage;
            bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
            if (vkCreateBuffer(m_device, &bufferInfo, nullptr, &buffer) != VK_SUCCESS)
                throw std::runtime_error("failed to create buffer!");

            // Get memory requirements for the vertex buffer. 
            VkMemoryRequirements memRequirements;
            vkGetBufferMemoryRequirements(m_device, buffer, &memRequirements);

            // Allocate memory to buffer
            /*
            properties = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT mean
            We use host coherent memory to ensure data is instantly visible in the buffer, this appraoch is a bit slower though
            */
            VkMemoryAllocateInfo allocInfo{};
            allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
            allocInfo.allocationSize = memRequirements.size;
            allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);

            if (vkAllocateMemory(m_device, &allocInfo, nullptr, &bufferMemory) != VK_SUCCESS) {
                throw std::runtime_error("failed to allocate buffer memory!");
            }
            // Last argument is offset, must be multiple of memRequirements.alignment
            vkBindBufferMemory(m_device, buffer, bufferMemory, 0);
        }

            //Helper for createBuffer to find the right type of memory for our requirements and hardware's requirements
            uint32_t findMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties)
            {
                VkPhysicalDeviceMemoryProperties memProperties; // Has two attributes: memoryTypes (types in every heap) and memoryHeaps (E.g. VRAM and RAM)
                vkGetPhysicalDeviceMemoryProperties(m_physicalDevice, &memProperties);

                // Check 1) if memory matches typeFilter and 2) if memory is visible and mapable from CPU
                for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++)
                    if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties)
                        return i;

                throw std::runtime_error("failed to find suitable memory type!");
            }

        // Helper for createTextureImage() Configure image, create, allocate memory and bind to handle
        void 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;

            // One dimensional images can be used to store an array of data or gradient, two dimensional images are mainly used for textures, 
            // and three dimensional images can be used to store voxel volumes,
            imageInfo.imageType = VK_IMAGE_TYPE_2D;

            imageInfo.extent.width = width;
            imageInfo.extent.height = height;
            imageInfo.extent.depth = 1;
            imageInfo.mipLevels = 1;
            imageInfo.arrayLayers = 1;

            // use the same format for the texels as the pixels in the buffer, otherwise the copy operation will fail.
            imageInfo.format = format;

            // Cannot be changed later, VK_IMAGE_TILING_LINEAR allows to acces pixels directly, 
            // VK_IMAGE_TILING_OPTIMAL is faster but no acces to pixels >> can be used when staginf buffer exists.
            imageInfo.tiling = tiling;

            // VK_IMAGE_LAYOUT_UNDEFINED: Not usable by the GPU and the very first transition will discard the texels.
            // VK_IMAGE_LAYOUT_PREINITIALIZED: Not usable by the GPU, but the first transition will preserve the texels >> only needed when used as staging image
            imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            imageInfo.usage = usage;
            imageInfo.samples = VK_SAMPLE_COUNT_1_BIT; // >> related to multi sampling
            imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE; // Image only used by one queue family!
            imageInfo.flags = 0; // Optional for sparse images that do not need the whole memory

            if (vkCreateImage(m_device, &imageInfo, nullptr, &image) != VK_SUCCESS)
                throw std::runtime_error("failed to create image!");

            VkMemoryRequirements memRequirements;
            vkGetImageMemoryRequirements(m_device, image, &memRequirements);

            VkMemoryAllocateInfo allocInfo{};
            allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
            allocInfo.allocationSize = memRequirements.size;
            allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);

            if (vkAllocateMemory(m_device, &allocInfo, nullptr, &imageMemory) != VK_SUCCESS)
                throw std::runtime_error("failed to allocate image memory!");

            vkBindImageMemory(m_device, image, imageMemory, 0);
        }

        // Helper for createTextureImage() Bring image into correct layout using barriers
        void transitionImageLayout(VkImage image, VkFormat format, VkImageLayout oldLayout, VkImageLayout newLayout)
        {
            VkCommandBuffer commandBuffer = beginSingleTimeCommands();
            // Use image memory barrier:
            /*
            One of the most common ways to perform layout transitions is using an image memory barrier.
            A pipeline barrier like that is generally used to synchronize access to resources, like ensuring that a write to a buffer completes before reading from it,
            but it can also be used to transition image layouts and transfer queue family ownership when VK_SHARING_MODE_EXCLUSIVE is used.
            There is an equivalent buffer memory barrier to do this for buffers.
            */
            VkImageMemoryBarrier barrier{};
            barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
            barrier.oldLayout = oldLayout;
            barrier.newLayout = newLayout;
            barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; //If you are using the barrier to transfer queue family ownership, then these two fields should be the indices of the queue families.
            barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barrier.image = image;
            barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            barrier.subresourceRange.baseMipLevel = 0;
            barrier.subresourceRange.levelCount = 1;
            barrier.subresourceRange.baseArrayLayer = 0;
            barrier.subresourceRange.layerCount = 1;

            //Handle Transitions:
            /*
                Undefined → transfer destination : transfer writes that don't need to wait on anything
                Transfer destination → shader reading : shader reads should wait on transfer writes, specifically the shader reads in the fragment shader, because that's where we're going to use the texture
                These rules are specified using the following access masks and pipeline stages
            */
            VkPipelineStageFlags sourceStage;
            VkPipelineStageFlags destinationStage;
            if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
                barrier.srcAccessMask = 0;
                barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;

                sourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
                destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
            }
            else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
                barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

                sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                destinationStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
            }
            else {
                throw std::invalid_argument("unsupported layout transition!");
            }

            /*
            All types of pipeline barriers are submitted using the same function.
            The first parameter after the command buffer specifies in which pipeline stage the operations occur that should happen before the barrier.
            The second parameter specifies the pipeline stage in which operations will wait on the barrier.
            The pipeline stages that you are allowed to specify before and after the barrier depend on how you use the resource before and after the barrier.
            The third parameter is either 0 or VK_DEPENDENCY_BY_REGION_BIT. The latter turns the barrier into a per-region condition.
            That means that the implementation is allowed to already begin reading from the parts of a resource that were written so far, for example.
            The last three pairs of parameters reference arrays of pipeline barriers of the three available types: memory barriers, buffer memory barriers, and image memory barriers like the one we're using here.
            */
            vkCmdPipelineBarrier(
                commandBuffer,
                sourceStage, destinationStage,
                0,
                0, nullptr,
                0, nullptr,
                1, &barrier
            );

            endSingleTimeCommands(commandBuffer);
        }

        // Helper for createTextureImage()
        void copyBufferToImage(VkBuffer buffer, VkImage image, uint32_t width, uint32_t height)
        {
            VkCommandBuffer commandBuffer = beginSingleTimeCommands();

            // Configure image region
            VkBufferImageCopy region{};
            region.bufferOffset = 0; //specifies the byte offset in the buffer at which the pixel values start
            region.bufferRowLength = 0; //specify how the pixels are laid out in memory.  0 for both indicates that the pixels are simply tightly packed 
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT; // Specify which part of image is copied (subresource, Offset, Extent)
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = { 0, 0, 0 };
            region.imageExtent =
            {
                width,
                height,
                1
            };

            // Copy the image: The fourth parameter indicates which layout the image is currently using, 
            // assuming here that the image has already been transitioned to the layout that is optimal for copying pixels to.
            vkCmdCopyBufferToImage
            (
                commandBuffer,
                buffer,
                image,
                VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                1,
                &region
            );

            endSingleTimeCommands(commandBuffer);
        }
    
            // Helper for copyBufferToImage: Record/ end single time Command buffers
            VkCommandBuffer beginSingleTimeCommands()
            {
                VkCommandBufferAllocateInfo allocInfo{};
                allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
                allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
                allocInfo.commandPool = m_commandPool;
                allocInfo.commandBufferCount = 1;
                VkCommandBuffer newCommandBuffer;
                vkAllocateCommandBuffers(m_device, &allocInfo, &newCommandBuffer);

                VkCommandBufferBeginInfo beginInfo{};
                beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
                beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

                vkBeginCommandBuffer(newCommandBuffer, &beginInfo);

                return newCommandBuffer;
            }

            // Helper for copyBufferToImage: Record/ end single time Command buffers
            void endSingleTimeCommands(VkCommandBuffer commandBuffer) {
                vkEndCommandBuffer(commandBuffer);

                VkSubmitInfo submitInfo{};
                submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
                submitInfo.commandBufferCount = 1;
                submitInfo.pCommandBuffers = &commandBuffer;

                vkQueueSubmit(m_graphicsQueue, 1, &submitInfo, VK_NULL_HANDLE);

                // Wait until copying is finished , could also be done with a fence
                vkQueueWaitIdle(m_graphicsQueue);

                vkFreeCommandBuffers(m_device, m_commandPool, 1, &commandBuffer);
            }

    // Creates Image view for textures using helper
    void createTextureImageView(VkImageView& in_textureImageView, VkImage& in_textureImage)
    {
        createImageView(in_textureImage, in_textureImageView, VK_FORMAT_R8G8B8A8_SRGB);
    }

    // General Helper: cretes Image views e.g. for textures
    void createImageView(VkImage& image, VkImageView& in_textureImageView, VkFormat format)
    {
        // Set up create Info for Image view
        VkImageViewCreateInfo viewInfo{};
        viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
        viewInfo.image = image;
        viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
        viewInfo.format = format;

        // Add info about purpose of image view and which part of image is accessed
        viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        viewInfo.subresourceRange.baseMipLevel = 0;
        viewInfo.subresourceRange.levelCount = 1;
        viewInfo.subresourceRange.baseArrayLayer = 0;
        viewInfo.subresourceRange.layerCount = 1;

        //Create ImageView
        if (vkCreateImageView(m_device, &viewInfo, nullptr, &in_textureImageView) != VK_SUCCESS)
            throw std::runtime_error("failed to create texture image view!");
    }

    void createTextureSampler(VkSampler& in_textureSampler)
    {
        // General Sampler Settings
        //addressMode:
        /*
        VK_SAMPLER_ADDRESS_MODE_REPEAT: Repeat the texture when going beyond the image dimensions.
        VK_SAMPLER_ADDRESS_MODE_MIRRORED_REPEAT : Like repeat, but inverts the coordinates to mirror the image when going beyond the dimensions.
        VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE : Take the color of the edge closest to the coordinate beyond the image dimensions.
        VK_SAMPLER_ADDRESS_MODE_MIRROR_CLAMP_TO_EDGE : Like clamp to edge, but instead uses the edge opposite to the closest edge.
        VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER : Return a solid color when sampling beyond the dimensions of the image.
        */
        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_FALSE;
        samplerInfo.unnormalizedCoordinates = VK_FALSE; // use [0, texWidth) and [0, texHeight) or [0,1) coordinates on the texture
        samplerInfo.compareEnable = VK_FALSE; // Mainly for percentage-closer-filtering: Should texels be compared to a value before using smaple?
        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;

        // Anisotropic filtering settings: Get the maximum setting from hardware
        VkPhysicalDeviceProperties properties{};
        vkGetPhysicalDeviceProperties(m_physicalDevice, &properties); // Gets a struct with all options and settings of the hardware
        samplerInfo.maxAnisotropy = properties.limits.maxSamplerAnisotropy;

        // Create Sampler
        if (vkCreateSampler(m_device, &samplerInfo, nullptr, &in_textureSampler) != VK_SUCCESS)
            throw std::runtime_error("failed to create texture sampler!");
    }
};