#include <fstream>
#include <destrum/Graphics/Init.h>
#include <destrum/Graphics/Util.h>

#include <volk.h>

#include "spdlog/spdlog.h"

namespace {
    VkImageAspectFlags AspectForLayout(VkImageLayout currentLayout, VkImageLayout newLayout)
    {
        return (currentLayout == VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL ||
                newLayout == VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL ||
                newLayout == VK_IMAGE_LAYOUT_DEPTH_READ_ONLY_STENCIL_ATTACHMENT_OPTIMAL)
            ? VK_IMAGE_ASPECT_DEPTH_BIT
            : VK_IMAGE_ASPECT_COLOR_BIT;
    }

    struct LayoutSync {
        VkPipelineStageFlags2 stage{VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT};
        VkAccessFlags2 access{VK_ACCESS_2_MEMORY_READ_BIT | VK_ACCESS_2_MEMORY_WRITE_BIT};
    };

    LayoutSync SyncForLayout(VkImageLayout layout, bool source)
    {
        switch (layout) {
        case VK_IMAGE_LAYOUT_UNDEFINED:
            return {VK_PIPELINE_STAGE_2_TOP_OF_PIPE_BIT, 0};
        case VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL:
            return {VK_PIPELINE_STAGE_2_TRANSFER_BIT, VK_ACCESS_2_TRANSFER_READ_BIT};
        case VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL:
            return {VK_PIPELINE_STAGE_2_TRANSFER_BIT, VK_ACCESS_2_TRANSFER_WRITE_BIT};
        case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
            return {
                VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
                VK_ACCESS_2_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT};
        case VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL:
            return {
                VK_PIPELINE_STAGE_2_EARLY_FRAGMENT_TESTS_BIT |
                    VK_PIPELINE_STAGE_2_LATE_FRAGMENT_TESTS_BIT,
                VK_ACCESS_2_DEPTH_STENCIL_ATTACHMENT_READ_BIT |
                    VK_ACCESS_2_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT};
        case VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL:
        case VK_IMAGE_LAYOUT_READ_ONLY_OPTIMAL:
            return {VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT, VK_ACCESS_2_SHADER_READ_BIT};
        case VK_IMAGE_LAYOUT_PRESENT_SRC_KHR:
            return {VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT, VK_ACCESS_2_MEMORY_READ_BIT};
        default:
            return source
                ? LayoutSync{}
                : LayoutSync{VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
                              VK_ACCESS_2_MEMORY_READ_BIT | VK_ACCESS_2_MEMORY_WRITE_BIT};
        }
    }
}

void vkutil::transitionImage(
    VkCommandBuffer cmd,
    VkImage image,
    VkImageLayout currentLayout,
    VkImageLayout newLayout)
{
    transitionImage(
        cmd,
        image,
        currentLayout,
        newLayout,
        vkinit::imageSubresourceRange(AspectForLayout(currentLayout, newLayout)));
}

void vkutil::transitionImage(
    VkCommandBuffer cmd,
    VkImage image,
    VkImageLayout currentLayout,
    VkImageLayout newLayout,
    const VkImageSubresourceRange& subresourceRange)
{
    const LayoutSync sourceSync = SyncForLayout(currentLayout, true);
    const LayoutSync destinationSync = SyncForLayout(newLayout, false);
    VkImageMemoryBarrier2 imageBarrier{
        .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
        .srcStageMask = sourceSync.stage,
        .srcAccessMask = sourceSync.access,
        .dstStageMask = destinationSync.stage,
        .dstAccessMask = destinationSync.access,
        .oldLayout = currentLayout,
        .newLayout = newLayout,
        .image = image,
        .subresourceRange = subresourceRange,
    };

    VkDependencyInfo depInfo{
        .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &imageBarrier,
    };

    vkCmdPipelineBarrier2(cmd, &depInfo);
}

void vkutil::copyImageToImage(VkCommandBuffer cmd, VkImage source, VkImage destination, VkExtent2D srcSize, VkExtent2D dstSize, VkFilter filter) {
    copyImageToImage(cmd, source, destination, srcSize, 0, 0, dstSize.width, dstSize.height, filter);
}

void vkutil::copyImageToImage(VkCommandBuffer cmd, VkImage source, VkImage destination, VkExtent2D srcSize, int destX, int destY, int destW, int destH, VkFilter filter) {
    const auto blitRegion = VkImageBlit2{
        .sType = VK_STRUCTURE_TYPE_IMAGE_BLIT_2,
        .srcSubresource =
        {
            .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
            .mipLevel = 0,
            .baseArrayLayer = 0,
            .layerCount = 1,
        },
        .srcOffsets =
        {
            {},
            {(std::int32_t)srcSize.width, (std::int32_t)srcSize.height, 1},
        },
        .dstSubresource =
        {
            .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
            .mipLevel = 0,
            .baseArrayLayer = 0,
            .layerCount = 1,
        },
        .dstOffsets =
        {
            {(std::int32_t)destX, (std::int32_t)destY},
            {(std::int32_t)(destX + destW), (std::int32_t)(destY + destH), 1},
        },
    };

    const auto blitInfo = VkBlitImageInfo2{
        .sType = VK_STRUCTURE_TYPE_BLIT_IMAGE_INFO_2,
        .srcImage = source,
        .srcImageLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
        .dstImage = destination,
        .dstImageLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        .regionCount = 1,
        .pRegions = &blitRegion,
        .filter = filter,
    };

    vkCmdBlitImage2(cmd, &blitInfo);
}

void vkutil::bufferHostWriteToShaderReadBarrier(
    VkCommandBuffer cmd, VkBuffer buffer, VkDeviceSize offset, VkDeviceSize size) {
    VkBufferMemoryBarrier2 barrier{VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2};
    barrier.srcStageMask = VK_PIPELINE_STAGE_2_HOST_BIT;
    barrier.srcAccessMask = VK_ACCESS_2_HOST_WRITE_BIT;
    barrier.dstStageMask = VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT |
                           VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT |
                           VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT;
    barrier.dstAccessMask = VK_ACCESS_2_SHADER_READ_BIT;
    barrier.buffer = buffer;
    barrier.offset = offset;
    barrier.size = size;

    VkDependencyInfo dep{VK_STRUCTURE_TYPE_DEPENDENCY_INFO};
    dep.bufferMemoryBarrierCount = 1;
    dep.pBufferMemoryBarriers = &barrier;

    vkCmdPipelineBarrier2(cmd, &dep);
}

void vkutil::bufferHostWriteToTransferReadBarrier(
    VkCommandBuffer cmd,
    VkBuffer buffer,
    VkDeviceSize offset,
    VkDeviceSize size)
{
    VkBufferMemoryBarrier2 barrier{VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER_2};
    barrier.srcStageMask = VK_PIPELINE_STAGE_2_HOST_BIT;
    barrier.srcAccessMask = VK_ACCESS_2_HOST_WRITE_BIT;
    barrier.dstStageMask = VK_PIPELINE_STAGE_2_TRANSFER_BIT;
    barrier.dstAccessMask = VK_ACCESS_2_TRANSFER_READ_BIT;
    barrier.buffer = buffer;
    barrier.offset = offset;
    barrier.size = size;

    VkDependencyInfo dep{VK_STRUCTURE_TYPE_DEPENDENCY_INFO};
    dep.bufferMemoryBarrierCount = 1;
    dep.pBufferMemoryBarriers = &barrier;

    vkCmdPipelineBarrier2(cmd, &dep);
}


void vkutil::addDebugLabel(VkDevice device, VkImage image, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_IMAGE,
        .objectHandle = (std::uint64_t)image,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void vkutil::addDebugLabel(VkDevice device, VkShaderModule shaderModule, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_SHADER_MODULE,
        .objectHandle = (std::uint64_t)shaderModule,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void vkutil::addDebugLabel(VkDevice device, VkPipeline pipeline, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_PIPELINE,
        .objectHandle = (std::uint64_t)pipeline,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void vkutil::addDebugLabel(VkDevice device, VkBuffer buffer, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_BUFFER,
        .objectHandle = (std::uint64_t)buffer,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void vkutil::addDebugLabel(VkDevice device, VkSampler sampler, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_SAMPLER,
        .objectHandle = (std::uint64_t)sampler,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

vkutil::RenderInfo vkutil::createRenderingInfo(const RenderingInfoParams& params) {
    assert(
        (params.colorImageView || params.depthImageView != nullptr) &&
        "Either draw or depth image should be present");
    assert(
        params.renderExtent.width != 0.f && params.renderExtent.height != 0.f &&
        "renderExtent not specified");

    RenderInfo ri;
    if (params.colorImageView) {
        ri.colorAttachment = VkRenderingAttachmentInfo{
            .sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
            .imageView = params.colorImageView,
            .imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
            .loadOp = params.colorImageClearValue ? VK_ATTACHMENT_LOAD_OP_CLEAR : VK_ATTACHMENT_LOAD_OP_LOAD,
            .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        };
        if (params.colorImageClearValue) {
            const auto col = params.colorImageClearValue.value();
            ri.colorAttachment.clearValue.color = {col[0], col[1], col[2], col[3]};
        }
    }

    if (params.depthImageView) {
        ri.depthAttachment = VkRenderingAttachmentInfo{
            .sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
            .imageView = params.depthImageView,
            .imageLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL,
            .loadOp = params.depthImageClearValue ? VK_ATTACHMENT_LOAD_OP_CLEAR : VK_ATTACHMENT_LOAD_OP_LOAD,
            .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        };
        if (params.depthImageClearValue) {
            ri.depthAttachment.clearValue.depthStencil.depth = params.depthImageClearValue.value();
        }
    }

    ri.renderingInfo = VkRenderingInfo{
        .sType = VK_STRUCTURE_TYPE_RENDERING_INFO,
        .renderArea =
        VkRect2D{
            .offset = {},
            .extent = params.renderExtent,
        },
        .layerCount = 1,
        .colorAttachmentCount = params.colorImageView ? 1u : 0u,
        .pColorAttachments = params.colorImageView ? &ri.colorAttachment : nullptr,
        .pDepthAttachment = params.depthImageView ? &ri.depthAttachment : nullptr,
    };

    return ri;
}

VkShaderModule vkutil::loadShaderModule(const std::filesystem::path& path, VkDevice device) {
    std::ifstream file(path, std::ios::ate | std::ios::binary);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open shader: " + path.string());
    }

    const auto fileSize = file.tellg();
    if (fileSize < 0 || fileSize % static_cast<std::streamoff>(sizeof(std::uint32_t)) != 0) {
        throw std::runtime_error("Shader file has an invalid size: " + path.string());
    }
    std::vector<std::uint32_t> buffer(fileSize / sizeof(std::uint32_t));

    file.seekg(0);
    if (!file.read(reinterpret_cast<char*>(buffer.data()), fileSize)) {
        throw std::runtime_error("Failed to read shader: " + path.string());
    }
    file.close();

    auto info = VkShaderModuleCreateInfo{
        .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
        .codeSize = buffer.size() * sizeof(std::uint32_t),
        .pCode = buffer.data(),
    };

    VkShaderModule shaderModule{VK_NULL_HANDLE};
    VK_CHECK(vkCreateShaderModule(device, &info, nullptr, &shaderModule));
    return shaderModule;
}

void addDebugLabel(VkDevice device, VkImage image, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_IMAGE,
        .objectHandle = (std::uint64_t)image,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void addDebugLabel(VkDevice device, VkImageView imageView, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_IMAGE_VIEW,
        .objectHandle = (std::uint64_t)imageView,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void addDebugLabel(VkDevice device, VkShaderModule shaderModule, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_SHADER_MODULE,
        .objectHandle = (std::uint64_t)shaderModule,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void addDebugLabel(VkDevice device, VkPipeline pipeline, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_PIPELINE,
        .objectHandle = (std::uint64_t)pipeline,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void addDebugLabel(VkDevice device, VkPipelineLayout layout, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_PIPELINE_LAYOUT,
        .objectHandle = (std::uint64_t)layout,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void addDebugLabel(VkDevice device, VkBuffer buffer, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_BUFFER,
        .objectHandle = (std::uint64_t)buffer,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

void addDebugLabel(VkDevice device, VkSampler sampler, const char* label) {
    const auto nameInfo = VkDebugUtilsObjectNameInfoEXT{
        .sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT,
        .objectType = VK_OBJECT_TYPE_SAMPLER,
        .objectHandle = (std::uint64_t)sampler,
        .pObjectName = label,
    };
    if (vkSetDebugUtilsObjectNameEXT) vkSetDebugUtilsObjectNameEXT(device, &nameInfo);
}

vkutil::RenderInfo createRenderingInfo(const vkutil::RenderingInfoParams& params) {
    return vkutil::createRenderingInfo(params);
}