#pragma once
#ifndef IMAGEVIEW_H
#define	IMAGEVIEW_H

#define STB_IMAGE_IMPLEMENTATION
#include "../include/stb-master/stb_image.h"

#include "VulkanLaunchpad.h"
/*!
 *	Struct used as the return value for loadImage
 */
struct ImageAndView {
    /*! Contains the created image handle */
    VkImage image;

    /*! Contains the created image view handle */
    VkImageView view;
};

VkImageView createImageViewForImage(VkDevice device, VkImage image, VkImageViewType view_type, VkFormat format) {
    /* --------------------------------------------- */
    // Subtask 5.6: Create an Image View for each Image
    /* --------------------------------------------- */
    VkImageViewCreateInfo image_view_create_info = {};
    image_view_create_info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    image_view_create_info.image = image;
    image_view_create_info.viewType = view_type;
    image_view_create_info.format = format;
    image_view_create_info.components.r = VK_COMPONENT_SWIZZLE_R;
    image_view_create_info.components.g = VK_COMPONENT_SWIZZLE_G;
    image_view_create_info.components.b = VK_COMPONENT_SWIZZLE_B;
    image_view_create_info.components.a = VK_COMPONENT_SWIZZLE_A;
    image_view_create_info.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    image_view_create_info.subresourceRange.baseMipLevel = 0u;
    image_view_create_info.subresourceRange.levelCount = VK_REMAINING_MIP_LEVELS;
    image_view_create_info.subresourceRange.baseArrayLayer = 0u;
    image_view_create_info.subresourceRange.layerCount = VK_IMAGE_VIEW_TYPE_CUBE == view_type ? 6u : 1u;

    VkImageView vk_image_view;
    VkResult result = vkCreateImageView(device, &image_view_create_info, nullptr, &vk_image_view);
    VKL_CHECK_VULKAN_RESULT(result);

    return vk_image_view;
}

ImageAndView loadImage(VkDevice device, VkQueue queue, VkCommandPool command_pool, std::string image_file_name, bool g_synchronization2_supported, PFN_vkCmdPipelineBarrier2KHR g_vkCmdPipelineBarrier2KHR) {
    std::vector<std::string> image_file_paths = gcgFindTextureFiles({image_file_name});

    /* --------------------------------------------- */
    // Subtask 5.5: Load DDS Textures into Images
    /* --------------------------------------------- */
    VklImageInfo image_info = vklGetDdsImageInfo(image_file_paths[0].c_str());
    const uint32_t image_layers = static_cast<uint32_t>(image_file_paths.size());
    assert(image_layers == 1u || image_layers == 6u);
    VkImageCreateFlags image_flags = image_layers == 6 ? VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT : VkImageCreateFlags{};
    VkResult result;

    // 0. Create the image:
    VkImage image = vklCreateDeviceLocalImageWithBackingMemory(image_info.extent.width, image_info.extent.height, image_info.imageFormat,
        VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_SAMPLED_BIT, image_layers, image_flags);

    // 1. Create a command buffer and start recording
    VkCommandBufferAllocateInfo command_buffer_allocate_info = {};
    command_buffer_allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    command_buffer_allocate_info.commandPool = command_pool;
    command_buffer_allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    command_buffer_allocate_info.commandBufferCount = 1u;

    VkCommandBuffer command_buffer;
    result = vkAllocateCommandBuffers(device, &command_buffer_allocate_info, &command_buffer);
    VKL_CHECK_VULKAN_RESULT(result);

    VkCommandBufferBeginInfo command_buffer_begin_info = {};
    command_buffer_begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    command_buffer_begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    vkBeginCommandBuffer(command_buffer, &command_buffer_begin_info);

    // Data structures for Synchronization2:
    VkImageMemoryBarrier2KHR image_memory_barrier2 = {}; // <-- Used only in the g_synchronization2_supported case
    image_memory_barrier2.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2_KHR;
    image_memory_barrier2.image = image;
    image_memory_barrier2.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    image_memory_barrier2.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;

    VkDependencyInfoKHR dependency_info = {};
    dependency_info.sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO_KHR;
    dependency_info.imageMemoryBarrierCount = 1u;
    dependency_info.pImageMemoryBarriers = &image_memory_barrier2;

    // Data structures for classical synchronization:
    VkImageMemoryBarrier image_memory_barrier = {}; // <-- Used only in the !g_synchronization2_supported case
    image_memory_barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    image_memory_barrier.image = image;
    image_memory_barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    image_memory_barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;

    /* --------------------------------------------- */
    // Subtask 5.10: Enable Mipmapping
    /* --------------------------------------------- */
    std::vector<VkBuffer> staging_buffers;
    uint32_t num_levels = static_cast<uint32_t>(std::log2(std::max(image_info.extent.width, image_info.extent.height))) + 1u;

    for (uint32_t layer = 0u; layer < image_layers; ++layer) {
        image_info = vklGetDdsImageInfo(image_file_paths[layer].c_str());
        uint32_t mipWidth = image_info.extent.width;
        uint32_t mipHeight = image_info.extent.height;
        for (uint32_t level = 0u; level < num_levels; ++level) {
            image_info = vklGetDdsImageLevelInfo(image_file_paths[layer].c_str(), level);
            // Sanity check to see we are loading an image with the correct size
            if (mipWidth != image_info.extent.width || mipHeight != image_info.extent.height)
                VKL_EXIT_WITH_ERROR("vklGetDdsImageLevelInfo for level " + std::to_string(level) +
                    " returned an image with width=" + std::to_string(image_info.extent.width) + ", height=" + std::to_string(image_info.extent.height) +
                    " instead of the expected width=" + std::to_string(mipWidth) + ", height=" + std::to_string(mipHeight) + ".");
            mipWidth = mipWidth > 1u ? mipWidth / 2u : 1u;
            mipHeight = mipHeight > 1u ? mipHeight / 2u : 1u;
            // Store in vector to keep alive until they are no longer needed, which is after the fence has been signaled:
            staging_buffers.push_back(vklLoadDdsImageLevelIntoHostCoherentBuffer(image_file_paths[layer].c_str(), level));

            if (g_synchronization2_supported) {
                // 2. Record an image layout transition into a format that is optimal for the image to get data copied into it:
                image_memory_barrier2.srcStageMask = VK_PIPELINE_STAGE_2_NONE_KHR;       // No need to wait on anything
                image_memory_barrier2.srcAccessMask = VK_ACCESS_2_NONE_KHR;              // Nothing required thanks to implicit memory guaranteed with host writes
                image_memory_barrier2.dstStageMask = VK_PIPELINE_STAGE_2_COPY_BIT_KHR;   // The subsequent command (which must wait) is a COPY command.
                image_memory_barrier2.dstAccessMask = VK_ACCESS_2_TRANSFER_READ_BIT_KHR; // Copy reads from the buffer. The layout transition must be available
                                                                                         // to that type of access.
                image_memory_barrier2.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
                image_memory_barrier2.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
                image_memory_barrier2.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                image_memory_barrier2.subresourceRange.baseMipLevel = level;
                image_memory_barrier2.subresourceRange.levelCount = 1u;
                image_memory_barrier2.subresourceRange.baseArrayLayer = layer;
                image_memory_barrier2.subresourceRange.layerCount = 1u;
                g_vkCmdPipelineBarrier2KHR(command_buffer, &dependency_info);
            } else { // => Synchronization2 is NOT supported => use oldschool style:
                // 2. Record an image layout transition into a format that is optimal for the image to get data copied into it:
                image_memory_barrier.srcAccessMask = {};
                image_memory_barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                image_memory_barrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
                image_memory_barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
                image_memory_barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                image_memory_barrier.subresourceRange.baseMipLevel = level;
                image_memory_barrier.subresourceRange.levelCount = 1u;
                image_memory_barrier.subresourceRange.baseArrayLayer = layer;
                image_memory_barrier.subresourceRange.layerCount = 1u;
                vkCmdPipelineBarrier(command_buffer, 
                    VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, /* --> */ VK_PIPELINE_STAGE_TRANSFER_BIT, {},
                    /* Memory barriers:        */ 0u, nullptr,
                    /* Buffer memory barriers: */ 0u, nullptr,
                    /* Image memory barriers:  */ 1u, &image_memory_barrier);
            }

            // 3. Record copying data from the (host-coherent) buffer into the (device-local) image:
            VkBufferImageCopy buffer_image_copy = {};
            buffer_image_copy.bufferOffset = 0;
            buffer_image_copy.bufferRowLength = 0u;
            buffer_image_copy.bufferImageHeight = 0u;
            buffer_image_copy.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            buffer_image_copy.imageSubresource.mipLevel = level;
            buffer_image_copy.imageSubresource.baseArrayLayer = layer;
            buffer_image_copy.imageSubresource.layerCount = 1u;
            buffer_image_copy.imageOffset.x = 0u;
            buffer_image_copy.imageOffset.y = 0u;
            buffer_image_copy.imageOffset.z = 0u;
            buffer_image_copy.imageExtent.width = image_info.extent.width;
            buffer_image_copy.imageExtent.height = image_info.extent.height;
            buffer_image_copy.imageExtent.depth = 1u;
            vkCmdCopyBufferToImage(command_buffer, staging_buffers.back(), image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1u, &buffer_image_copy);

            if (g_synchronization2_supported) {
                // 4. Record an image layout transition into a format that is optimal for rendering:
                // Re-use the VkImageMemoryBarrier from above, but modify a few parameters:
                image_memory_barrier2.srcStageMask = VK_PIPELINE_STAGE_2_COPY_BIT_KHR;    // It is the preceding COPY which must have completed before moving on.
                image_memory_barrier2.srcAccessMask = VK_ACCESS_2_TRANSFER_WRITE_BIT_KHR; // The writes of that COPY must have completed.
                image_memory_barrier2.dstStageMask =
                    VK_PIPELINE_STAGE_2_NONE_KHR; // No dst stage required in the barrier, because afterwards we sync with the fence.
                image_memory_barrier2.dstAccessMask = VK_ACCESS_2_NONE_KHR;
                image_memory_barrier2.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
                image_memory_barrier2.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; // Transform into right format for being sampled from in shaders.
                g_vkCmdPipelineBarrier2KHR(command_buffer, &dependency_info);
            } else { // => Synchronization2 is NOT supported => use oldschool style:
                // 4. Record an image layout transition into a format that is optimal for rendering:
                // Re-use the VkImageMemoryBarrier from above, but modify a few parameters:
                image_memory_barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                image_memory_barrier.dstAccessMask = {};
                image_memory_barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
                image_memory_barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;

                vkCmdPipelineBarrier(command_buffer, 
                    VK_PIPELINE_STAGE_TRANSFER_BIT, /* --> */ VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, {},
                    /* Memory barriers:        */ 0u, nullptr,
                    /* Buffer memory barriers: */ 0u, nullptr,
                    /* Image memory barriers:  */ 1u, &image_memory_barrier);
            }
        }
    }

    // 5. End recording and submit to queue; use a fence to wait for the operation to complete:
    vkEndCommandBuffer(command_buffer);

    VkFenceCreateInfo fence_create_info = {};
    fence_create_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;

    VkFence fence;
    result = vkCreateFence(device, &fence_create_info, nullptr, &fence);
    VKL_CHECK_VULKAN_RESULT(result);

    VkSubmitInfo submit_info = {};
    submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submit_info.waitSemaphoreCount = 0u;
    submit_info.commandBufferCount = 1u;
    submit_info.pCommandBuffers = &command_buffer;
    submit_info.signalSemaphoreCount = 0u;
    result = vkQueueSubmit(queue, 1u, &submit_info, fence);
    VKL_CHECK_VULKAN_RESULT(result);

    // 6. Wait for the operation to complete, then clean up:
    result = vkWaitForFences(device, 1u, &fence, VK_TRUE, UINT64_MAX);
    VKL_CHECK_VULKAN_RESULT(result);

    vkDestroyFence(device, fence, nullptr);
    vkFreeCommandBuffers(device, command_pool, 1u, &command_buffer);
    for (auto staging_buffer : staging_buffers) {
        vklDestroyHostCoherentBufferAndItsBackingMemory(staging_buffer);
    }

    return ImageAndView{
        image, createImageViewForImage(device, image, image_layers == 6 ? VK_IMAGE_VIEW_TYPE_CUBE : VK_IMAGE_VIEW_TYPE_2D, image_info.imageFormat)
    };
}







#endif

