#include "CommandPool.hpp"

#include "shared.hpp"
#include "../Renderer.hpp"

using namespace render::vk;

CommandPool::CommandPool(VkDevice device, const std::vector<VkImageView>& views, PipelineRef pipe, const PhysicalDeviceInfo& info, const renderOptions& opt):
        device(device) {

    VkCommandPoolCreateInfo poolInfo{};
    poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.queueFamilyIndex = info.queueIndices.graphicsFamily.value();
    poolInfo.flags = 0; // Optional

    if (vkCreateCommandPool(device, &poolInfo, ALLOC, &graphicsPool) != VK_SUCCESS) {
        FATAL("Failed to create command pool!");
    }

    allocate(views, pipe, info.swapDetails.capabilities.currentExtent, opt);
}
CommandPool::~CommandPool() {
    if(!freed)
        free();

    vkDestroyCommandPool(device, graphicsPool, ALLOC);
}
void CommandPool::allocate(const std::vector<VkImageView>& views, PipelineRef pipe, VkExtent2D extent, const renderOptions& opt) {
    assert(freed);

    framebuffers.resize(views.size());

    for (size_t i = 0; i < views.size(); i++) {
        VkImageView attachments[] = { views[i] };

        VkFramebufferCreateInfo framebufferInfo{};
        framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
        framebufferInfo.renderPass = pipe.first;
        framebufferInfo.attachmentCount = 1;
        framebufferInfo.pAttachments = attachments;
        framebufferInfo.width = extent.width;
        framebufferInfo.height = extent.height;
        framebufferInfo.layers = 1;

        if (vkCreateFramebuffer(device, &framebufferInfo, ALLOC, &framebuffers[i]) != VK_SUCCESS) {
            FATAL("Failed to create framebuffer!");
        }
    }

    graphicsBuffers.resize(framebuffers.size());

    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.commandPool = graphicsPool;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = (uint32_t) graphicsBuffers.size();

    if (vkAllocateCommandBuffers(device, &allocInfo, graphicsBuffers.data()) != VK_SUCCESS) {
        FATAL("Failed to allocate command buffers!");
    }

    for (size_t i = 0; i < graphicsBuffers.size(); i++) {
        VkCommandBufferBeginInfo beginInfo{};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        beginInfo.flags = 0;
        beginInfo.pInheritanceInfo = nullptr;

        if (vkBeginCommandBuffer(graphicsBuffers[i], &beginInfo) != VK_SUCCESS) {
            FATAL("Failed to begin recording command buffer!");
        }

        VkRenderPassBeginInfo renderPassInfo{};
        renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
        renderPassInfo.renderPass = pipe.first;
        renderPassInfo.framebuffer = framebuffers[i];
        renderPassInfo.renderArea.offset = {0, 0};
        renderPassInfo.renderArea.extent = extent;

        VkClearValue clearColor = {opt.clear_color.x, opt.clear_color.y, opt.clear_color.z, opt.clear_color.a};
        renderPassInfo.clearValueCount = 1; //TODO: clear depth
        renderPassInfo.pClearValues = &clearColor;

        vkCmdBeginRenderPass(graphicsBuffers[i], &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);
        vkCmdBindPipeline(graphicsBuffers[i], VK_PIPELINE_BIND_POINT_GRAPHICS, pipe.second);
        vkCmdDraw(graphicsBuffers[i], 3, 1, 0, 0);
        vkCmdEndRenderPass(graphicsBuffers[i]);

        if (vkEndCommandBuffer(graphicsBuffers[i]) != VK_SUCCESS) {
            FATAL("Failed to record command buffer!");
        }
    }

    freed = false;
}
void CommandPool::free() {
    assert(!freed);
    for (size_t i = 0; i < framebuffers.size(); i++) {
        vkDestroyFramebuffer(device, framebuffers[i], nullptr);
    }

    vkFreeCommandBuffers(device, graphicsPool, static_cast<uint32_t>(graphicsBuffers.size()), graphicsBuffers.data());

    freed = true;
}

void CommandPool::submitGraphics(uint32_t idx, VkQueue graphicsQueue, VkSemaphore waitSemaphore, VkSemaphore signalSemaphore, VkFence submittedFence) {
    assert(!freed);

    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

    VkSemaphore waitSemaphores[] = {waitSemaphore};
    VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
    submitInfo.waitSemaphoreCount = 1;
    submitInfo.pWaitSemaphores = waitSemaphores;
    submitInfo.pWaitDstStageMask = waitStages;

    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &graphicsBuffers[idx];

    VkSemaphore signalSemaphores[] = {signalSemaphore};
    submitInfo.signalSemaphoreCount = 1;
    submitInfo.pSignalSemaphores = signalSemaphores;

    vkResetFences(device, 1, &submittedFence);
    if (vkQueueSubmit(graphicsQueue, 1, &submitInfo, submittedFence) != VK_SUCCESS) {
        FATAL("Failed to submit draw command buffer!");
    }
}