#include <PipelineStage/DrawStage/DrawStage.hpp>
#include <FrameDrawer/FrameDrawer.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <PipelineStage/PipelineStage.hpp>
#include <Vk/Framebuffer/Framebuffer.hpp>
#include <Vk/QueryPool/QueryPool.hpp>

void DrawStage::onFrameEnds()
{
    // NOP
}

VkExtent2D DrawStage::getFramebufferExtent()
{
    return frameDrawer.worldDisplay.framebufferExtent;
}

DrawStage::DrawStage(
    std::string label,
    FrameDrawer &frameDrawer,
    Vk::RenderPass &renderPass,
    Vk::Framebuffer &framebuffer,
    Vk::Pipeline &pipeline,
    std::vector<VkClearValue> clearValues
) : PipelineStage(label),
    frameDrawer(frameDrawer),
    renderPass(renderPass),
    getFramebuffer([&framebuffer]() -> Vk::Framebuffer& { return framebuffer; }),
    pipeline(pipeline),
    clearValues(clearValues)
{
    cmdbuf = frameDrawer.createStageCommandBuffer();
}

DrawStage::DrawStage(
    std::string label,
    FrameDrawer &frameDrawer,
    Vk::RenderPass &renderPass,
    std::function<Vk::Framebuffer&()> getFramebuffer,
    Vk::Pipeline &pipeline,
    std::vector<VkClearValue> clearValues
) : PipelineStage(label),
    frameDrawer(frameDrawer),
    renderPass(renderPass),
    getFramebuffer(getFramebuffer),
    pipeline(pipeline),
    clearValues(clearValues)
{
    cmdbuf = frameDrawer.createStageCommandBuffer();
}

void DrawStage::initFrame()
{
    cmdbuf->reset();
    cmdbuf->begin(renderPass, 0, getFramebuffer());

    VkExtent2D extent = getFramebufferExtent();

    cmdbuf->bindPipeline(VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
    cmdbuf->setViewport(
        0.0f, 0.0f,
        (float) extent.width, (float) extent.height
    );
    cmdbuf->setScissor({{0, 0}, extent});

    onFrameBegins();
}

void DrawStage::recordStageCommands(Vk::CommandBuffer &primary)
{
    onFrameEnds();
    cmdbuf->end();

    primary.beginRenderPass(
        renderPass,
        getFramebuffer(),
        getFramebufferExtent(),
        clearValues,
        VK_SUBPASS_CONTENTS_SECONDARY_COMMAND_BUFFERS
    );

    primary.execute(*cmdbuf);
    primary.endRenderPass();
}

Vk::CommandBuffer& DrawStage::getStageBuffer()
{
    return *cmdbuf;
}