#ifndef DRAW_STAGE_HPP_
#define DRAW_STAGE_HPP_

#include <PipelineStage/PipelineStage.hpp>
#include <Vk/Framebuffer/Framebuffer.hpp>
#include <Vk/Pipeline/Pipeline.hpp>
#include <Vk/QueryPool/QueryPool.hpp>
#include <Vk/RenderPass/RenderPass.hpp>
#include <functional>
#include <memory>
#include <Vk/CommandBuffer/CommandBuffer.hpp>
#include <AverageTimeTracker/AverageTimeTracker.hpp>

class FrameDrawer;

/**
 * @brief Base class for drawing stages.
 * 
 * A drawing stage records a secondary command buffer which will be executed within a specific
 * render pass of a game frame. They are orchestrated and executed by the `FrameDrawer`.
 */
class DrawStage : public PipelineStage
{
protected:
    FrameDrawer &frameDrawer;
    Vk::RenderPass &renderPass;
    std::function<Vk::Framebuffer&()> getFramebuffer;
    Vk::Pipeline &pipeline;
    std::shared_ptr<Vk::CommandBuffer> cmdbuf;
    std::vector<VkClearValue> clearValues;

    /**
     * @brief Called when the frame is about to begin.
     * 
     * The pipeline is already bound and the viewport and scissor are set.
     */
    virtual void onFrameBegins() = 0;

    /**
     * @brief Called when the frame is about to end.
     * 
     * You can add additional commands to the command buffer here. Does nothing by
     * default. 
     */
    virtual void onFrameEnds();

    /**
     * @brief Returns the extent of the framebuffer to which this stage draws.
     */
    virtual VkExtent2D getFramebufferExtent();

public:
    DrawStage(
        std::string label,
        FrameDrawer &frameDrawer,
        Vk::RenderPass &renderPass,
        Vk::Framebuffer &framebuffer,
        Vk::Pipeline &pipeline,
        std::vector<VkClearValue> clearValues
    );

    DrawStage(
        std::string label,
        FrameDrawer &frameDrawer,
        Vk::RenderPass &renderPass,
        std::function<Vk::Framebuffer&()> getFramebuffer,
        Vk::Pipeline &pipeline,
        std::vector<VkClearValue> clearValues
    );

    virtual void initFrame() override;
    virtual void recordStageCommands(Vk::CommandBuffer &cmdbuf) override;

    /**
     * @brief Returns a reference to the command buffer for this stage.
     */
    Vk::CommandBuffer& getStageBuffer();
};

#endif