#include <BlockRegistry/BlockRegistry.hpp>
#include <PipelineStage/DrawStage/TerrainStage/TerrainStage.hpp>
#include <FrameDrawer/FrameDrawer.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Settings/Settings.hpp>
#include <Vk/ShaderConst.hpp>
#include <VoxelTree/VoxelTree.hpp>
#include <ChunkWindow/ChunkWindow.hpp>
#include <Cubemap/Cubemap.hpp>

TerrainStage::TerrainStage(FrameDrawer &frameDrawer, std::shared_ptr<Vk::DescriptorSet> dataDescSet) :
    DrawStage(
        "Terrain",
        frameDrawer,
        *frameDrawer.worldDisplay.terrainRenderPass,
        *frameDrawer.worldDisplay.terrainFramebuffer,
        *frameDrawer.worldDisplay.terrainPipeline,
        {}
    ),
    dataDescSet(dataDescSet),
    tintManager(frameDrawer)
{
    std::shared_ptr<Vk::DescriptorPool> skyboxDescriptorPool = Vk::DescriptorPool::Builder()
        .addDescriptorType(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1)
        .addDescriptorType(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 2)
        .build(GraphicsContext::instance->getDevice(), 1);

    skyboxDataBuffer = std::make_shared<Vk::Buffer>(
        GraphicsContext::instance->getDevice(),
        sizeof(SkyboxData),
        VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    skyboxDescSet = std::make_shared<Vk::DescriptorSet>(
        GraphicsContext::instance->getDevice(),
        skyboxDescriptorPool,
        frameDrawer.worldDisplay.skyboxDescriptorSetLayout
    );

    skyboxDescSet->setUniformBuffer(GLSL_SKYBOX_DATA_INDEX, *skyboxDataBuffer, 0, sizeof(SkyboxData));
    skyboxDescSet->setImageSampler(GLSL_SKYBOX_STAR_NOISE_INDEX, Cubemap::starNoise.imageView(), Vk::Sampler::cubemap());
    skyboxDescSet->setImageSampler(GLSL_SKYBOX_SUNMAP_INDEX, Cubemap::sunmap.imageView(), Vk::Sampler::cubemap());
}

void TerrainStage::onFrameBegins()
{
    SkyboxData *skyboxData = (SkyboxData*) skyboxDataBuffer->hostdata();
    skyboxData->model = frameDrawer.worldDisplay.getPlayer()->getSkyboxTransform();
    skyboxData->antiModel = glm::inverse(skyboxData->model);
    skyboxData->color = frameDrawer.worldDisplay.getWorld()->getSkyColor();
    skyboxData->screenSize = glm::vec2(GraphicsContext::instance->getWindowWidth(), GraphicsContext::instance->getWindowHeight());
    skyboxData->playerZ = frameDrawer.worldDisplay.getPlayer()->getEyePos().z;
    skyboxData->farZ = (float) (settings.horizontalViewDistance * 16);
    skyboxData->debugMode = frameDrawer.worldDisplay.gpuDebugMode;
    skyboxData->fxaaEnabled = settings.fxaa;

    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        GLSL_TERRAIN_SKYBOX_DESCRIPTOR_SET,
        *skyboxDescSet
    );

    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        GLSL_TERRAIN_FRAME_DESCRIPTOR_SET,
        *dataDescSet
    );
}

void TerrainStage::onFrameEnds()
{
    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        GLSL_TERRAIN_SVO_DESCRIPTOR_SET,
        frameDrawer.worldDisplay.getWorld()->window->svo.getDescSet(GraphicsContext::instance->cmdbuf())
    );

    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        GLSL_TERRAIN_VOXEL_PALETTE_DESCRIPTOR_SET,
        blockRegistry.getVoxelPaletteDescSet()
    );

    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        GLSL_TERRAIN_TINT_DESCRIPTOR_SET,
        *tintManager.getDescSet()
    );

    cmdbuf->draw(6, 1, 0, 0);
}