#include <Vk/CommandBuffer/CommandBuffer.hpp>
#include <Utilities/Utilities.hpp>
#include <vulkan/vulkan_core.h>

Vk::CommandBuffer::CommandBuffer(
    std::shared_ptr<Vk::Device> device,
    std::shared_ptr<Vk::CommandPool> pool,
    VkCommandBufferLevel level
) : m_device(device), m_pool(pool)
{
    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.commandPool = pool->inner();
    allocInfo.level = level;
    allocInfo.commandBufferCount = 1;

    if (vkAllocateCommandBuffers(device->inner(), &allocInfo, &m_commandBuffer) != VK_SUCCESS)
    {
        crash("Failed to allocate command buffers!");
    }
}

VkCommandBuffer Vk::CommandBuffer::inner()
{
    return m_commandBuffer;
}

void Vk::CommandBuffer::reset(VkCommandBufferResetFlags flags)
{
    vkResetCommandBuffer(m_commandBuffer, flags);
    m_waitForSemaphores.clear();
    m_waitStages.clear();
    m_signalSemaphores.clear();
}

void Vk::CommandBuffer::waitFor(Vk::Semaphore &semaphore, VkPipelineStageFlags stages)
{
    m_waitForSemaphores.push_back(semaphore.inner());
    m_waitStages.push_back(stages);
}

void Vk::CommandBuffer::signal(Vk::Semaphore &semaphore)
{
    m_signalSemaphores.push_back(semaphore.inner());
}

void Vk::CommandBuffer::begin(VkCommandBufferUsageFlags usage)
{
    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = usage;

    if (vkBeginCommandBuffer(m_commandBuffer, &beginInfo) != VK_SUCCESS)
    {
        crash("Failed to begin recording command buffer!");
    }
}

void Vk::CommandBuffer::begin(Vk::RenderPass &renderPass, uint32_t subpassIndex, Vk::Framebuffer &framebuffer)
{
    VkCommandBufferInheritanceInfo inheritanceInfo = {};
    inheritanceInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_INHERITANCE_INFO;
    inheritanceInfo.renderPass = renderPass.inner();
    inheritanceInfo.subpass = subpassIndex;
    inheritanceInfo.framebuffer = framebuffer.inner();

    VkCommandBufferBeginInfo beginInfo = {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_RENDER_PASS_CONTINUE_BIT;
    beginInfo.pInheritanceInfo = &inheritanceInfo;

    if (vkBeginCommandBuffer(m_commandBuffer, &beginInfo) != VK_SUCCESS)
    {
        crash("Failed to begin recording command buffer!");
    }
}

void Vk::CommandBuffer::beginRenderPass(
    Vk::RenderPass &renderPass,
    Vk::Framebuffer &framebuffer,
    VkExtent2D renderExtent,
    const std::vector<VkClearValue> &clearValues,
    VkSubpassContents subpassContents
)
{
    VkRenderPassBeginInfo renderPassInfo{};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    renderPassInfo.renderPass = renderPass.inner();
    renderPassInfo.framebuffer = framebuffer.inner();
    renderPassInfo.renderArea.offset = {0, 0};
    renderPassInfo.renderArea.extent = renderExtent;
    renderPassInfo.clearValueCount = clearValues.size();
    renderPassInfo.pClearValues = clearValues.data();

    vkCmdBeginRenderPass(m_commandBuffer, &renderPassInfo, subpassContents);
}

void Vk::CommandBuffer::bindPipeline(VkPipelineBindPoint bindPoint, Vk::Pipeline &pipeline)
{
    vkCmdBindPipeline(m_commandBuffer, bindPoint, pipeline.inner());
}

void Vk::CommandBuffer::bindDescriptorSet(
    VkPipelineBindPoint bindPoint,
    Vk::Pipeline &pipeline,
    uint32_t index,
    Vk::DescriptorSet &descSet
)
{
    VkDescriptorSet innerSet = descSet.inner();
    vkCmdBindDescriptorSets(m_commandBuffer, bindPoint, pipeline.getInnerLayout(), index, 1, &innerSet, 0, nullptr);
}

void Vk::CommandBuffer::setViewport(float x, float y, float width, float height, float minDepth, float maxDepth)
{
    VkViewport viewport{};
    viewport.x = x;
    viewport.y = y;
    viewport.width = width;
    viewport.height = height;
    viewport.minDepth = minDepth;
    viewport.maxDepth = maxDepth;
    vkCmdSetViewport(m_commandBuffer, 0, 1, &viewport);
}

void Vk::CommandBuffer::setScissor(VkRect2D rect)
{
    vkCmdSetScissor(m_commandBuffer, 0, 1, &rect);
};

void Vk::CommandBuffer::pushConstants(
    Vk::Pipeline &pipeline,
    VkShaderStageFlags stages,
    uint32_t offset, uint32_t size,
    const void *data
)
{
    vkCmdPushConstants(m_commandBuffer, pipeline.getInnerLayout(), stages, offset, size, data);
}

void Vk::CommandBuffer::draw(uint32_t vertexCount, uint32_t instanceCount, uint32_t firstVertex, uint32_t firstInstance)
{
    vkCmdDraw(m_commandBuffer, vertexCount, instanceCount, firstVertex, firstInstance);
}

void Vk::CommandBuffer::drawIndexed(
    uint32_t indexCount,
    uint32_t instanceCount,
    uint32_t firstIndex,
    int32_t vertexOffset,
    uint32_t firstInstance
)
{
    vkCmdDrawIndexed(m_commandBuffer, indexCount, instanceCount, firstIndex, vertexOffset, firstInstance);
}

void Vk::CommandBuffer::endRenderPass()
{
    vkCmdEndRenderPass(m_commandBuffer);
}

void Vk::CommandBuffer::transitionBufferToTransferDst(Vk::Buffer &buffer)
{
    VkBufferMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
    barrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;
    barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.buffer = buffer.inner();
    barrier.offset = 0;
    barrier.size = VK_WHOLE_SIZE;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT | VK_PIPELINE_STAGE_HOST_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        0,
        0, nullptr,
        1, &barrier,
        0, nullptr
    );
}

void Vk::CommandBuffer::transitionBufferToDraw(Vk::Buffer &buffer)
{
    VkBufferMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
    barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    barrier.dstAccessMask = VK_ACCESS_INDEX_READ_BIT | VK_ACCESS_VERTEX_ATTRIBUTE_READ_BIT;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.buffer = buffer.inner();
    barrier.offset = 0;
    barrier.size = VK_WHOLE_SIZE;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        VK_PIPELINE_STAGE_VERTEX_INPUT_BIT,
        0,
        0, nullptr,
        1, &barrier,
        0, nullptr
    );
}

void Vk::CommandBuffer::transitionBufferToShaderRead(Vk::Buffer &buffer)
{
    VkBufferMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
    barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.buffer = buffer.inner();
    barrier.offset = 0;
    barrier.size = VK_WHOLE_SIZE;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        VK_PIPELINE_STAGE_VERTEX_SHADER_BIT | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        0,
        0, nullptr,
        1, &barrier,
        0, nullptr
    );
}

void Vk::CommandBuffer::barrier(
    Vk::Buffer &buffer,
    VkPipelineStageFlags srcStage,
    VkAccessFlags srcAccessMask,
    VkPipelineStageFlags destStage,
    VkAccessFlags destAccessMask
)
{
    VkBufferMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
    barrier.srcAccessMask = srcAccessMask;
    barrier.dstAccessMask = destAccessMask;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.buffer = buffer.inner();
    barrier.offset = 0;
    barrier.size = VK_WHOLE_SIZE;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        srcStage,
        destStage,
        0,
        0, nullptr,
        1, &barrier,
        0, nullptr
    );
}

void Vk::CommandBuffer::barrier(
    Vk::Image& image,
    VkPipelineStageFlags srcStage,
    VkAccessFlags srcAccessMask,
    VkPipelineStageFlags dstStage,
    VkAccessFlags dstAccessMask,
    VkImageLayout oldLayout,
    VkImageLayout newLayout,
    VkImageAspectFlags aspectMask
)
{
    VkImageMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;

    barrier.srcAccessMask = srcAccessMask;
    barrier.dstAccessMask = dstAccessMask;

    barrier.oldLayout = oldLayout;
    barrier.newLayout = newLayout;

    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;

    barrier.image = image.inner();
    barrier.subresourceRange.aspectMask = aspectMask;
    barrier.subresourceRange.baseMipLevel = 0;
    barrier.subresourceRange.levelCount = 1;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount = VK_REMAINING_ARRAY_LAYERS;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        srcStage,
        dstStage,
        0,
        0, nullptr,
        0, nullptr,
        1, &barrier
    );
}

void Vk::CommandBuffer::transitionImageLayout(VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout)
{
    VkImageMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.oldLayout = oldLayout;
    barrier.newLayout = newLayout;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    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 = VK_REMAINING_ARRAY_LAYERS;

    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 if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_PRESENT_SRC_KHR)
    {
        barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
        barrier.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT;

        sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
        destinationStage = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
    }
    else if (oldLayout == VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
    {
        barrier.srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
        sourceStage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;

        barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        destinationStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    }
    else
    {
        crash("Unsupported image layout transition!");
    }

    vkCmdPipelineBarrier(
        m_commandBuffer,
        sourceStage, destinationStage,
        0,
        0, nullptr,
        0, nullptr,
        1, &barrier
    );
}

void Vk::CommandBuffer::transitionImageLayout(Vk::Image &image, VkImageLayout oldLayout, VkImageLayout newLayout)
{
    transitionImageLayout(image.inner(), oldLayout, newLayout);
}

void Vk::CommandBuffer::copyBufferToImage(Vk::Buffer &buffer, Vk::Image &image, uint32_t width, uint32_t height, uint32_t baseLayer)
{
    VkBufferImageCopy region{};
    region.bufferOffset = 0;
    region.bufferRowLength = 0;
    region.bufferImageHeight = 0;
    region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    region.imageSubresource.mipLevel = 0;
    region.imageSubresource.baseArrayLayer = baseLayer;
    region.imageSubresource.layerCount = 1;
    region.imageOffset = {0, 0, 0};
    region.imageExtent = {
        width,
        height,
        1
    };

    vkCmdCopyBufferToImage(m_commandBuffer, buffer.inner(), image.inner(), VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);
}

void Vk::CommandBuffer::copyImageToBuffer(VkImage image, Vk::Buffer &buffer, uint32_t width, uint32_t height, uint32_t baseLayer)
{
    VkBufferImageCopy region = {};
    region.bufferOffset = 0;
    region.bufferRowLength = 0; // Tightly packed
    region.bufferImageHeight = 0; // Tightly packed
    region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    region.imageSubresource.mipLevel = 0;
    region.imageSubresource.baseArrayLayer = 0;
    region.imageSubresource.layerCount = 1;
    region.imageOffset = {0, 0, 0};
    region.imageExtent = {width, height, 1};

    vkCmdCopyImageToBuffer(m_commandBuffer, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, buffer.inner(), 1, &region);
}

void Vk::CommandBuffer::copyImageToBuffer(Vk::Image &image, Vk::Buffer &buffer, uint32_t width, uint32_t height, uint32_t baseLayer)
{
    copyImageToBuffer(image.inner(), buffer, width, height, baseLayer);
}

void Vk::CommandBuffer::copyBuffer(Vk::Buffer &src, size_t srcOffset, Vk::Buffer &dest, size_t destOffset, size_t size)
{
    VkBufferCopy copyRegion{};
    copyRegion.srcOffset = srcOffset;
    copyRegion.dstOffset = destOffset;
    copyRegion.size = size;
    
    vkCmdCopyBuffer(m_commandBuffer, src.inner(), dest.inner(), 1, &copyRegion);
}

void Vk::CommandBuffer::copyBuffer(Vk::Buffer &src, Vk::Buffer &dest, const std::vector<VkBufferCopy> &regions)
{
    vkCmdCopyBuffer(m_commandBuffer, src.inner(), dest.inner(), regions.size(), regions.data());
}

void Vk::CommandBuffer::bindVertexBuffer(int binding, Vk::Buffer &buffer, size_t offset)
{
    VkBuffer innerBuffer = buffer.inner();
    VkDeviceSize devOffset = offset;

    vkCmdBindVertexBuffers(m_commandBuffer, binding, 1, &innerBuffer, &devOffset);
}

void Vk::CommandBuffer::bindIndexBuffer(Vk::Buffer &buffer, size_t offset, VkIndexType indexType)
{
    vkCmdBindIndexBuffer(m_commandBuffer, buffer.inner(), offset, indexType);
}

void Vk::CommandBuffer::execute(Vk::CommandBuffer &secondary)
{
    VkCommandBuffer secondaryInner = secondary.inner();
    vkCmdExecuteCommands(m_commandBuffer, 1, &secondaryInner);
}

void Vk::CommandBuffer::writeTimestamp(Vk::QueryPool &queryPool, uint32_t index)
{
    vkCmdWriteTimestamp(m_commandBuffer, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, queryPool.inner(), index);
}

void Vk::CommandBuffer::resetQueryPool(Vk::QueryPool &queryPool, uint32_t firstIndex, uint32_t count)
{
    vkCmdResetQueryPool(m_commandBuffer, queryPool.inner(), firstIndex, count);
}

void Vk::CommandBuffer::dispatch(glm::uvec3 workGroupCounts)
{
    vkCmdDispatch(m_commandBuffer, workGroupCounts.x, workGroupCounts.y, workGroupCounts.z);
}

void Vk::CommandBuffer::end()
{
    if (vkEndCommandBuffer(m_commandBuffer) != VK_SUCCESS)
    {
        crash("Failed to record the command buffer!");
    }
}

void Vk::CommandBuffer::submit(Vk::Fence *fence)
{
    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.waitSemaphoreCount = m_waitForSemaphores.size();
    submitInfo.pWaitSemaphores = m_waitForSemaphores.data();
    submitInfo.pWaitDstStageMask = m_waitStages.data();
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &m_commandBuffer;
    submitInfo.signalSemaphoreCount = m_signalSemaphores.size();
    submitInfo.pSignalSemaphores = m_signalSemaphores.data();

    VkFence innerFence = fence == nullptr ? VK_NULL_HANDLE : fence->inner();
    if (vkQueueSubmit(m_device->graphicsQueue.inner(), 1, &submitInfo, innerFence) != VK_SUCCESS)
    {
        crash("Failed to submit a command buffer!");
    }
}