#ifndef LIGHTING_STAGE_HPP_
#define LIGHTING_STAGE_HPP_

#include <PipelineStage/PipelineStage.hpp>
#include <Vk/CommandBuffer/CommandBuffer.hpp>
#include <Vk/DescriptorPool/DescriptorPool.hpp>
#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <GPUMirroredArray/GPUMirroredArray.hpp>
#include <memory>

class FrameDrawer;

/**
 * @brief The drawing stage at which we add up lights by drawing light volumes.
 */
class LightingStage : public PipelineStage
{
public:
    typedef uint32_t LightID;

    /**
     * @brief Represents a light source for the light shader.
     */
    struct LightSource
    {
        alignas(16) glm::uvec4 sourceVoxelAndRadius;
        alignas(16) glm::vec4 colorAndVariance;

        LightSource(glm::uvec3 sourceVoxel, glm::vec3 color, float variance = 0.0f);
    };

    /**
     * @brief Represents an on-screen cluster.
     *
     * The cluster specifies an offset into the light index buffer, and the number of lights
     * in the cluster.
     */
    struct Cluster
    {
        uint32_t offset;
        uint32_t count;
    };

    /**
     * @brief Represents the bounding box of a light in cluster coordinates.
     */
    struct LightScope
    {
        glm::uvec4 minScope;
        glm::uvec4 maxScope;
    };

private:
    struct Session
    {
        Session();

        GPUMirroredArray<LightSource> lightSources;
        std::vector<LightID> lightReclaimStack;

        std::shared_ptr<Vk::DescriptorPool> lightStorageDescPool;
        std::shared_ptr<Vk::Buffer> lightScopeBuffer;
        std::shared_ptr<Vk::DescriptorSet> lightScopeDescSet;
        size_t currentLightScopeCapacity = 0;

        std::shared_ptr<Vk::Buffer> lightIndexBuffer;
        std::shared_ptr<Vk::DescriptorSet> lightIndexDescSet;
    };

    static constexpr size_t maxEntityLights = 256;

    std::shared_ptr<Vk::DescriptorSet> dataDescSet;
    std::shared_ptr<Session> session;

    size_t numEntityLights = 0;

    FrameDrawer &frameDrawer;

    void checkScopeBuffer();
    void generateTileDepthBuffer(Vk::CommandBuffer &cmdbuf, glm::uvec2 screenTiles);
    void resetLightClusters(Vk::CommandBuffer &cmdbuf, glm::uvec2 screenTiles);
    void computeLightScopes(Vk::CommandBuffer &cmdbuf, Vk::DescriptorSet &sourcesDescSet);
    void countLights(Vk::CommandBuffer &cmdbuf, glm::uvec2 screenTiles);
    void applyLights(Vk::CommandBuffer &cmdbuf, glm::uvec2 screenTiles, Vk::DescriptorSet &sourcesDescSet);

public:
    LightingStage(
        FrameDrawer &frameDrawer,
        std::shared_ptr<Vk::DescriptorSet> dataDescSet,
        std::shared_ptr<Vk::DescriptorPool> frameDataDescriptorPool
    );

    /**
     * @brief Add an entity light source.
     *
     * This method can be called from an entity's `draw()` implementation, to add
     * a light source which should be emitted this frame.
     */
    void addEntityLight(LightSource lightSource);

    LightingStage::LightID addBlockLight(LightingStage::LightSource lightSource);
    void setBlockLight(LightingStage::LightID id, LightingStage::LightSource lightSource);
    void releaseBlockLight(LightingStage::LightID id);
    size_t getActiveLights();

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

#endif