#ifndef VULKAN_COMMAND_BUFFER_HPP_
#define VULKAN_COMMAND_BUFFER_HPP_

#include <Utilities/Extern.hpp>
#include <Vk/CommandPool/CommandPool.hpp>
#include <Vk/RenderPass/RenderPass.hpp>
#include <Vk/Framebuffer/Framebuffer.hpp>
#include <Vk/Pipeline/Pipeline.hpp>
#include <Vk/Semaphore/Semaphore.hpp>
#include <Vk/Fence/Fence.hpp>
#include <Vk/Buffer/Buffer.hpp>
#include <Vk/Image/Image.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <Vk/QueryPool/QueryPool.hpp>

namespace Vk
{
    /**
     * @brief Represents a Vulkan command pool.
     */
    class CommandBuffer
    {
    private:
        std::shared_ptr<Vk::Device> m_device;
        std::shared_ptr<Vk::CommandPool> m_pool;
        VkCommandBuffer m_commandBuffer;
        std::vector<VkSemaphore> m_waitForSemaphores;
        std::vector<VkPipelineStageFlags> m_waitStages;
        std::vector<VkSemaphore> m_signalSemaphores;
    
    public:
        /**
         * @brief Construct a command buffer.
         * 
         * This object takes shared ownership of the command pool, to ensure the pool
         * is not destroyed until all command buffer references are.
         * 
         * @param device The device.
         * @param pool The command pool to allocate from.
         * @param level The command buffer level.
         */
        CommandBuffer(
            std::shared_ptr<Vk::Device> device,
            std::shared_ptr<Vk::CommandPool> pool,
            VkCommandBufferLevel level
        );
        CommandBuffer(const CommandBuffer &) = delete;

        /**
         * @brief Return the wrapped command buffer.
         */
        VkCommandBuffer inner();

        /**
         * @brief Reset the command buffer.
         * 
         * @param flags The flags.
         */
        void reset(VkCommandBufferResetFlags flags = 0);

        /**
         * @brief Wait for the specified semaphore.
         * 
         * This function indicates that the GPU should wait for this semaphore before
         * performing any commands from this buffer. It does not cause the CPU to wait.
         * 
         * @param semaphore The semaphore.
         * @param stages The stages at which the wait shall occur.
         */
        void waitFor(Vk::Semaphore &semaphore, VkPipelineStageFlags stages);

        /**
         * @brief Signal the specified semaphore once the commands complete.
         * 
         * @param semaphore The semaphore to be signalled.
         */
        void signal(Vk::Semaphore &semaphore);

        /**
         * @brief Begin recording commands.
         * 
         * @param usage The usage flags (none by defualt).
         */
        void begin(VkCommandBufferUsageFlags usage = 0);

        /**
         * @brief Begin a secondary command buffer for the specified render pass and framebuffer.
         */
        void begin(Vk::RenderPass &renderPass, uint32_t subpassIndex, Vk::Framebuffer &framebuffer);

        /**
         * @brief Record a command to begin the render pass.
         */
        void beginRenderPass(
            Vk::RenderPass &renderPass,
            Vk::Framebuffer &framebuffer,
            VkExtent2D renderExtent,
            const std::vector<VkClearValue> &clearValues,
            VkSubpassContents subpassContents
        );

        /**
         * @brief Record a command to bind a pipeline.
         */
        void bindPipeline(VkPipelineBindPoint bindPoint, Vk::Pipeline &pipeline);

        /**
         * @brief Record a command to bind a descriptor set.
         */
        void bindDescriptorSet(
            VkPipelineBindPoint bindPoint,
            Vk::Pipeline &pipeline,
            uint32_t index,
            Vk::DescriptorSet &descSet
        );

        /**
         * @brief Record a command to set the viewport.
         */
        void setViewport(float x, float y, float width, float height, float minDepth = 0.0f, float maxDepth = 1.0f);

        /**
         * @brief Record a command to set the scissor rectangle.
         */
        void setScissor(VkRect2D rect);

        /**
         * @brief Record a command to update a range of push constants.
         */
        void pushConstants(
            Vk::Pipeline &pipeline,
            VkShaderStageFlags stages,
            uint32_t offset, uint32_t size,
            const void *data
        );

        /**
         * @brief Record a command to draw.
         */
        void draw(uint32_t vertexCount, uint32_t instanceCount, uint32_t firstVertex, uint32_t firstInstance);

        /**
         * @brief Record a command to draw using an index buffer.
         */
        void drawIndexed(
            uint32_t indexCount,
            uint32_t instanceCount,
            uint32_t firstIndex,
            int32_t vertexOffset,
            uint32_t firstInstance
        );

        /**
         * @brief Record a command to end the render pass.
         */
        void endRenderPass();

        /**
         * @brief Prepare a buffer to be a transfer destination from host memory.
         */
        void transitionBufferToTransferDst(Vk::Buffer &buffer);

        /**
         * @brief Transisiton a buffer to be used as vertex/index inputs to a draw.
         */
        void transitionBufferToDraw(Vk::Buffer &buffer);

        /**
         * @brief Transition a buffer to be used as a shader input (uniform buffer, storage buffer, etc).
         */
        void transitionBufferToShaderRead(Vk::Buffer &buffer);

        /**
         * @brief Issue a barrier.
         */
        void barrier(
            Vk::Buffer &buffer,
            VkPipelineStageFlags srcStage,
            VkAccessFlags srcAccessMask,
            VkPipelineStageFlags destStage,
            VkAccessFlags destAccessMask
        );

        /**
         * @brief Issue a barrier.
         */
        void barrier(
            Vk::Image& image,
            VkPipelineStageFlags srcStage,
            VkAccessFlags srcAccessMask,
            VkPipelineStageFlags dstStage,
            VkAccessFlags dstAccessMask,
            VkImageLayout oldLayout,
            VkImageLayout newLayout,
            VkImageAspectFlags aspectMask
        );

        /**
         * @brief Record a command to change the layout of an image.
         */
        void transitionImageLayout(VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout);

        /**
         * @brief Record a command to change the layout of an image.
         */
        void transitionImageLayout(Vk::Image &image, VkImageLayout oldLayout, VkImageLayout newLayout);

        /**
         * @brief Copy data from a buffer to an image.
         */
        void copyBufferToImage(Vk::Buffer &buffer, Vk::Image &image, uint32_t width, uint32_t height, uint32_t baseLayer = 0);

        /**
         * @brief Copy data from an image to a buffer.
         */
        void copyImageToBuffer(VkImage image, Vk::Buffer &buffer, uint32_t width, uint32_t height, uint32_t baseLayer = 0);

        /**
         * @brief Copy data from an image to a buffer.
         */
        void copyImageToBuffer(Vk::Image &image, Vk::Buffer &buffer, uint32_t width, uint32_t height, uint32_t baseLayer = 0);

        /**
         * @brief Copy data from a source buffer to a destination buffer.
         */
        void copyBuffer(Vk::Buffer &src, size_t srcOffset, Vk::Buffer &dest, size_t destOffset, size_t size);
    
        /**
         * @brief Copy data from a source buffer to a destination buffer.
         */
        void copyBuffer(Vk::Buffer &src, Vk::Buffer &dest, const std::vector<VkBufferCopy> &regions);

        /**
         * @brief Record a command to bind a vertex buffer.
         */
        void bindVertexBuffer(int binding, Vk::Buffer &buffer, size_t offset);

        /**
         * @brief Record a command to bind an index buffer.
         */
        void bindIndexBuffer(Vk::Buffer &buffer, size_t offset, VkIndexType indexType);

        /**
         * @brief Record a command to execute a command buffer.
         */
        void execute(Vk::CommandBuffer &secondary);

        /**
         * @brief Record a command to write to a timestamp from a query pool.
         */
        void writeTimestamp(Vk::QueryPool &queryPool, uint32_t index);

        /**
         * @brief Record a command to reset the specified queries in a query pool.
         */
        void resetQueryPool(Vk::QueryPool &queryPool, uint32_t firstIndex, uint32_t count);

        /**
         * @brief Dispatch the currently-bound compute shader.
         * 
         * @param workGroupCounts The number of workgroups in each dimension.
         */
        void dispatch(glm::uvec3 workGroupCounts);

        /**
         * @brief End the command buffer recording.
         */
        void end();

        /**
         * @brief Submit the buffer to the graphics queue.
         * 
         * @param fence The fence to signal once this buffer completes, or `nullptr`.
         */
        void submit(Vk::Fence *fence = nullptr);
    };
}

#endif